[mlir][sparse] add create-sparse-deallocs options to match the create-deallocs in BufferizationOption.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

78–81

Should document how this option interacts with enableRuntimeLibrary. That is, I think this is only used/needed for the codegen path rather than the runtime path?

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

78

by the sparse compiler

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

224

by the sparse compiler

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

1521

how about not adding rule when disabled?

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

185–188

Any particular reason you're putting createSparseDeallocs there, rather than at the end?

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

1521

my mistake, we still erase...

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

78–81

Yeah, you are right.

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

185–188

by alphabet order (c < e)?

[Peiming]

address comments.

[aartbik]

mlir/test/Dialect/SparseTensor/codegen_sparse_dealloc.mlir

17

user-requested

18

can you add a

CHECK-NO-DEALLOC-LABEL: @sparse_convert_permuted

and

CHECK-DEALLOC-LABEL: @sparse_convert_permuted

It clearly works without, but it feels a bit cleaner to avoid matching output elsewhere if this every changes

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

185–188

fair enough :)

[Peiming]

address comments.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

78–83

I would rephrase this to "Specifies whether to deallocate the temporary sparse buffers created by the sparse compiler when enable-runtime-library=false. For compatibility with core bufferization passes. See also..."

mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h

struct SparseCompilerOptions
   PassOptions::Option<bool> testBufferizationAnalysisOnly{
       *this, "test-bufferization-analysis-only",
       desc("Run only the inplacability analysis"), init(false)};
 
   PassOptions::Option<bool> enableBufferInitialization{
       *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)};
 
   // These options must be kept in sync with `SparseTensorConversionBase`.
   PassOptions::Option<int32_t> sparseToSparse{
       *this, "s2s-strategy",
       desc("Set the strategy for sparse-to-sparse conversion"), init(0)};

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

 class SparseTensorTypeToBufferConverter : public TypeConverter {
 public:
   SparseTensorTypeToBufferConverter();
 };
 
 /// 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.
 //===----------------------------------------------------------------------===//
 
 void populatePreSparsificationRewriting(RewritePatternSet &patterns);
 
 std::unique_ptr<Pass> createPreSparsificationRewritePass();

 //===----------------------------------------------------------------------===//
 // Other rewriting rules and passes.
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
     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);
 
 std::unique_ptr<Pass> createSparseBufferRewritePass();
 std::unique_ptr<Pass>
 createSparseBufferRewritePass(bool enableBufferInitialization);
 
 void populateSparseVectorizationPatterns(RewritePatternSet &patterns,

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

let dependentDialects = [
     "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
   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.">,
   ];
 }
 
 def SparseBufferRewrite : Pass<"sparse-buffer-rewrite", "ModuleOp"> {
   let summary = "Rewrite sparse primitives on buffers to actual code";
   let description = [{
     A pass that rewrites sparse primitives on buffers to the MLIR implementation
     of the primitives. For example, sparse_tensor.sort operator is implemented

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass());
   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)
     return;
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
   // TODO(springerm): Add sparse support to the BufferDeallocation pass and add

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

private:
   bool enableBufferInitialization;
 };
 
 /// Sparse codegen rule for the dealloc operator.
 class SparseTensorDeallocConverter
     : 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 {
     auto enc = getSparseTensorEncoding(op.getTensor().getType());
     if (!enc)
       return failure();
 
+    // 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.
 class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(LoadOp op, OpAdaptor adaptor,

 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// 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,
                                             StorageSpecifierKind::DimStride>,
                SparseToPositionsConverter, SparseToCoordinatesConverter,
                SparseToCoordinatesBufferConverter, SparseToValuesConverter,
                SparseConvertConverter, SparseNewOpConverter,
                SparseNumberOfEntriesConverter>(typeConverter,
                                                patterns.getContext());
+  patterns.add<SparseTensorDeallocConverter>(
+      typeConverter, patterns.getContext(), createSparseDeallocs);
   patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                            enableBufferInitialization);
 }

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

struct SparseTensorConversionPass
   }
 };
 
 struct SparseTensorCodegenPass
     : public impl::SparseTensorCodegenBase<SparseTensorCodegenPass> {
 
   SparseTensorCodegenPass() = default;
   SparseTensorCodegenPass(const SparseTensorCodegenPass &pass) = default;
-  SparseTensorCodegenPass(bool enableInit) {
+  SparseTensorCodegenPass(bool createDeallocs, bool enableInit) {
+    createSparseDeallocs = createDeallocs;
     enableBufferInitialization = enableInit;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseTensorTypeToBufferConverter converter;
     ConversionTarget target(*ctx);
target.addLegalDialect<
         arith::ArithDialect, bufferization::BufferizationDialect,
         complex::ComplexDialect, memref::MemRefDialect, scf::SCFDialect>();
     target.addLegalOp<UnrealizedConversionCastOp>();
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    converter);
     scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
                                                          target);
-    populateSparseTensorCodegenPatterns(converter, patterns,
-                                        enableBufferInitialization);
+    populateSparseTensorCodegenPatterns(
+        converter, patterns, createSparseDeallocs, enableBufferInitialization);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
 struct SparseBufferRewritePass
     : public impl::SparseBufferRewriteBase<SparseBufferRewritePass> {
std::unique_ptr<Pass> mlir::createSparseTensorConversionPass(
   return std::make_unique<SparseTensorConversionPass>(options);
 }
 
 std::unique_ptr<Pass> mlir::createSparseTensorCodegenPass() {
   return std::make_unique<SparseTensorCodegenPass>();
 }
 
 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() {
   return std::make_unique<SparseBufferRewritePass>();
 }
 
 std::unique_ptr<Pass>
 mlir::createSparseBufferRewritePass(bool enableBufferInitialization) {

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

 class SparsificationAndBufferizationPass
     : public PassWrapper<SparsificationAndBufferizationPass,
                          OperationPass<ModuleOp>> {
 public:
   SparsificationAndBufferizationPass(
       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),
         enableVLAVectorization(enableVLAVectorization),
         enableSIMDIndex32(enableSIMDIndex32) {}
 
   /// Bufferize all dense ops. This assumes that no further analysis is needed
   /// and that all required buffer copies were already inserted by
// of `bufferization.alloc_tensor` ops.
         pm.addPass(mlir::createLoopInvariantCodeMotionPass());
         pm.addPass(createSparseVectorizationPass(
             vectorLength, enableVLAVectorization, enableSIMDIndex32));
       }
       if (enableRuntimeLibrary) {
         pm.addPass(
             createSparseTensorConversionPass(sparseTensorConversionOptions));
       } else {
-        pm.addPass(createSparseTensorCodegenPass(enableBufferInitialization));
+        pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,
+                                                 enableBufferInitialization));
         pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
         pm.addPass(createStorageSpecifierToLLVMPass());
       }
       if (failed(runPipeline(pm, getOperation())))
         return signalPassFailure();
     }
 
     // Bufferize all dense ops.
     if (failed(runDenseBufferization()))
       signalPassFailure();
   }
 
 private:
   bufferization::OneShotBufferizationOptions bufferizationOptions;
   SparsificationOptions sparsificationOptions;
   SparseTensorConversionOptions sparseTensorConversionOptions;
+  bool createSparseDeallocs;
   bool enableRuntimeLibrary;
   bool enableBufferInitialization;
   unsigned vectorLength;
   bool enableVLAVectorization;
   bool enableSIMDIndex32;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 std::unique_ptr<Pass> mlir::createSparsificationAndBufferizationPass(
     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);
 }

mlir/test/Dialect/SparseTensor/codegen_sparse_dealloc.mlir

+// 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>
+}