*this, "enable-buffer-initialization",
desc("Enable zero-initialization of memory buffers"), init(false)};
+ PassOptions::Option<bool> createSparseDeallocs{
+ *this, "create-sparse-deallocs",
+ desc("Specify if the temporary sparse buffer created by the sparse "
+ "compiler should be deallocated. For compatibility with core "
+ "bufferization passes. "
+ "It only takes effect when enable-runtime-library=false, otherwise "
+ "the memory storage for sparse tensors are managed by the runtime "
+ "library. See also create-deallocs for BufferizationOption."),
+ init(true)};
+
PassOptions::Option<int32_t> vectorLength{
*this, "vl", desc("Set the vector length (0 disables vectorization)"),
init(0)};
/// Sets up sparse tensor conversion rules.
void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
RewritePatternSet &patterns,
+ bool createSparseDeallocs,
bool enableBufferInitialization);
std::unique_ptr<Pass> createSparseTensorCodegenPass();
std::unique_ptr<Pass>
-createSparseTensorCodegenPass(bool enableBufferInitialization);
+createSparseTensorCodegenPass(bool createSparseDeallocs,
+ bool enableBufferInitialization);
//===----------------------------------------------------------------------===//
// The PreSparsificationRewriting pass.
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions,
- bool enableRuntimeLibrary, bool enableBufferInitialization,
- unsigned vectorLength, bool enableVLAVectorization, bool enableSIMDIndex32);
+ bool createSparseDeallocs, bool enableRuntimeLibrary,
+ bool enableBufferInitialization, unsigned vectorLength,
+ bool enableVLAVectorization, bool enableSIMDIndex32);
void populateSparseBufferRewriting(RewritePatternSet &patterns,
bool enableBufferInitialization);
let options = [
Option<"enableBufferInitialization", "enable-buffer-initialization", "bool",
"false", "Enable zero-initialization of the memory buffers">,
+ Option<"createSparseDeallocs", "create-sparse-deallocs", "bool",
+ "true", "Specify if the temporary sparse buffer created by the sparse "
+ "compiler should be deallocated. For compatibility with core "
+ "bufferization passes. "
+ "It only takes effect when enable-runtime-library=false, otherwise "
+ "the memory storage for sparse tensors are managed by the runtime "
+ "library. See also create-deallocs for BufferizationOption.">,
];
}
pm.addPass(createSparsificationAndBufferizationPass(
getBufferizationOptions(options.testBufferizationAnalysisOnly),
options.sparsificationOptions(), options.sparseTensorConversionOptions(),
- options.enableRuntimeLibrary, options.enableBufferInitialization,
- options.vectorLength,
+ options.createSparseDeallocs, options.enableRuntimeLibrary,
+ options.enableBufferInitialization, options.vectorLength,
/*enableVLAVectorization=*/options.armSVE,
/*enableSIMDIndex32=*/options.force32BitVectorIndices));
if (options.testBufferizationAnalysisOnly)
: public OpConversionPattern<bufferization::DeallocTensorOp> {
public:
using OpConversionPattern::OpConversionPattern;
+ SparseTensorDeallocConverter(TypeConverter &typeConverter,
+ MLIRContext *context, bool createDeallocs)
+ : OpConversionPattern(typeConverter, context),
+ createDeallocs(createDeallocs) {}
+
LogicalResult
matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
if (!enc)
return failure();
- // Replace the sparse tensor deallocation with field deallocations.
- Location loc = op.getLoc();
- auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
- for (auto input : desc.getMemRefFields())
- // Deallocate every buffer used to store the sparse tensor handler.
- rewriter.create<memref::DeallocOp>(loc, input);
-
+ // If user requests not to deallocate sparse tensors, simply erase the
+ // operation.
+ if (createDeallocs) {
+ // Replace the sparse tensor deallocation with field deallocations.
+ Location loc = op.getLoc();
+ auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+ for (auto input : desc.getMemRefFields())
+ // Deallocate every buffer used to store the sparse tensor handler.
+ rewriter.create<memref::DeallocOp>(loc, input);
+ }
rewriter.eraseOp(op);
return success();
}
+
+private:
+ bool createDeallocs;
};
/// Sparse codegen rule for tensor rematerialization.
/// the sparsification of linear algebra operations.
void mlir::populateSparseTensorCodegenPatterns(
TypeConverter &typeConverter, RewritePatternSet &patterns,
- bool enableBufferInitialization) {
+ bool createSparseDeallocs, bool enableBufferInitialization) {
patterns.add<SparsePackOpConverter, SparseUnpackOpConverter,
SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
- SparseCastConverter, SparseTensorDeallocConverter,
- SparseExtractSliceConverter, SparseTensorLoadConverter,
- SparseExpandConverter, SparseCompressConverter,
- SparseInsertConverter,
+ SparseCastConverter, SparseExtractSliceConverter,
+ SparseTensorLoadConverter, SparseExpandConverter,
+ SparseCompressConverter, SparseInsertConverter,
SparseSliceGetterOpConverter<ToSliceOffsetOp,
StorageSpecifierKind::DimOffset>,
SparseSliceGetterOpConverter<ToSliceStrideOp,
SparseConvertConverter, SparseNewOpConverter,
SparseNumberOfEntriesConverter>(typeConverter,
patterns.getContext());
+ patterns.add<SparseTensorDeallocConverter>(
+ typeConverter, patterns.getContext(), createSparseDeallocs);
patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
enableBufferInitialization);
}
SparseTensorCodegenPass() = default;
SparseTensorCodegenPass(const SparseTensorCodegenPass &pass) = default;
- SparseTensorCodegenPass(bool enableInit) {
+ SparseTensorCodegenPass(bool createDeallocs, bool enableInit) {
+ createSparseDeallocs = createDeallocs;
enableBufferInitialization = enableInit;
}
converter);
scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
target);
- populateSparseTensorCodegenPatterns(converter, patterns,
- enableBufferInitialization);
+ populateSparseTensorCodegenPatterns(
+ converter, patterns, createSparseDeallocs, enableBufferInitialization);
if (failed(applyPartialConversion(getOperation(), target,
std::move(patterns))))
signalPassFailure();
}
std::unique_ptr<Pass>
-mlir::createSparseTensorCodegenPass(bool enableBufferInitialization) {
- return std::make_unique<SparseTensorCodegenPass>(enableBufferInitialization);
+mlir::createSparseTensorCodegenPass(bool createSparseDeallocs,
+ bool enableBufferInitialization) {
+ return std::make_unique<SparseTensorCodegenPass>(createSparseDeallocs,
+ enableBufferInitialization);
}
std::unique_ptr<Pass> mlir::createSparseBufferRewritePass() {
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions,
- bool enableRuntimeLibrary, bool enableBufferInitialization,
- unsigned vectorLength, bool enableVLAVectorization,
- bool enableSIMDIndex32)
+ bool createSparseDeallocs, bool enableRuntimeLibrary,
+ bool enableBufferInitialization, unsigned vectorLength,
+ bool enableVLAVectorization, bool enableSIMDIndex32)
: bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions),
sparseTensorConversionOptions(sparseTensorConversionOptions),
+ createSparseDeallocs(createSparseDeallocs),
enableRuntimeLibrary(enableRuntimeLibrary),
enableBufferInitialization(enableBufferInitialization),
vectorLength(vectorLength),
pm.addPass(
createSparseTensorConversionPass(sparseTensorConversionOptions));
} else {
- pm.addPass(createSparseTensorCodegenPass(enableBufferInitialization));
+ pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,
+ enableBufferInitialization));
pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
pm.addPass(createStorageSpecifierToLLVMPass());
}
bufferization::OneShotBufferizationOptions bufferizationOptions;
SparsificationOptions sparsificationOptions;
SparseTensorConversionOptions sparseTensorConversionOptions;
+ bool createSparseDeallocs;
bool enableRuntimeLibrary;
bool enableBufferInitialization;
unsigned vectorLength;
const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions,
- bool enableRuntimeLibrary, bool enableBufferInitialization,
- unsigned vectorLength, bool enableVLAVectorization,
- bool enableSIMDIndex32) {
+ bool createSparseDeallocs, bool enableRuntimeLibrary,
+ bool enableBufferInitialization, unsigned vectorLength,
+ bool enableVLAVectorization, bool enableSIMDIndex32) {
return std::make_unique<
mlir::sparse_tensor::SparsificationAndBufferizationPass>(
bufferizationOptions, sparsificationOptions,
- sparseTensorConversionOptions, enableRuntimeLibrary,
+ sparseTensorConversionOptions, createSparseDeallocs, enableRuntimeLibrary,
enableBufferInitialization, vectorLength, enableVLAVectorization,
enableSIMDIndex32);
}
--- /dev/null
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false" \
+// RUN: --sparse-tensor-codegen=create-sparse-deallocs=false \
+// RUN: --canonicalize --cse | FileCheck %s -check-prefix=CHECK-NO-DEALLOC
+
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false" \
+// RUN: --sparse-tensor-codegen=create-sparse-deallocs=true \
+// RUN: --canonicalize --cse | FileCheck %s -check-prefix=CHECK-DEALLOC
+
+#CSR = #sparse_tensor.encoding<{ dimLevelType = ["dense", "compressed"]}>
+#CSC = #sparse_tensor.encoding<{
+ dimLevelType = ["dense", "compressed"],
+ dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+//
+// No memref.dealloc is user-requested so
+// CHECK-NO-DEALLOC-LABEL: @sparse_convert_permuted
+// CHECK-NO-DEALLOC-NOT: memref.dealloc
+//
+// Otherwise memref.dealloc is created to free temporary sparse buffers.
+// CHECK-DEALLOC-LABEL: @sparse_convert_permuted
+// CHECK-DEALLOC: memref.dealloc
+//
+func.func @sparse_convert_permuted(%arg0: tensor<?x?xf32, #CSR>) -> tensor<?x?xf32, #CSC> {
+ %0 = sparse_tensor.convert %arg0 : tensor<?x?xf32, #CSR> to tensor<?x?xf32, #CSC>
+ return %0 : tensor<?x?xf32, #CSC>
+}