[mlir][vector] Relax transfer_read/transfer_write restriction on memref operand

[ThomasRaoux]

Summary

Relax the verifier for transfer_read/transfer_write operation so that it can take a memref with a different element type than the vector being read/written.
This is based on the discourse discussion:
https://llvm.discourse.group/t/memref-cast/1514

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

1516–1517

It looks like this TODO can be remove too now, since you removed the suffix part?

[aartbik]

also, with the restrictions out of the way, did you make sure this does not crash the lowering to llvm part (try out a few of your new examples). you may have to return failure() in the lowering logic itself if it breaks anything....

[ThomasRaoux]

In D85244#2194729, @aartbik wrote:

also, with the restrictions out of the way, did you make sure this does not crash the lowering to llvm part (try out a few of your new examples). you may have to return failure() in the lowering logic itself if it breaks anything....

It looks like it mostly works out of the box. The translation to llvm does a pointer bitcast so it doesn't seem to rely on type matching. I added a test for conversion to LLVM.

[aartbik]

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

925

ah, this seems to be in the old style (note that our llvm type syntax has changed this morning)

You will need to rebase this CL, otherwise you will get conflicts trying to submit this

[ThomasRaoux]

rebase

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

925

Thanks for the heads up. I rebased.

[aartbik]

We probably want to add some sort of sanity checking in later to make sure we don't do very strange bitcasts, but this is good for now to get the ball rolling.

[nicolasvasilache]

The relaxing is too brutal IMO and will likely create multiple issues related to lowering to SCF and LLVM.
To reduce the pain in making this work with transformations I'd suggest to add restrictions on the type.
Basically you're only allowed to:

1. take elements immediately to the left of your element type in the memref and stick these into the vector.
2. cast to vector types that have the same number fo bits and a power of 2 size.

Anything else will be a significant amount of pain until we have a solid DataLayout in MLIR.

So basically you can allow the existing semantic + exactly 1 vector type_cast that changes vector type.

memref<axbxcxdxvector<txuxvxf32>> -> memref<axbx | cxdxvector<txuxvxf32>> -> memref<axbx vector<cxdxtxuxvxf32>> -> memref<axbx vector<whatever>> -> vector<whatever>, where:

1. | is a separator that you can put anywhere on the memref (just like the current vector.type_cast)
2. whatever has the same bitsize and alignment as cxdxtxuxvxf32.

To make 2. reasonably simple I'd say just allow the most minor of whatever and v x f32 that have a power of 2 size if you want to cast the element type.

I'd recommend very slowly and very carefully relaxing the constraints to match your needs without stepping out of bounds for now.

[ThomasRaoux]

In D85244#2197244, @nicolasvasilache wrote:

The relaxing is too brutal IMO and will likely create multiple issues related to lowering to SCF and LLVM.
To reduce the pain in making this work with transformations I'd suggest to add restrictions on the type.
Basically you're only allowed to:

1. take elements immediately to the left of your element type in the memref and stick these into the vector.
2. cast to vector types that have the same number fo bits and a power of 2 size.

Anything else will be a significant amount of pain until we have a solid DataLayout in MLIR.

So basically you can allow the existing semantic + exactly 1 vector type_cast that changes vector type.

memref<axbxcxdxvector<txuxvxf32>> -> memref<axbx | cxdxvector<txuxvxf32>> -> memref<axbx vector<cxdxtxuxvxf32>> -> memref<axbx vector<whatever>> -> vector<whatever>, where:

1. | is a separator that you can put anywhere on the memref (just like the current vector.type_cast)
2. whatever has the same bitsize and alignment as cxdxtxuxvxf32.

To make 2. reasonably simple I'd say just allow the most minor of whatever and v x f32 that have a power of 2 size if you want to cast the element type.

I'd recommend very slowly and very carefully relaxing the constraints to match your needs without stepping out of bounds for now.

Those rules are not enough for the cases I'm trying to solve. For instance I want to support:
vector.transfer_read ... : memref<128x2xvector<4xi32>>, vector<16x32xi8> and
vector.transfer_read ... : memref<4096x1024xvector<4xi32>>, vector<16x16xi32>
vector.transfer_read ... : memref<32x8xvector<4xi32>>, vector<32x16xi8>

