[mlir][sparse] Add has_runtime_library test op
#85355
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
matthias-springer
requested review from
aartbik,
PeimingLiu and
yinying-lisa-li
as code owners
|
@llvm/pr-subscribers-mlir-sparse @llvm/pr-subscribers-mlir Author: Matthias Springer (matthias-springer) ChangesThis commit adds a new test-only op: 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 Alternatives considered:
Full diff: https://github.com/llvm/llvm-project/pull/85355.diff 4 Files Affected:
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
}
}
|
aartbik
approved these changes
Mar 15, 2024
matthias-springer
force-pushed
the
users/matthias-springer/sparse_leak_next
branch
from
889cee2 to
b2294d3
Compare
Base automatically changed from
users/matthias-springer/sparse_leak_next
to
main
March 15, 2024 04:31
matthias-springer
force-pushed
the
users/matthias-springer/has_runtime_lib
branch
from
c4fefa0 to
e921086
Compare
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`.
matthias-springer
force-pushed
the
users/matthias-springer/has_runtime_lib
branch
from
e921086 to
a5b2ebd
Compare
joker-eph
reviewed
Mar 15, 2024
| test cases that require different IR depending on runtime/codegen mode. | ||
| }]; | ||
| let assemblyFormat = "attr-dict"; | ||
| } |
There was a problem hiding this comment.
I'm a bit concerned about "test only" ops mixed up in the non-test dialects, and how this gets into "test patterns" mixed into the normal compilation flow.
Is the TODO in the test indicative that this will be reverted soon?
There was a problem hiding this comment.
I plan to revert this once the buffer deallocation pass is working properly with the sparse compiler. The plan is to fix all memory leaks first, so that I can easily spot issues/regressions while working on the buffer deallocation pass.
I could move the op to Dialect/Test/TestOps.td in the mean time and rename it to test.sparse_tensor.has_runtime_library. But so far we don't seem to have dialect-specific test ops in there. Alternatively, I could also rename it to sparse_tensor.test.has_runtime_library.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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 usessparse_tensor.assembleto 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:
sparse_tensor.assemblecorrectly. 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.sparse_tensor.runtime_onlyattribute tobufferization.dealloc_tensor. Verifying that the attribute can only be attached tobufferization.dealloc_tensormay introduce an unwanted dependency ofMLIRSparseTensorDialectonMLIRBufferizationDialect.