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

[mlir][sparse] refactoring: split translateIndices.

Browse Source 
TranslateIndicesArray take an array of SSA value and convert them into another array of SSA values based on reassociation. Which makes it easier to be reused by `foreach` operator (as the indices array are given as an array of SSA values).

Reviewed By: aartbik, bixia

Differential Revision: https://reviews.llvm.org/D134918
This commit is contained in:
Peiming Liu
2022-09-29 21:59:02 +00:00
parent 073534cb6f
commit 11069cbcb4
4 changed files with 77 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -475,44 +475,21 @@ static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,
ArrayRef<Value> srcShape) {
unsigned dstRank = dstTp.getRank();
unsigned srcRank = srcTp.getRank();
unsigned start = 0;
unsigned i = 0;
bool isExpand = srcRank > dstRank;
ArrayRef<Value> shape = isExpand ? srcShape : dstShape;
// Iterate over reassociation map.
for (const auto &map : llvm::enumerate(reassociation)) {
// Prepare strides information in dimension slice.
Value linear = constantIndex(rewriter, loc, 1);
for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
linear = rewriter.create<arith::MulIOp>(loc, linear, shape[j]);
}
// Start collapse.
Value idx = constantIndex(rewriter, loc, i++);
Value val;
if (!isExpand)
val = rewriter.create<memref::LoadOp>(loc, srcIdx, idx);
// Iterate over dimension slice.
for (unsigned j = start, end = start + map.value().size(); j < end; j++) {
linear = rewriter.create<arith::DivUIOp>(loc, linear, shape[j]);
Value jdx = constantIndex(rewriter, loc, j);
if (isExpand) {
Value old = rewriter.create<memref::LoadOp>(loc, srcIdx, jdx);
Value mul = rewriter.create<arith::MulIOp>(loc, old, linear);
val = val ? rewriter.create<arith::AddIOp>(loc, val, mul) : mul;
} else {
Value old = val;
val = rewriter.create<arith::DivUIOp>(loc, val, linear);
rewriter.create<memref::StoreOp>(loc, val, dstIdx, jdx);
val = rewriter.create<arith::RemUIOp>(loc, old, linear);
}
}
// Finalize expansion.
if (isExpand)
rewriter.create<memref::StoreOp>(loc, val, dstIdx, idx);
start += map.value().size();
SmallVector<Value, 4> srcIndices;
for (unsigned i = 0; i < srcRank; i++) {
Value idx = rewriter.create<memref::LoadOp>(
loc, srcIdx, constantIndex(rewriter, loc, i));
srcIndices.push_back(idx);
}
// Sanity.
assert((isExpand && i == dstRank) || (!isExpand && i == srcRank));
SmallVector<Value, 4> dstIndices;
translateIndicesArray(rewriter, loc, reassociation, srcIndices, srcShape,
dstShape, dstIndices);
for (unsigned i = 0; i < dstRank; i++)
rewriter.create<memref::StoreOp>(loc, dstIndices[i], dstIdx,
constantIndex(rewriter, loc, i));
}
/// Helper method to compute the shape of destination tensor of a reshape