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

[mlir][sparse] reuse tensor.insert operation to insert elements into … (#84987)

Browse Source 
…a sparse tensor.
This commit is contained in:
Peiming Liu
2024-03-12 16:59:17 -07:00
committed by GitHub
parent 1c3b15e9f5
commit 94e27c265a
26 changed files with 106 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -580,17 +580,24 @@ public:
};
/// Sparse conversion rule for the insertion operator.
class SparseTensorInsertConverter : public OpConversionPattern<InsertOp> {
class SparseTensorInsertConverter
: public OpConversionPattern<tensor::InsertOp> {
public:
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(InsertOp op, OpAdaptor adaptor,
matchAndRewrite(tensor::InsertOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
// Note that the current regime only allows for strict lexicographic
// coordinate order. All values are passed by reference through stack
// allocated memrefs.
Location loc = op->getLoc();
const auto stt = getSparseTensorType(op.getTensor());
const auto stt = getSparseTensorType(op.getDest());
// Dense tensor insertion.
if (!stt.hasEncoding())
return failure();
assert(stt.isIdentity() && "Run reinterpret-map before conversion.");
const auto elemTp = stt.getElementType();
const Level lvlRank = stt.getLvlRank();
Value lvlCoords, vref;
@@ -608,12 +615,12 @@ public:
lvlCoords = genAlloca(rewriter, loc, lvlRank, rewriter.getIndexType());
vref = genAllocaScalar(rewriter, loc, elemTp);
}
storeAll(rewriter, loc, lvlCoords, adaptor.getLvlCoords());
rewriter.create<memref::StoreOp>(loc, adaptor.getValue(), vref);
storeAll(rewriter, loc, lvlCoords, adaptor.getIndices());
rewriter.create<memref::StoreOp>(loc, adaptor.getScalar(), vref);
SmallString<12> name{"lexInsert", primaryTypeFunctionSuffix(elemTp)};
createFuncCall(rewriter, loc, name, {},
{adaptor.getTensor(), lvlCoords, vref}, EmitCInterface::On);
rewriter.replaceOp(op, adaptor.getTensor());
{adaptor.getDest(), lvlCoords, vref}, EmitCInterface::On);
rewriter.replaceOp(op, adaptor.getDest());
return success();
}
};