[mlir][vector] Fold inbound attribute for transfer op with permutations

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorOps.cpp

2643

s/Broadcassted/Broadcasted

[nicolasvasilache]

do you still plan on landing this?

[ThomasRaoux]

In D103253#2909607, @nicolasvasilache wrote:

do you still plan on landing this?

No turns out Matthias had a fix in flight already. I'm closing this patch.

mlir/lib/Dialect/Vector/VectorOps.cpp

static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
   int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
 
   return cstOp.getValue() + vectorSize <= sourceSize;
 }
 
 template <typename TransferOp>
 static LogicalResult foldTransferInBoundsAttribute(TransferOp op) {
   AffineMap permutationMap = op.permutation_map();
-  if (!permutationMap.isMinorIdentity())
-    return failure();
   bool changed = false;
   SmallVector<bool, 4> newInBounds;
   newInBounds.reserve(op.getTransferRank());
-  op.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
+  for (int64_t resultIdx = 0, eResult = op.getTransferRank();
+       resultIdx < eResult; ++resultIdx) {
     // Already marked as in-bounds, nothing to see here.
     if (op.isDimInBounds(resultIdx)) {
       newInBounds.push_back(true);
-      return;
+      continue;
     }
+    // Broadcassted dimensions are always in bound.
+    if (permutationMap.getResult(resultIdx).isa<AffineConstantExpr>()) {
+      newInBounds.push_back(true);
+      changed = true;
+      continue;
+    }
+    int64_t indicesIdx = permutationMap.getDimPosition(resultIdx);
     // Currently out-of-bounds, check whether we can statically determine it is
     // inBounds.
     auto inBounds = isInBounds(op, resultIdx, indicesIdx);
     newInBounds.push_back(inBounds);
     // We commit the pattern if it is "more inbounds".
     changed |= inBounds;
-  });
+  }
   if (!changed)
     return failure();
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(op.getContext());
   op->setAttr(TransferOp::getInBoundsAttrName(),
               b.getBoolArrayAttr(newInBounds));
   return success();
 }

mlir/test/Dialect/Vector/canonicalize.mlir

func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>,
   %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32>
   %r = vector.extract %0[1] : vector<2x4x2xf32>
   return %r : vector<4x2xf32>
 }
 
 
 // -----
 
+#map0 = affine_map<(d0, d1, d2, d3) -> (d3, 0, d1)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)>
+
 // CHECK-LABEL: fold_vector_transfers
-func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
+func @fold_vector_transfers(%A: memref<?x8xf32>, %B: memref<4x3x5x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
   %c0 = constant 0 : index
   %f0 = constant 0.0 : f32
 
   // CHECK: vector.transfer_read %{{.*}} {in_bounds = [false, true]}
   %1 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x8xf32>, vector<4x8xf32>
 
   // CHECK: vector.transfer_write %{{.*}} {in_bounds = [false, true]}
   vector.transfer_write %1, %A[%c0, %c0] : vector<4x8xf32>, memref<?x8xf32>
 
   // Both dims may be out-of-bounds, attribute is elided.
   // CHECK: vector.transfer_read %{{.*}}
   // CHECK-NOT: in_bounds
   %2 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x8xf32>, vector<4x9xf32>
 
   // Both dims may be out-of-bounds, attribute is elided.
   // CHECK: vector.transfer_write %{{.*}}
   // CHECK-NOT: in_bounds
   vector.transfer_write %2, %A[%c0, %c0] : vector<4x9xf32>, memref<?x8xf32>
 
+  // CHECK: vector.transfer_read %{{.*}} {in_bounds = [false, true, true]
+  %3 = vector.transfer_read %B[%c0, %c0, %c0, %c0], %f0 {permutation_map = #map0} : memref<4x3x5x8xf32>, vector<9x12x3xf32>
+
+  // CHECK: vector.transfer_write %{{.*}} {in_bounds = [false, false, true]
+  vector.transfer_write %3, %B[%c0, %c0, %c0, %c0] {permutation_map = #map1} : vector<9x12x3xf32>, memref<4x3x5x8xf32>
+
   // CHECK: return
   return %1, %2 : vector<4x8xf32>, vector<4x9xf32>
 }
 
 // -----
 
 // CHECK-LABEL: bitcast_folding
 //  CHECK-SAME:   %[[A:.*]]: vector<4x8xf32>