[llvm-branch-commits] [mlir] [mlir][sparse] Add `has_runtime_library` test op (PR #85355)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Thu Mar 14 20:20:26 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Matthias Springer (matthias-springer)
<details>
<summary>Changes</summary>
This commit adds a new test-only op:
`sparse_tensor.has_runtime_library`. The op returns "1" if the sparse compiler runs in runtime library mode.
This op is useful for writing test cases that require different IR depending on whether the sparse compiler runs in runtime library or codegen mode.
This commit fixes a memory leak in `sparse_pack_d.mlir`. This test case uses `sparse_tensor.assemble` to create a sparse tensor SSA value from existing buffers. This runtime library reallocates+copies the existing buffers; the codegen path does not. Therefore, the test requires additional deallocations when running in runtime library mode.
Alternatives considered:
- Make the codegen path allocate. "Codegen" is the "default" compilation mode and it is handling `sparse_tensor.assemble` correctly. The issue is with the runtime library path, which should not allocate. Therefore, it is better to put a workaround in the runtime library path than to work around the issue with a new flag in the codegen path.
- Add a `sparse_tensor.runtime_only` attribute to `bufferization.dealloc_tensor`. Verifying that the attribute can only be attached to `bufferization.dealloc_tensor` may introduce an unwanted dependency of `MLIRSparseTensorDialect` on `MLIRBufferizationDialect`.
---
Full diff: https://github.com/llvm/llvm-project/pull/85355.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td (+12-1)
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp (+28-15)
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp (+16-2)
- (modified) mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir (+10-4)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
index 0498576fcffc51..1b2981ac751f3e 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
@@ -1419,7 +1419,7 @@ def SparseTensor_ForeachOp : SparseTensor_Op<"foreach",
}
//===----------------------------------------------------------------------===//
-// Sparse Tensor Debugging Operations.
+// Sparse Tensor Debugging and Test-Only Operations.
//===----------------------------------------------------------------------===//
def SparseTensor_PrintOp : SparseTensor_Op<"print">,
@@ -1440,4 +1440,15 @@ def SparseTensor_PrintOp : SparseTensor_Op<"print">,
let assemblyFormat = "$tensor attr-dict `:` type($tensor)";
}
+def SparseTensor_HasRuntimeLibraryOp
+ : SparseTensor_Op<"has_runtime_library", []>, Results<(outs I1:$result)> {
+ string summary = "Indicates whether running in runtime/codegen mode";
+ string description = [{
+ Returns a boolean value that indicates whether the sparse compiler runs in
+ runtime library mode or not. For testing only: This op is useful for writing
+ test cases that require different IR depending on runtime/codegen mode.
+ }];
+ let assemblyFormat = "attr-dict";
+}
+
#endif // SPARSETENSOR_OPS
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 7ff2fc25328a6c..5679f277e14866 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1561,6 +1561,19 @@ struct SparseNewConverter : public OpConversionPattern<NewOp> {
}
};
+struct SparseHasRuntimeLibraryConverter
+ : public OpConversionPattern<HasRuntimeLibraryOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(HasRuntimeLibraryOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto i1Type = rewriter.getI1Type();
+ rewriter.replaceOpWithNewOp<arith::ConstantOp>(
+ op, i1Type, rewriter.getIntegerAttr(i1Type, 0));
+ return success();
+ }
+};
+
} // namespace
//===----------------------------------------------------------------------===//
@@ -1572,21 +1585,21 @@ struct SparseNewConverter : public OpConversionPattern<NewOp> {
void mlir::populateSparseTensorCodegenPatterns(
TypeConverter &typeConverter, RewritePatternSet &patterns,
bool createSparseDeallocs, bool enableBufferInitialization) {
- patterns.add<SparseAssembleOpConverter, SparseDisassembleOpConverter,
- SparseReturnConverter, SparseCallConverter, SparseLvlOpConverter,
- SparseCastConverter, SparseExtractSliceConverter,
- SparseTensorLoadConverter, SparseExpandConverter,
- SparseCompressConverter, SparseInsertConverter,
- SparseReorderCOOConverter, SparseReMapConverter,
- SparseSliceGetterOpConverter<ToSliceOffsetOp,
- StorageSpecifierKind::DimOffset>,
- SparseSliceGetterOpConverter<ToSliceStrideOp,
- StorageSpecifierKind::DimStride>,
- SparseToPositionsConverter, SparseToCoordinatesConverter,
- SparseToCoordinatesBufferConverter, SparseToValuesConverter,
- SparseConvertConverter, SparseNewConverter,
- SparseNumberOfEntriesConverter>(typeConverter,
- patterns.getContext());
+ patterns.add<
+ SparseAssembleOpConverter, SparseDisassembleOpConverter,
+ SparseReturnConverter, SparseCallConverter, SparseLvlOpConverter,
+ SparseCastConverter, SparseExtractSliceConverter,
+ SparseTensorLoadConverter, SparseExpandConverter, SparseCompressConverter,
+ SparseInsertConverter, SparseReorderCOOConverter, SparseReMapConverter,
+ SparseSliceGetterOpConverter<ToSliceOffsetOp,
+ StorageSpecifierKind::DimOffset>,
+ SparseSliceGetterOpConverter<ToSliceStrideOp,
+ StorageSpecifierKind::DimStride>,
+ SparseToPositionsConverter, SparseToCoordinatesConverter,
+ SparseToCoordinatesBufferConverter, SparseToValuesConverter,
+ SparseConvertConverter, SparseNewConverter,
+ SparseNumberOfEntriesConverter, SparseHasRuntimeLibraryConverter>(
+ typeConverter, patterns.getContext());
patterns.add<SparseTensorDeallocConverter>(
typeConverter, patterns.getContext(), createSparseDeallocs);
patterns.add<SparseTensorAllocConverter, SparseTensorEmptyConverter>(
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index 0937c10f257283..92c98b34af6027 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -840,6 +840,19 @@ class SparseTensorDisassembleConverter
}
};
+struct SparseHasRuntimeLibraryConverter
+ : public OpConversionPattern<HasRuntimeLibraryOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(HasRuntimeLibraryOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto i1Type = rewriter.getI1Type();
+ rewriter.replaceOpWithNewOp<arith::ConstantOp>(
+ op, i1Type, rewriter.getIntegerAttr(i1Type, 1));
+ return success();
+ }
+};
+
} // namespace
//===----------------------------------------------------------------------===//
@@ -868,6 +881,7 @@ void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
SparseTensorToValuesConverter, SparseNumberOfEntriesConverter,
SparseTensorLoadConverter, SparseTensorInsertConverter,
SparseTensorExpandConverter, SparseTensorCompressConverter,
- SparseTensorAssembleConverter, SparseTensorDisassembleConverter>(
- typeConverter, patterns.getContext());
+ SparseTensorAssembleConverter, SparseTensorDisassembleConverter,
+ SparseHasRuntimeLibraryConverter>(typeConverter,
+ patterns.getContext());
}
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
index aa1bd04fde87dc..42401d94ca02c8 100755
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
@@ -140,10 +140,16 @@ module {
sparse_tensor.print %s1 : tensor<4x3x2xf32, #BatchedCSR>
sparse_tensor.print %s2 : tensor<4x3x2xf32, #CSRDense>
- // FIXME: doing this explicitly crashes runtime
- // bufferization.dealloc_tensor %s0 : tensor<4x3x2xf32, #CCC>
- // bufferization.dealloc_tensor %s1 : tensor<4x3x2xf32, #BatchedCSR>
- // bufferization.dealloc_tensor %s2 : tensor<4x3x2xf32, #CSRDense>
+ %has_runtime = sparse_tensor.has_runtime_library
+ scf.if %has_runtime {
+ // sparse_tensor.assemble copies buffers when running with the runtime
+ // library. Deallocations are needed not needed when running in codgen
+ // mode.
+ bufferization.dealloc_tensor %s0 : tensor<4x3x2xf32, #CCC>
+ bufferization.dealloc_tensor %s1 : tensor<4x3x2xf32, #BatchedCSR>
+ bufferization.dealloc_tensor %s2 : tensor<4x3x2xf32, #CSRDense>
+ }
+
return
}
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/85355
More information about the llvm-branch-commits
mailing list