[mlir] Register tensor dialect for transfer_read conversion

Make sure to register tensor dialect as tensor.transfer_read can be dependent on its parameter. It resolves the issue causing the unregisterd dialect error to convert vector to SCF reported https://github.com/llvm/llvm-project/issues/60197.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D142866

diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 024c6f4..e8b61b4 100644 (file)
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -946,7 +946,8 @@ def ConvertVectorToSCF : Pass<"convert-vector-to-scf"> {
   let dependentDialects = [
     "AffineDialect",
     "memref::MemRefDialect",
-    "scf::SCFDialect"
+    "scf::SCFDialect",
+    "tensor::TensorDialect"
   ];
   let options = [
     Option<"fullUnroll", "full-unroll", "bool", /*default=*/"false",

diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
index 401e227..931a85b 100644 (file)
--- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
+++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
@@ -19,6 +19,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"

diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
index ca353a0..dd7e525 100644 (file)
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -498,3 +498,16 @@ func.func @transfer_read_within_async_execute(%A : memref<2x2xf32>) -> !async.to
   }
   return %token : !async.token
 }
+
+// -----
+
+// CHECK-LABEL: transfer_read_with_tensor
+func.func @transfer_read_with_tensor(%arg: tensor<f32>) -> vector<1xf32> {
+    // CHECK:      %[[EXTRACTED:.*]] = tensor.extract %{{.*}}[] : tensor<f32>
+    // CHECK-NEXT: %[[RESULT:.*]] = vector.broadcast %[[EXTRACTED]] : f32 to vector<1xf32>
+    // CHECK-NEXT: return %[[RESULT]] : vector<1xf32>
+    %f0 = arith.constant 0.0 : f32
+    %0 = vector.transfer_read %arg[], %f0 {permutation_map = affine_map<()->(0)>} :
+      tensor<f32>, vector<1xf32>
+    return %0: vector<1xf32>
+}