[mlir] Improve TransferOp verifier: broadcasts are in_bounds

[springerm]

Summary

Broadcast dimensions of vector transfer ops are always in-bounds. This is consistent with the fact that the starting position of a transfer is always in-bounds.

Depends On D102548

[springerm]

rebase

[springerm]

no change

mlir/include/mlir/Dialect/Vector/VectorOps.td

let description = [{
     provided to specify a fallback value in the case of out-of-bounds accesses
     and/or masking.
 
     An optional SSA value `mask` of the same shape as the vector type may be
     specified to mask out elements. Such elements will be replaces with
     `padding`. Elements whose corresponding mask element is `0` are masked out.
 
     An optional boolean array attribute is provided to specify which dimensions
-    of the transfer are guaranteed to be within bounds. The absence of this
-    `in_bounds` attribute signifies that any dimension of the transfer may be
-    out-of-bounds. A `vector.transfer_read` can be lowered to a simple load if
-    all dimensions are specified to be within bounds and no `mask` was
-    specified.
+    of the transfer are guaranteed to be within bounds. The length of the array
+    must equal the rank of the vector type. Broadcast dimensions must always be
+    in-bounds. The absence of this optional `in_bounds` attribute signifies that
+    any dimension of the transfer (except for broadcasts) may be out-of-bounds.
+    A `vector.transfer_read` can be lowered to a simple load if all dimensions
+    are specified to be within bounds and no `mask` was specified.
+
+    Note that `in_bounds` is specified for result dimensions and not input
+    dimensions. The starting point of the transfer, i.e.,
+    `%A[%expr1, %expr2, %expr3, %expr4]` in the example below, is expected to
+    be in-bounds and as indices are increasing, accesses may run out-of-bounds.
 
     This operation is called 'read' by opposition to 'load' because the
     super-vector granularity is generally not representable with a single
     hardware register. A `vector.transfer_read` is thus a mid-level abstraction
     that supports super-vectorization with non-effecting padding for full-tile
     only operations.
 
     More precisely, let's dive deeper into the permutation_map for the following
let description = [{
 
     The slice is further defined by a full-rank index within the MemRef/Tensor,
     supplied as the operands `3 .. 2 + rank(memref/tensor)`.
 
     The permutation_map [attribute](../LangRef.md#attributes) is an
     [affine-map](Affine.md#affine-maps) which specifies the transposition on the
     slice to match the vector shape. The permutation map may be implicit and
     omitted from parsing and printing if it is the canonical minor identity map
-    (i.e. if it does not permute or broadcast any dimension).
+    (i.e. if it does not permute any dimension). In contrast to `transfer_read`,
+    write ops cannot have broadcast dimensions.
 
     The size of the slice is specified by the size of the vector.
 
     An optional SSA value `mask` of the same shape as the vector type may be
     specified to mask out elements. Elements whose corresponding mask element
     is `0` are masked out.
 
     An optional boolean array attribute is provided to specify which dimensions
     of the transfer are guaranteed to be within bounds. The absence of this
     `in_bounds` attribute signifies that any dimension of the transfer may be
     out-of-bounds. A `vector.transfer_write` can be lowered to a simple store
     if all dimensions are specified to be within bounds and no `mask` was
     specified.
 
+    An optional boolean array attribute is provided to specify which dimensions
+    of the transfer are guaranteed to be within bounds. The length of the array
+    must equal the rank of the vector type. The absence of this optional
+    `in_bounds` attribute signifies that any dimension of the transfer
+    may be out-of-bounds. A `vector.transfer_write` can be lowered to a simple
+    store if all dimensions are specified to be within bounds and no `mask` was
+    specified.
+
+    Note that `in_bounds` is specified for result dimensions and not input
+    dimensions. The starting point of the transfer, i.e.,
+    `%A[%expr1, %expr2, %expr3, %expr4]` in the example below, is expected to
+    be in-bounds and as indices are increasing, accesses may run out-of-bounds.
+
     This operation is called 'write' by opposition to 'store' because the
     super-vector granularity is generally not representable with a single
     hardware register. A `vector.transfer_write` is thus a
     mid-level abstraction that supports super-vectorization with non-effecting
     padding for full-tile-only code. It is the responsibility of
     `vector.transfer_write`'s implementation to ensure the memory writes are
     valid. Different lowerings may be pertinent depending on the hardware
     support.

mlir/include/mlir/Interfaces/VectorInterfaces.td

StaticInterfaceMethod<
       /*desc=*/"Return the `permutation_map` attribute name.",
       /*retTy=*/"StringRef",
       /*methodName=*/"getPermutationMapAttrName",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/ [{ return "permutation_map"; }]
     >,
     InterfaceMethod<
-      /*desc=*/[{
-      Return `true` when the `in_bounds` attribute at dimension
-      `dim` is set to `true`. Return `false` otherwise.}],
+      /*desc=*/[{ Return `true` if dimension `dim` is in-bounds. Return `false`
+                 otherwise. }],
       /*retTy=*/"bool",
       /*methodName=*/"isDimInBounds",
       /*args=*/(ins "unsigned":$dim),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return $_op.in_bounds() &&
-          $_op.in_bounds()->template cast<ArrayAttr>()[dim]
-                          .template cast<BoolAttr>().getValue();
+        return $_op.isBroadcastDim(dim)
+            || ($_op.in_bounds()
+                && $_op.in_bounds()->template cast<ArrayAttr>()[dim]
+                                    .template cast<BoolAttr>().getValue());
       }]
     >,
     InterfaceMethod<
       /*desc=*/"Return the memref or ranked tensor operand.",
       /*retTy=*/"Value",
       /*methodName=*/"source",
       /*args=*/(ins),
       /*methodBody=*/"return $_op.source();"

mlir/lib/Dialect/Vector/VectorOps.cpp

if (permutationMap.getNumInputs() != shapedType.getRank())
     return op->emitOpError("requires a permutation_map with input dims of the "
                            "same rank as the source type");
 
   if (inBounds) {
     if (permutationMap.getNumResults() != static_cast<int64_t>(inBounds.size()))
       return op->emitOpError("expects the optional in_bounds attr of same rank "
                              "as permutation_map results: ")
              << AffineMapAttr::get(permutationMap);
+    for (unsigned int i = 0; i < permutationMap.getNumResults(); ++i)
+      if (permutationMap.getResult(i).isa<AffineConstantExpr>()
+          && !inBounds.getValue()[i].cast<BoolAttr>().getValue())
+        return op->emitOpError("requires broadcast dimensions to be in-bounds");
   }
 
   return success();
 }
 
 /// Builder that sets padding to zero.
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,

mlir/test/Dialect/Vector/invalid.mlir

 }
 
 // -----
 
 func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) {
   %c3 = constant 3 : index
   %f0 = constant 0.0 : f32
   %vf0 = splat %f0 : vector<2x3xf32>
+  // expected-error@+1 {{requires broadcast dimensions to be in-bounds}}
+  %0 = vector.transfer_read %arg0[%c3, %c3], %vf0 {in_bounds = [false, true], permutation_map = affine_map<(d0, d1)->(0, d1)>} : memref<?x?xvector<2x3xf32>>, vector<1x1x2x3xf32>
+}
+
+// -----
+
+func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) {
+  %c3 = constant 3 : index
+  %f0 = constant 0.0 : f32
+  %vf0 = splat %f0 : vector<2x3xf32>
   %mask = splat %c1 : vector<2x3xi1>
   // expected-error@+1 {{does not support masks with vector element type}}
   %0 = vector.transfer_read %arg0[%c3, %c3], %vf0, %mask {permutation_map = affine_map<(d0, d1)->(d0, d1)>} : memref<?x?xvector<2x3xf32>>, vector<1x1x2x3xf32>
 }
 
 // -----
 
 func @test_vector.transfer_write(%arg0: memref<?x?xf32>) {