Refine how bcast dims are handled

Only mark bcast dims as "in bounds" when all non-bcast dims are "in
bounds".

Author: banach-space

mlir/include/mlir/IR/AffineMap.h

@@ -146,6 +146,14 @@ class AffineMap {
   /// affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
   bool isMinorIdentity() const;
 
+  /// Returns the list of broadcast dimensions (i.e. dims indicated by value 0
+  /// in the result).
+  /// Ex:
+  ///  * (d0, d1, d2) -> (0, d1) gives [0]
+  ///  * (d0, d1, d2) -> (d2, d1) gives []
+  ///  * (d0, d1, d2, d4) -> (d0, 0, d1, 0) gives [1, 3]
+  SmallVector<unsigned> getBroadcastDims() const;
+
   /// Returns true if this affine map is a minor identity up to broadcasted
   /// dimensions which are indicated by value 0 in the result. If
   /// `broadcastedDims` is not null, it will be populated with the indices of

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

@@ -4134,6 +4134,7 @@ static LogicalResult foldTransferInBoundsAttribute(TransferOp op) {
   bool changed = false;
   SmallVector<bool, 4> newInBounds;
   newInBounds.reserve(op.getTransferRank());
+  SmallVector<unsigned> nonBcastDims;
   for (unsigned i = 0; i < op.getTransferRank(); ++i) {
     // 1. Already marked as in-bounds, nothing to see here.
     if (op.isDimInBounds(i)) {
@@ -4148,15 +4149,27 @@ static LogicalResult foldTransferInBoundsAttribute(TransferOp op) {
       // 2.a Non-broadcast dim
       inBounds = isInBounds(op, /*resultIdx=*/i,
                             /*indicesIdx=*/dimExpr.getPosition());
-    } else {
-      // 2.b Broadcast dim
-      inBounds = true;
+      // 2.b Broadcast dims are handled after processing non-bcast dims
+      // FIXME: constant expr != 0 are not broadcasts - should such
+      // constants be allowed at all?
+      nonBcastDims.push_back(i);
     }
 
     newInBounds.push_back(inBounds);
     // We commit the pattern if it is "more inbounds".
     changed |= inBounds;
   }
+
+  // Handle broadcast dims: if all non-broadcast dims are "in
+  // bounds", then all bcast dims should be "in bounds" as well.
+  bool allNonBcastDimsInBounds = llvm::all_of(
+      nonBcastDims, [&newInBounds](unsigned idx) { return newInBounds[idx]; });
+  if (allNonBcastDimsInBounds)
+    llvm::for_each(permutationMap.getBroadcastDims(), [&](unsigned idx) {
+      changed |= !newInBounds[idx];
+      newInBounds[idx] = true;
+    });
+
   if (!changed)
     return failure();
   // OpBuilder is only used as a helper to build an I64ArrayAttr.

mlir/lib/IR/AffineMap.cpp

@@ -158,6 +158,19 @@ bool AffineMap::isMinorIdentity() const {
              getMinorIdentityMap(getNumDims(), getNumResults(), getContext());
 }
 
+SmallVector<unsigned> AffineMap::getBroadcastDims() const {
+  SmallVector<unsigned> broadcastedDims;
+  for (const auto &[resIdx, expr] : llvm::enumerate(getResults())) {
+    if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {
+      if (constExpr.getValue() != 0)
+        continue;
+      broadcastedDims.push_back(resIdx);
+    }
+  }
+
+  return broadcastedDims;
+}
+
 /// Returns true if this affine map is a minor identity up to broadcasted
 /// dimensions which are indicated by value 0 in the result.
 bool AffineMap::isMinorIdentityWithBroadcasting(

mlir/test/Dialect/Vector/vector-transfer-unroll.mlir

@@ -207,7 +207,7 @@ func.func @transfer_read_unroll_permutation(%mem : memref<6x4xf32>) -> vector<4x
 func.func @transfer_read_unroll_broadcast(%mem : memref<6x4xf32>) -> vector<6x4xf32> {
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
-  %res = vector.transfer_read %mem[%c0, %c0], %cf0 permutation_map = #map0} : memref<6x4xf32>, vector<6x4xf32>
+  %res = vector.transfer_read %mem[%c0, %c0], %cf0 {permutation_map = #map0} : memref<6x4xf32>, vector<6x4xf32>
   return %res : vector<6x4xf32>
 }
 
@@ -234,7 +234,7 @@ func.func @transfer_read_unroll_broadcast(%mem : memref<6x4xf32>) -> vector<6x4x
 func.func @transfer_read_unroll_broadcast_permuation(%mem : memref<6x4xf32>) -> vector<4x6xf32> {
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
-  %res = vector.transfer_read %mem[%c0, %c0], %cf0 permutation_map = #map0} : memref<6x4xf32>, vector<4x6xf32>
+  %res = vector.transfer_read %mem[%c0, %c0], %cf0 {permutation_map = #map0} : memref<6x4xf32>, vector<4x6xf32>
   return %res : vector<4x6xf32>
 }