[Mlir-commits] [mlir] a795955 - [mlir] Register tensor dialect for transfer_read conversion

[Kai Sasaki]

Author: Kai Sasaki
Date: 2023-02-01T09:17:08+09:00
New Revision: a795955f5251856bddc3f895c527e084073169c1

URL: https://github.com/llvm/llvm-project/commit/a795955f5251856bddc3f895c527e084073169c1
DIFF: https://github.com/llvm/llvm-project/commit/a795955f5251856bddc3f895c527e084073169c1.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Conversion/Passes.td
    mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
    mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 024c6f4280c29..e8b61b43e1e4c 100644
--- 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 401e227629540..931a85b636f9d 100644
--- 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 ca353a07ad761..dd7e52552f0df 100644
--- 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>
+}