What I need is basically to be able to read as if the memref was scalar. One restriction I could add is to make sure that each dimension of the vector is smaller or equal to the associated dimension of the memref and for the last dimension that the size in bits is smaller or equal to the size in bits of the last dimension of the memref.
That only ensure that no dimension of the vector is bigger than the memref though.

Do you have any suggestions on the rules to pick to support those cases?

[ThomasRaoux]

Add extra restrictions.

[ThomasRaoux]

In D85244#2197244, @nicolasvasilache wrote:

The relaxing is too brutal IMO and will likely create multiple issues related to lowering to SCF and LLVM.
To reduce the pain in making this work with transformations I'd suggest to add restrictions on the type.
Basically you're only allowed to:

1. take elements immediately to the left of your element type in the memref and stick these into the vector.
2. cast to vector types that have the same number fo bits and a power of 2 size.

Anything else will be a significant amount of pain until we have a solid DataLayout in MLIR.

So basically you can allow the existing semantic + exactly 1 vector type_cast that changes vector type.

memref<axbxcxdxvector<txuxvxf32>> -> memref<axbx | cxdxvector<txuxvxf32>> -> memref<axbx vector<cxdxtxuxvxf32>> -> memref<axbx vector<whatever>> -> vector<whatever>, where:

1. | is a separator that you can put anywhere on the memref (just like the current vector.type_cast)
2. whatever has the same bitsize and alignment as cxdxtxuxvxf32.

To make 2. reasonably simple I'd say just allow the most minor of whatever and v x f32 that have a power of 2 size if you want to cast the element type.

I'd recommend very slowly and very carefully relaxing the constraints to match your needs without stepping out of bounds for now.

I added a restriction to enforce that the inner dimension of the destination vector is a multiple of the inner dimension of the memref element type. This way the transfer_read/transfer_write can always be lowered to several load of the memref element type.

[nicolasvasilache]

Cool, just as we discussed offline.

Could you please make sure conversion to LLVM barfs when seeing this, we don't want to miscompile and we likely will right now.

mlir/lib/Dialect/Vector/VectorOps.cpp

1509–1510

We need a better error message here and below.
How about:

requires the bitwidth of the minor 1-D vector to be an integral multiple of the bitwidth of the minor 1-D vector of the memref

?

1520

Please turn this into an assert into the ConvertVectorToLLVM pattern that wants to lower that to LLVM.
Normally we would go with a failure to convert but this is so pernicious that I want it to fail very loudly.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorOps.cpp

1528–1529

We need a better error message here and above.
How about:

requires the bitwidth of the minor 1-D vector to be an integral multiple of the bitwidth of the memref element type

?

[ThomasRaoux]

Address review comments

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

if (xferOp.permutation_map() !=
       return failure();
 
     auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); };
 
     Location loc = op->getLoc();
     Type i64Type = rewriter.getIntegerType(64);
     MemRefType memRefType = xferOp.getMemRefType();
 
+    if (auto memrefVectorElementType =
+            memRefType.getElementType().dyn_cast<VectorType>()) {
+      // Memref has vector element type.
+      if (memrefVectorElementType.getElementType() !=
+          xferOp.getVectorType().getElementType())
+        return failure();
+      // Check that memref vector type is a suffix of 'vectorType.
+      unsigned memrefVecEltRank = memrefVectorElementType.getRank();
+      unsigned resultVecRank = xferOp.getVectorType().getRank();
+      assert(memrefVecEltRank <= resultVecRank);
+      // TODO: Move this to isSuffix in Vector/Utils.h.
+      unsigned rankOffset = resultVecRank - memrefVecEltRank;
+      auto memrefVecEltShape = memrefVectorElementType.getShape();
+      auto resultVecShape = xferOp.getVectorType().getShape();
+      for (unsigned i = 0; i < memrefVecEltRank; ++i)
+        assert(memrefVecEltShape[i] != resultVecShape[rankOffset + i] &&
+               "memref vector element shape should match suffix of vector "
+               "result shape.");
+    }
+
     // 1. Get the source/dst address as an LLVM vector pointer.
     //    The vector pointer would always be on address space 0, therefore
     //    addrspacecast shall be used when source/dst memrefs are not on
     //    address space 0.
     // TODO: support alignment when possible.
     Value dataPtr = getDataPtr(loc, memRefType, adaptor.memref(),
                                adaptor.indices(), rewriter);
     auto vecTy =

