[Mlir-commits] [mlir] [mlir][SparseTensor] Add missing dependent dialect to pass (PR #88870)

[Matthias Springer]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/88870

>From 4a35018bcb9303b016342fa262b5ec0d992ad131 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Tue, 16 Apr 2024 14:47:15 +0000
Subject: [PATCH] [mlir][SparseTensor] Add missing dependent dialect to pass

This commit fixes the following error when stopping the sparse compiler pipeline after bufferization (e.g., with `test-analysis-only`):

```
LLVM ERROR: Building op `vector.print` but it isn't known in this MLIRContext: the dialect may not be loaded or this operation hasn't been added by the dialect. See also https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management
```
---
 mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td   | 1 +
 .../Transforms/SparsificationAndBufferizationPass.cpp         | 1 +
 .../Integration/Dialect/SparseTensor/CPU/sparse_block3d.mlir  | 4 ++++
 3 files changed, 6 insertions(+)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 4706d5ba2f218c..2f844cee5ff528 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -460,6 +460,7 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
+    "vector::VectorDialect"
   ];
 }
 
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index f497be6e48eba1..3a8972072ac3b1 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -24,6 +24,7 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_block3d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_block3d.mlir
index 2ff73923c8327d..467b671500e173 100755
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_block3d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_block3d.mlir
@@ -30,6 +30,10 @@
 // Do the same run, but now with direct IR generation and VLA vectorization.
 // RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %}
 
+// Test that test-bufferization-analysis-only works. This option is useful
+// for understanding why buffer copies were inserted.
+// RUN: mlir-opt %s --sparsifier="test-bufferization-analysis-only" -o /dev/null
+
 #Sparse1 = #sparse_tensor.encoding<{
   map = (i, j, k) -> (
     j : compressed,