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

[mlir][sparse] add a "release" operation to sparse tensor dialect

Browse Source 
We have several ways to materialize sparse tensors (new and convert) but no explicit operation to release the underlying sparse storage scheme at runtime (other than making an explicit delSparseTensor() library call). To simplify memory management, a sparse_tensor.release operation has been introduced that lowers to the runtime library call while keeping tensors, opague pointers, and memrefs transparent in the initial IR.

*Note* There is obviously some tension between the concept of immutable tensors and memory management methods. This tension is addressed by simply stating that after the "release" call, no further memref related operations are allowed on the tensor value. We expect the design to evolve over time, however, and arrive at a more satisfactory view of tensors and buffers eventually.

Bug:
http://llvm.org/pr52046

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D111099
This commit is contained in:
Aart Bik
2021-10-04 13:13:24 -07:00
parent 7a4e9a0c73
commit 16b8f4ddae
24 changed files with 191 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -99,8 +99,8 @@ static Value getTensor(ConversionPatternRewriter &rewriter, unsigned width,
/// the "_emit_c_interface" on the function declaration when requested,
/// so that LLVM lowering generates a wrapper function that takes care
/// of ABI complications with passing in and returning MemRefs to C functions.
static FlatSymbolRefAttr getFunc(Operation *op, StringRef name, Type resultType,
ValueRange operands,
static FlatSymbolRefAttr getFunc(Operation *op, StringRef name,
TypeRange resultType, ValueRange operands,
bool emitCInterface = false) {
MLIRContext *context = op->getContext();
auto module = op->getParentOfType<ModuleOp>();
@@ -471,6 +471,23 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
}
};
/// Sparse conversion rule for the release operator.
class SparseTensorReleaseConverter : public OpConversionPattern<ReleaseOp> {
public:
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(ReleaseOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
StringRef name = "delSparseTensor";
TypeRange none;
rewriter.create<CallOp>(op.getLoc(), none,
getFunc(op, name, none, adaptor.getOperands()),
adaptor.getOperands());
rewriter.eraseOp(op);
return success();
}
};
/// Sparse conversion rule for pointer accesses.
class SparseTensorToPointersConverter
: public OpConversionPattern<ToPointersOp> {
@@ -483,7 +500,7 @@ public:
Type eltType = resType.cast<ShapedType>().getElementType();
StringRef name;
if (eltType.isIndex())
name = "sparsePointers";
name = "sparsePointers"; // 64-bit, but its own name for unique signature
else if (eltType.isInteger(64))
name = "sparsePointers64";
else if (eltType.isInteger(32))
@@ -514,7 +531,7 @@ public:
Type eltType = resType.cast<ShapedType>().getElementType();
StringRef name;
if (eltType.isIndex())
name = "sparseIndices";
name = "sparseIndices"; // 64-bit, but its own name for unique signature
else if (eltType.isInteger(64))
name = "sparseIndices64";
else if (eltType.isInteger(32))
@@ -609,7 +626,8 @@ void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
RewritePatternSet &patterns) {
patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
SparseTensorNewConverter, SparseTensorConvertConverter,
SparseTensorToPointersConverter, SparseTensorToIndicesConverter,
SparseTensorToValuesConverter, SparseTensorToTensorConverter>(
typeConverter, patterns.getContext());
SparseTensorReleaseConverter, SparseTensorToPointersConverter,
SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
SparseTensorToTensorConverter>(typeConverter,
patterns.getContext());
}