[Mlir-commits] [mlir] e00ff44 - [mlir][sparse] Register bufferization dialect in stage pass (#227550)

[llvmlistbot at llvm.org]

Author: Ziliang Zhang
Date: 2026-10-01T09:50:23+09:00
New Revision: e00ff4440bb6c30621b33e3b90f92547f79a71ff

URL: https://github.com/llvm/llvm-project/commit/e00ff4440bb6c30621b33e3b90f92547f79a71ff
DIFF: https://github.com/llvm/llvm-project/commit/e00ff4440bb6c30621b33e3b90f92547f79a71ff.diff

LOG: [mlir][sparse] Register bufferization dialect in stage pass (#227550)

StageSparseOperations may create bufferization.dealloc_tensor while
staging conversions. Register the dialect as a pass dependency and add a
standalone regression test.

Added: 
    mlir/test/Dialect/SparseTensor/stage_sparse_ops.mlir

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index c4c83819f1e0e..0960254a31469 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -201,6 +201,7 @@ def StageSparseOperations : Pass<"stage-sparse-ops", "func::FuncOp"> {
   }];
   let constructor = "mlir::createStageSparseOperationsPass()";
   let dependentDialects = [
+    "bufferization::BufferizationDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
 }

diff  --git a/mlir/test/Dialect/SparseTensor/stage_sparse_ops.mlir b/mlir/test/Dialect/SparseTensor/stage_sparse_ops.mlir
new file mode 100644
index 0000000000000..2cfc5a5b2cb75
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/stage_sparse_ops.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s --stage-sparse-ops | FileCheck %s
+
+#SparseTensor = #sparse_tensor.encoding<{
+  map = (d0, d1, d2) -> (d2 : dense, d0 : compressed, d1 : compressed)
+}>
+
+// CHECK-LABEL: func.func @sparse_convert_3d
+// CHECK:         %[[TMP:.*]] = sparse_tensor.reorder_coo
+// CHECK:         bufferization.dealloc_tensor %[[TMP]]
+func.func @sparse_convert_3d(%arg0: tensor<?x?x?xf64>) -> tensor<?x?x?xf64, #SparseTensor> {
+  %0 = sparse_tensor.convert %arg0 : tensor<?x?x?xf64> to tensor<?x?x?xf64, #SparseTensor>
+  return %0 : tensor<?x?x?xf64, #SparseTensor>
+}