mlir/lib/Dialect/Vector/VectorOps.cpp

 static LogicalResult verifyTransferOp(Operation *op, MemRefType memrefType,
                                       VectorType vectorType,
                                       AffineMap permutationMap,
                                       ArrayAttr optionalMasked) {
   auto memrefElementType = memrefType.getElementType();
   if (auto memrefVectorElementType = memrefElementType.dyn_cast<VectorType>()) {
     // Memref has vector element type.
 
-    // Check that 'memrefVectorElementType' and vector element types match.
-    if (memrefVectorElementType.getElementType() != vectorType.getElementType())
+    unsigned memrefVecSize = memrefVectorElementType.getElementTypeBitWidth() *
+                             memrefVectorElementType.getShape().back();
+    unsigned resultVecSize =
+        vectorType.getElementTypeBitWidth() * vectorType.getShape().back();
+    if (resultVecSize % memrefVecSize != 0)
       return op->emitOpError(
-          "requires memref and vector types of the same elemental type");
+          "requires the bitwidth of the minor 1-D vector to be an integral "
+          "multiple of the bitwidth of the minor 1-D vector of the memref");
 
-    // Check that memref vector type is a suffix of 'vectorType.
     unsigned memrefVecEltRank = memrefVectorElementType.getRank();
     unsigned resultVecRank = vectorType.getRank();
     if (memrefVecEltRank > resultVecRank)
       return op->emitOpError(
           "requires memref vector element and vector result ranks to match.");
-    // TODO: Move this to isSuffix in Vector/Utils.h.
     unsigned rankOffset = resultVecRank - memrefVecEltRank;
-    auto memrefVecEltShape = memrefVectorElementType.getShape();
-    auto resultVecShape = vectorType.getShape();
-    for (unsigned i = 0; i < memrefVecEltRank; ++i)
-      if (memrefVecEltShape[i] != resultVecShape[rankOffset + i])
-        return op->emitOpError(
-            "requires memref vector element shape to match suffix of "
-            "vector result shape.");
     // Check that permutation map results match 'rankOffset' of vector type.
     if (permutationMap.getNumResults() != rankOffset)
       return op->emitOpError("requires a permutation_map with result dims of "
                              "the same rank as the vector type");
   } else {
     // Memref has scalar element type.
-
-    // Check that memref and vector element types match.
-    if (memrefType.getElementType() != vectorType.getElementType())
+    unsigned resultVecSize =
+        vectorType.getElementTypeBitWidth() * vectorType.getShape().back();
+    if (resultVecSize % memrefElementType.getIntOrFloatBitWidth() != 0)
       return op->emitOpError(
-          "requires memref and vector types of the same elemental type");
+          "requires the bitwidth of the minor 1-D vector to be an integral "
+          "multiple of the bitwidth of the memref element type");
 
     // Check that permutation map results match rank of vector type.
     if (permutationMap.getNumResults() != vectorType.getRank())
       return op->emitOpError("requires a permutation_map with result dims of "
                              "the same rank as the vector type");
   }
 
   if (permutationMap.getNumSymbols() != 0)
