[Mlir-commits] [mlir] [mlir][sparse] Load the vector dialect in pre-sparsification-rewrite (PR #216491)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sparse

Author: Hamza Qureshi (hamzaqureshi5)

<details>
<summary>Changes</summary>

Fixes #<!-- -->216273

`mlir-opt --pre-sparsification-rewrite` aborts on this input:

```mlir
#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }>
func.func @<!-- -->print(%arg0: !llvm.ptr) {
  %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor<?x?xf32, #sparse>
  sparse_tensor.print %0 : tensor<?x?xf32, #sparse>
  llvm.return
}
```

```
LLVM ERROR: can't create Attribute 'mlir::vector::PrintPunctuationAttr' because
storage uniquer isn't initialized: the dialect was likely not loaded, ...
```

`PrintRewriter` lowers `sparse_tensor.print` into `vector.print` operations, but `PreSparsificationRewrite` does not list the vector dialect in `dependentDialects`. When the input does not use the dialect either, it is never loaded and building the punctuation attribute aborts.

The pass already declares the other dialects it needs, so this just adds the missing one.

Testing: added a `sparse_tensor.print` case to `mlir/test/Dialect/SparseTensor/pre_rewriting.mlir`, which does not otherwise mention the vector dialect. The test fails on main and passes with the change, and `mlir/test/Dialect/SparseTensor` stays at 113/113.

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


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td (+1) 
- (modified) mlir/test/Dialect/SparseTensor/pre_rewriting.mlir (+12) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 207fab1366294..c4c83819f1e0e 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -105,6 +105,7 @@ def PreSparsificationRewrite : Pass<"pre-sparsification-rewrite", "ModuleOp"> {
     "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
+    "vector::VectorDialect",
   ];
 }
 
diff --git a/mlir/test/Dialect/SparseTensor/pre_rewriting.mlir b/mlir/test/Dialect/SparseTensor/pre_rewriting.mlir
index 5fdbb46e3d8b0..79c0bb93802e5 100644
--- a/mlir/test/Dialect/SparseTensor/pre_rewriting.mlir
+++ b/mlir/test/Dialect/SparseTensor/pre_rewriting.mlir
@@ -99,3 +99,15 @@ func.func @sparse_select(%cond: tensor<4x4xi1>,
   } -> tensor<4x4xf64, #DCSR>
   return %0 : tensor<4x4xf64, #DCSR>
 }
+
+// The print rewrite generates vector.print, so the pass has to load the vector
+// dialect even when the input does not use it.
+//
+// CHECK-LABEL: func.func @sparse_print(
+//  CHECK-SAME: %[[A:.*]]: tensor<?x?xf32, #sparse{{[0-9]*}}>)
+//       CHECK:   vector.print
+//       CHECK:   return
+func.func @sparse_print(%arg0: tensor<?x?xf32, #DCSR>) {
+  sparse_tensor.print %arg0 : tensor<?x?xf32, #DCSR>
+  return
+}

``````````

</details>


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