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

[mlir][sparse] Breaking up openSparseTensor to better support non-permutations

Browse Source 
This commit updates how the `SparseTensorConversion` pass handles `NewOp`.  It breaks up the underlying `openSparseTensor` function into two parts (`SparseTensorReader::create` and `SparseTensorReader::readSparseTensor`) so that the pass can inject code for constructing `lvlSizes` between those two parts.  Migrating the construction of `lvlSizes` out of the runtime and into the pass is a necessary first step toward fully supporting non-permutations.  (The alternative would be for the pass to generate a `FuncOp` for performing the construction and then passing that to the runtime; which doesn't seem to have any benefits over the design of this commit.)  And since the pass now generates the code to call these two functions, this change also removes the `Action::kFromFile` value from the enum used by `_mlir_ciface_newSparseTensor`.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D138363
This commit is contained in:
wren romano
2022-12-01 18:18:33 -08:00
parent ca23b7ca47
commit 2af2e4dbb7
10 changed files with 447 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
View File

@@ -913,9 +913,8 @@ struct NewRewriter : public OpRewritePattern<NewOp> {
Location loc = op.getLoc();
auto dstTp = op.getResult().getType().template cast<RankedTensorType>();
SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
if (!encDst) {
if (!encDst)
return failure();
}
// Create a sparse tensor reader.
Value fileName = op.getSource();
@@ -933,7 +932,7 @@ struct NewRewriter : public OpRewritePattern<NewOp> {
// the sparse tensor reader.
SmallVector<Value> dynSizesArray;
if (!dstTp.hasStaticShape()) {
createFuncCall(rewriter, loc, "getSparseTensorReaderDimSizes", {},
createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {},
{reader, dimSizes}, EmitCInterface::On)
.getResult(0);
ArrayRef<int64_t> dstShape = dstTp.getShape();
@@ -977,7 +976,7 @@ struct NewRewriter : public OpRewritePattern<NewOp> {
ArrayRef<Value>(cooBuffer));
rewriter.setInsertionPointToStart(forOp.getBody());
SmallString<18> getNextFuncName{"getSparseTensorReaderNext",
SmallString<29> getNextFuncName{"getSparseTensorReaderNext",
primaryTypeFunctionSuffix(eltTp)};
Value indices = dimSizes; // Reuse the indices memref to store indices.
createFuncCall(rewriter, loc, getNextFuncName, {}, {reader, indices, value},
@@ -1060,7 +1059,7 @@ struct OutRewriter : public OpRewritePattern<OutOp> {
Value indices = dimSizes; // Reuse the dimSizes buffer for indices.
Type eltTp = srcTp.getElementType();
SmallString<18> outNextFuncName{"outSparseTensorWriterNext",
SmallString<29> outNextFuncName{"outSparseTensorWriterNext",
primaryTypeFunctionSuffix(eltTp)};
Value value = genAllocaScalar(rewriter, loc, eltTp);
ModuleOp module = op->getParentOfType<ModuleOp>();