static LogicalResult verifyTransferOp(Operation *op, MemRefType memrefType,
   return success();
 }
 
 /// Builder that sets padding to zero.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vector, Value memref, ValueRange indices,
                            AffineMap permutationMap,
                            ArrayRef<bool> maybeMasked) {
-  Type elemType = vector.cast<VectorType>().getElementType();
+  Type elemType = memref.getType().cast<MemRefType>().getElementType();
   Value padding = builder.create<ConstantOp>(result.location, elemType,
                                              builder.getZeroAttr(elemType));
   if (maybeMasked.empty())
     return build(builder, result, vector, memref, indices, permutationMap,
                  padding, ArrayAttr());
   ArrayAttr maskedArrayAttr = builder.getBoolArrayAttr(maybeMasked);
   build(builder, result, vector, memref, indices, permutationMap, padding,
         maskedArrayAttr);
if (memrefVectorElementType != paddingType)
           "requires memref element type and padding type to match.");
 
   } else {
     // Check that 'paddingType' is valid to store in a vector type.
     if (!VectorType::isValidElementType(paddingType))
       return op.emitOpError("requires valid padding vector elemental type");
 
     // Check that padding type and vector element types match.
-    if (paddingType != vectorType.getElementType())
+    if (paddingType != memrefElementType)
       return op.emitOpError(
-          "requires formal padding and vector of the same elemental type");
+          "requires formal padding and memref of the same elemental type");
   }
 
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 /// This is a common class used for patterns of the form
 /// ```

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 //       CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :
 //  CHECK-SAME: (!llvm.ptr<float>, !llvm.i64) -> !llvm.ptr<float>
 //       CHECK: %[[vecPtr:.*]] = llvm.bitcast %[[gep]] :
 //  CHECK-SAME: !llvm.ptr<float> to !llvm.ptr<vec<17 x float>>
 //
 // 2. Rewrite as a load.
 //       CHECK: %[[loaded:.*]] = llvm.load %[[vecPtr]] {alignment = 4 : i64} : !llvm.ptr<vec<17 x float>>
 
+func @transfer_read_1d_cast(%A : memref<?xi32>, %base: index) -> vector<12xi8> {
+  %c0 = constant 0: i32
+  %v = vector.transfer_read %A[%base], %c0 {masked = [false]} :
+    memref<?xi32>, vector<12xi8>
+  return %v: vector<12xi8>
+}
+// CHECK-LABEL: func @transfer_read_1d_cast
+//  CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: !llvm.i64) -> !llvm.vec<12 x i8>
+//
+// 1. Bitcast to vector form.
+//       CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :
+//  CHECK-SAME: (!llvm.ptr<i32>, !llvm.i64) -> !llvm.ptr<i32>
+//       CHECK: %[[vecPtr:.*]] = llvm.bitcast %[[gep]] :
+//  CHECK-SAME: !llvm.ptr<i32> to !llvm.ptr<vec<12 x i8>>
+//
+// 2. Rewrite as a load.
+//       CHECK: %[[loaded:.*]] = llvm.load %[[vecPtr]] {alignment = 4 : i64} : !llvm.ptr<vec<12 x i8>>
+
 func @genbool_1d() -> vector<8xi1> {
   %0 = vector.constant_mask [4] : vector<8xi1>
   return %0 : vector<8xi1>
 }
 // CHECK-LABEL: func @genbool_1d
 // CHECK: %[[C1:.*]] = llvm.mlir.constant(dense<[true, true, true, true, false, false, false, false]> : vector<8xi1>) : !llvm.vec<8 x i1>
 // CHECK: llvm.return %[[C1]] : !llvm.vec<8 x i1>
 

mlir/test/Dialect/Vector/invalid.mlir

 }
 
 // -----
 
 func @test_vector.transfer_read(%arg0: memref<?x?xvector<4x3xf32>>) {
   %c3 = constant 3 : index
   %f0 = constant 0.0 : f32
   %vf0 = splat %f0 : vector<4x3xf32>
-  // expected-error@+1 {{requires memref and vector types of the same elemental type}}
-  %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xi32>
-}
-
-// -----
-
-func @test_vector.transfer_read(%arg0: memref<?x?xvector<4x3xf32>>) {
-  %c3 = constant 3 : index
-  %f0 = constant 0.0 : f32
-  %vf0 = splat %f0 : vector<4x3xf32>
   // expected-error@+1 {{requires memref vector element and vector result ranks to match}}
   %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<4x3xf32>>, vector<3xf32>
 }
 
 // -----
 
