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

[Ziliang Zhang]

https://github.com/ziliangzl created https://github.com/llvm/llvm-project/pull/227550

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

>From 57b8a217e457943e8e4e66f7dce7a068b62fbfe4 Mon Sep 17 00:00:00 2001
From: Ziliang Zhang <zzl.coding at gmail.com>
Date: Wed, 30 Sep 2026 11:58:01 +0800
Subject: [PATCH] [mlir][sparse] Register bufferization dialect in stage pass

StageSparseOperations may create bufferization.dealloc_tensor while staging conversions. Register the dialect as a pass dependency and add a standalone regression test.
---
 .../mlir/Dialect/SparseTensor/Transforms/Passes.td  |  1 +
 .../test/Dialect/SparseTensor/stage_sparse_ops.mlir | 13 +++++++++++++
 2 files changed, 14 insertions(+)
 create mode 100644 mlir/test/Dialect/SparseTensor/stage_sparse_ops.mlir

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