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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sparse

Author: Ziliang Zhang (ziliangzl)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/227550.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+1) 
- (added) mlir/test/Dialect/SparseTensor/stage_sparse_ops.mlir (+13) 


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

``````````

</details>


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