caio/clang-p2996
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[mlir][sparse] Refactor the conversion of the tensor reshape operators.

Browse Source 
Move genReshapeDstShape to codegen utils to support the rewriting of the tensor
reshape operators for the codegen path.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D135074
This commit is contained in:
bixia1
2022-10-02 20:19:44 -07:00
parent c2df707666
commit 349bceba65
3 changed files with 64 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
View File

@@ -492,65 +492,6 @@ static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
constantIndex(rewriter, loc, i));
}
/// Helper method to compute the shape of destination tensor of a reshape
/// operator. This is only used when operands have dynamic shape. The shape of
/// the destination is stored into dstShape.
void genReshapeDstShape(Location loc, ConversionPatternRewriter &rewriter,
SmallVector<Value, 4> &dstShape,
ArrayRef<Value> srcShape,
ArrayRef<int64_t> staticDstShape,
ArrayRef<ReassociationIndices> reassociation) {
// Collapse shape.
if (reassociation.size() < srcShape.size()) {
unsigned start = 0;
for (const auto &map : llvm::enumerate(reassociation)) {
auto dstDim = constantIndex(rewriter, loc, 1);
for (unsigned i = start; i < start + map.value().size(); i++) {
dstDim = rewriter.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
}
dstShape.push_back(dstDim);
start = start + map.value().size();
}
assert(start == srcShape.size());
return;
}
// Expand shape.
assert(reassociation.size() == srcShape.size());
unsigned start = 0;
// Expand the i-th dimension in srcShape.
for (unsigned i = 0, size = srcShape.size(); i < size; i++) {
auto map = reassociation[i];
auto srcDim = srcShape[i];
// Iterate through dimensions expanded from the i-th dimension.
for (unsigned j = start; j < start + map.size(); j++) {
// There can be only one dynamic sized dimension among dimensions expanded
// from the i-th dimension in srcShape. For example, if srcDim = 8, then
// the expanded shape could be <2x?x2>, but not <2x?x?>.
if (staticDstShape[j] == ShapedType::kDynamicSize) {
// The expanded dimension has dynamic size. We compute the dimension
// by dividing srcDim by the product of the static dimensions.
int64_t product = 1;
for (unsigned k = start; k < start + map.size(); k++) {
if (staticDstShape[k] != ShapedType::kDynamicSize) {
product *= staticDstShape[k];
}
}
// Compute the dynamic dimension size.
Value productVal = constantIndex(rewriter, loc, product);
Value dynamicSize =
rewriter.create<arith::DivUIOp>(loc, srcDim, productVal);
dstShape.push_back(dynamicSize);
} else {
// The expanded dimension is statically known.
dstShape.push_back(constantIndex(rewriter, loc, staticDstShape[j]));
}
}
start = start + map.size();
}
assert(start == staticDstShape.size());
}
/// Generate code for a general sparse to sparse reshaping operation.
/// Note that unlike dense reshaping (which can be done with a "cheap"
/// change of view), sparse reshaping is currently done with actual