[mlir][vector] Fix parser of vector.transfer_read

Author: douyixuan

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -152,22 +152,18 @@ static bool isSupportedCombiningKind(CombiningKind combiningKind,
 
 AffineMap mlir::vector::getTransferMinorIdentityMap(ShapedType shapedType,
                                                     VectorType vectorType) {
+  int64_t elementVectorRank = 0;
+  VectorType elementVectorType =
+      llvm::dyn_cast<VectorType>(shapedType.getElementType());
+  if (elementVectorType)
+    elementVectorRank += elementVectorType.getRank();
   // 0-d transfers are to/from tensor<t>/memref<t> and vector<1xt>.
   // TODO: replace once we have 0-d vectors.
   if (shapedType.getRank() == 0 &&
       vectorType.getShape() == ArrayRef<int64_t>{1})
     return AffineMap::get(
         /*numDims=*/0, /*numSymbols=*/0,
         getAffineConstantExpr(0, shapedType.getContext()));
-  int64_t elementVectorRank = 0;
-  VectorType elementVectorType =
-      llvm::dyn_cast<VectorType>(shapedType.getElementType());
-  if (elementVectorType)
-    elementVectorRank += elementVectorType.getRank();
-  if (shapedType.getRank() < vectorType.getRank() - elementVectorRank) {
-    // Not enough dimensions in the shaped type to form a minor identity map.
-    return AffineMap();
-  }
   return AffineMap::getMinorIdentityMap(
       shapedType.getRank(), vectorType.getRank() - elementVectorRank,
       shapedType.getContext());
@@ -4263,18 +4259,16 @@ ParseResult TransferReadOp::parse(OpAsmParser &parser, OperationState &result) {
   Attribute permMapAttr = result.attributes.get(permMapAttrName);
   AffineMap permMap;
   if (!permMapAttr) {
+    int64_t elementVectorRank = 0;
+    VectorType elementVectorType =
+        llvm::dyn_cast<VectorType>(shapedType.getElementType());
+    if (elementVectorType)
+      elementVectorRank += elementVectorType.getRank();
+    if (shapedType.getRank() < vectorType.getRank() - elementVectorRank)
+      return parser.emitError(typesLoc,
+                              "expected a custom permutation_map when "
+                              "rank(source) != rank(destination)");
     permMap = getTransferMinorIdentityMap(shapedType, vectorType);
-    if (!permMap) {
-      int64_t elementVectorRank = 0;
-      VectorType elementVectorType =
-          llvm::dyn_cast<VectorType>(shapedType.getElementType());
-      if (elementVectorType)
-        elementVectorRank += elementVectorType.getRank();
-      if (shapedType.getRank() < vectorType.getRank() - elementVectorRank)
-        return parser.emitError(typesLoc,
-                                "expected a custom permutation_map when source "
-                                "rank is less than required for vector rank");
-    }
     result.attributes.set(permMapAttrName, AffineMapAttr::get(permMap));
   } else {
     permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
@@ -4682,18 +4676,16 @@ ParseResult TransferWriteOp::parse(OpAsmParser &parser,
   auto permMapAttr = result.attributes.get(permMapAttrName);
   AffineMap permMap;
   if (!permMapAttr) {
+    int64_t elementVectorRank = 0;
+    VectorType elementVectorType =
+        llvm::dyn_cast<VectorType>(shapedType.getElementType());
+    if (elementVectorType)
+      elementVectorRank += elementVectorType.getRank();
+    if (shapedType.getRank() < vectorType.getRank() - elementVectorRank)
+      return parser.emitError(typesLoc,
+                              "expected a custom permutation_map when "
+                              "rank(source) != rank(destination)");
     permMap = getTransferMinorIdentityMap(shapedType, vectorType);
-    if (!permMap) {
-      int64_t elementVectorRank = 0;
-      VectorType elementVectorType =
-          llvm::dyn_cast<VectorType>(shapedType.getElementType());
-      if (elementVectorType)
-        elementVectorRank += elementVectorType.getRank();
-      if (shapedType.getRank() < vectorType.getRank() - elementVectorRank)
-        return parser.emitError(typesLoc,
-                                "expected a custom permutation_map when result "
-                                "rank is less than required for vector rank");
-    }
     result.attributes.set(permMapAttrName, AffineMapAttr::get(permMap));
   } else {
     permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();

mlir/test/Dialect/Vector/invalid.mlir

@@ -527,13 +527,22 @@ func.func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) {
 
 func.func @test_vector.transfer_read(%arg1: memref<?xindex>) -> vector<3x4xi32> {
   %c3 = arith.constant 3 : index
-  // expected-error@+1 {{expected a custom permutation_map when source rank is less than required for vector rank}}
+  // expected-error@+1 {{expected a custom permutation_map when rank(source) != rank(destination)}}
   %0 = vector.transfer_read %arg1[%c3, %c3], %c3 : memref<?xindex>, vector<3x4xi32>
   return %0 : vector<3x4xi32>
 }
 
 // -----
 
+func.func @test_vector.transfer_read(%arg1: memref<?xindex>) -> vector<3x4xindex> {
+  %c3 = arith.constant 3 : index
+  // expected-error@+1 {{expected a custom permutation_map when rank(source) != rank(destination)}}
+  %0 = vector.transfer_read %arg1[%c3, %c3], %c3 : memref<?xindex>, vector<3x4xindex>
+  return %0 : vector<3x4xindex>
+}
+
+// -----
+
 func.func @test_vector.transfer_write(%arg0: memref<?x?xf32>) {
   %c3 = arith.constant 3 : index
   %cst = arith.constant 3.0 : f32
@@ -657,12 +666,20 @@ func.func @test_vector.transfer_write(%arg0: memref<?xf32>, %arg1: vector<7xf32>
 
 func.func @test_vector.transfer_write(%vec_to_write: vector<3x4xi32>, %output_memref: memref<?xindex>) {
   %c3 = arith.constant 3 : index
-  // expected-error@+1 {{expected a custom permutation_map when result rank is less than required for vector rank}}
+  // expected-error@+1 {{expected a custom permutation_map when rank(source) != rank(destination)}}
   vector.transfer_write %vec_to_write, %output_memref[%c3, %c3] : vector<3x4xi32>, memref<?xindex>
 }
 
 // -----
 
+func.func @test_vector.transfer_write(%vec_to_write: vector<3x4xindex>, %output_memref: memref<?xindex>) {
+  %c3 = arith.constant 3 : index
+  // expected-error@+1 {{expected a custom permutation_map when rank(source) != rank(destination)}}
+  vector.transfer_write %vec_to_write, %output_memref[%c3, %c3] : vector<3x4xindex>, memref<?xindex>
+}
+
+// -----
+
 func.func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected offsets of same size as destination vector rank}}
   %1 = vector.insert_strided_slice %a, %b {offsets = [100], strides = [1, 1]} : vector<4x4xf32> into vector<4x8x16xf32>