-func @test_vector.transfer_read(%arg0: memref<?x?xvector<4x3xf32>>) {
+func @test_vector.transfer_read(%arg0: memref<?x?xvector<6xf32>>) {
   %c3 = constant 3 : index
   %f0 = constant 0.0 : f32
-  %vf0 = splat %f0 : vector<4x3xf32>
-  // expected-error@+1 {{ requires memref vector element shape to match suffix of vector result shape}}
-  %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<4x3xf32>>, vector<1x1x2x3xf32>
+  %vf0 = splat %f0 : vector<6xf32>
+  // expected-error@+1 {{requires the bitwidth of the minor 1-D vector to be an integral multiple of the bitwidth of the minor 1-D vector of the memref}}
+  %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 : memref<?x?xvector<6xf32>>, vector<3xf32>
 }
 
 // -----
 
 func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) {
   %c3 = constant 3 : index
   %f0 = constant 0.0 : f32
   %vf0 = splat %f0 : vector<2x3xf32>

mlir/test/Dialect/Vector/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
 // CHECK-DAG: #[[MAP0:map[0-9]+]] = affine_map<(d0, d1) -> (d0, d1)>
 
 // CHECK-LABEL: func @vector_transfer_ops(
 func @vector_transfer_ops(%arg0: memref<?x?xf32>,
-                          %arg1 : memref<?x?xvector<4x3xf32>>) {
+                          %arg1 : memref<?x?xvector<4x3xf32>>,
+                          %arg2 : memref<?x?xvector<4x3xi32>>) {
   // CHECK: %[[C3:.*]] = constant 3 : index
   %c3 = constant 3 : index
   %cst = constant 3.0 : f32
   %f0 = constant 0.0 : f32
+  %c0 = constant 0 : i32
   %vf0 = splat %f0 : vector<4x3xf32>
+  %v0 = splat %c0 : vector<4x3xi32>
 
   //
   // CHECK: vector.transfer_read
   %0 = vector.transfer_read %arg0[%c3, %c3], %f0 {permutation_map = affine_map<(d0, d1)->(d0)>} : memref<?x?xf32>, vector<128xf32>
   // CHECK: vector.transfer_read
   %1 = vector.transfer_read %arg0[%c3, %c3], %f0 {permutation_map = affine_map<(d0, d1)->(d1, d0)>} : memref<?x?xf32>, vector<3x7xf32>
   // CHECK: vector.transfer_read
   %2 = vector.transfer_read %arg0[%c3, %c3], %cst {permutation_map = affine_map<(d0, d1)->(d0)>} : memref<?x?xf32>,  vector<128xf32>
   // CHECK: vector.transfer_read
   %3 = vector.transfer_read %arg0[%c3, %c3], %cst {permutation_map = affine_map<(d0, d1)->(d1)>} : memref<?x?xf32>,  vector<128xf32>
   // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
   %4 = vector.transfer_read %arg1[%c3, %c3], %vf0 {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
   // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}} {masked = [true, false]} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
   %5 = vector.transfer_read %arg1[%c3, %c3], %vf0 {masked = [true, false]} : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
+  // CHECK: vector.transfer_read %{{.*}}[%[[C3]], %[[C3]]], %{{.*}} : memref<?x?xvector<4x3xi32>>, vector<5x24xi8>
+  %6 = vector.transfer_read %arg2[%c3, %c3], %v0 : memref<?x?xvector<4x3xi32>>, vector<5x24xi8>
+
 
   // CHECK: vector.transfer_write
   vector.transfer_write %0, %arg0[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d0)>} : vector<128xf32>, memref<?x?xf32>
   // CHECK: vector.transfer_write
   vector.transfer_write %1, %arg0[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d1, d0)>} : vector<3x7xf32>, memref<?x?xf32>
   // CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%[[C3]], %[[C3]]] : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
   vector.transfer_write %4, %arg1[%c3, %c3] {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
   // CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%[[C3]], %[[C3]]] : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
   vector.transfer_write %5, %arg1[%c3, %c3] {masked = [true, true]} : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
+  // CHECK: vector.transfer_write %{{.*}}, %{{.*}}[%[[C3]], %[[C3]]] : vector<5x24xi8>, memref<?x?xvector<4x3xi32>>
+  vector.transfer_write %6, %arg2[%c3, %c3] : vector<5x24xi8>, memref<?x?xvector<4x3xi32>>
 
   return
 }
 
 // CHECK-LABEL: @vector_broadcast
 func @vector_broadcast(%a: f32, %b: vector<16xf32>, %c: vector<1x16xf32>, %d: vector<8x1xf32>) -> vector<8x16xf32> {
   // CHECK: vector.broadcast %{{.*}} : f32 to vector<16xf32>
   %0 = vector.broadcast %a : f32 to vector<16xf32>