[Mlir-commits] [mlir] [mlir][vector] Clean up VectorTransferOpInterface (PR #72353)

[Diego Caballero]

================
@@ -46,246 +48,249 @@ def VectorUnrollOpInterface : OpInterface<"VectorUnrollOpInterface"> {
 
 def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
   let description = [{
-    Encodes properties of a transfer read or write operation.
+    Encodes properties of a `vector.transfer_read` or `vector.transfer_write`
+    operation. Vector transfer ops have:
+
+    - A shaped value that the op reads from/writes to: a memref or a tensor.
+    - A vector, either as a result or as an operand.
+    - Indicies that describe where the transfer from/to the shaped value starts.
+    - An optional mask.
+    - An optional in_bounds array to indicate transfer dimensions that are
+      guaranteed to be in-bounds.
+    - A permutation map to indicate transposes and broadcasts.
+
+    The "vector rank" is the rank of the vector type. E.g.:
+    ```
+    // Transfer with shaped value rank 2 and vector (transfer) rank 1.
+    %0 = vector.transfer_read %arg0[%c3, %c3], %f0
+        {permutation_map = affine_map<(d0, d1) -> (d0)>}
+        : memref<?x?xf32>, vector<128xf32>
+    ```
+
+    The "vector transfer rank" is the number of dimensions that participate in
+    the transfer and matches the number of results in the permutation map. In
+    most cases, the vector rank matches the vector transfer rank; the only
+    exception is when a vector is flattened as part of the transfer (see
+    `getPermutationMap`).
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
-    StaticInterfaceMethod<
-      /*desc=*/"Return the `in_bounds` attribute name.",
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the `in_bounds` attribute name.
+      }],
       /*retTy=*/"::mlir::StringRef",
-      /*methodName=*/"getInBoundsAttrStrName",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/ [{ return "in_bounds"; }]
+      /*methodName=*/"getInBoundsAttrName",
+      /*args=*/(ins)
     >,
-    StaticInterfaceMethod<
-      /*desc=*/"Return the `permutation_map` attribute name.",
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the `permutation_map` attribute name.
+      }],
       /*retTy=*/"::mlir::StringRef",
-      /*methodName=*/"getPermutationMapAttrStrName",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/ [{ return "permutation_map"; }]
+      /*methodName=*/"getPermutationMapAttrName",
+      /*args=*/(ins)
     >,
     InterfaceMethod<
-      /*desc=*/[{ Return `true` if dimension `dim` is in-bounds. Return `false`
-                 otherwise. }],
-      /*retTy=*/"bool",
-      /*methodName=*/"isDimInBounds",
-      /*args=*/(ins "unsigned":$dim),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        return $_op.isBroadcastDim(dim)
-            || ($_op.getInBounds()
-                && cast<::mlir::BoolAttr>(cast<::mlir::ArrayAttr>(*$_op.getInBounds())[dim]).getValue());
-      }]
+      /*desc=*/[{
+        Return the optional in_bounds attribute that specifies for each vector
+        dimension whether it is in-bounds or not. (Broadcast dimensions are
+        always in-bounds).
+      }],
+      /*retTy=*/"::std::optional<::mlir::ArrayAttr>",
+      /*methodName=*/"getInBounds",
+      /*args=*/(ins)
     >,
     InterfaceMethod<
-      /*desc=*/"Return the memref or ranked tensor operand.",
+      /*desc=*/[{
+        Return the memref or ranked tensor operand that this operation operates
+        on. In case of a "read" operation, that's the source from which the
+        operation reads. In case of a "write" operation, that's the destination
+        into which the operation writes.
+        TODO: Change name of operand, which is not accurate for xfer_write.
+      }],
       /*retTy=*/"::mlir::Value",
-      /*methodName=*/"source",
-      /*args=*/(ins),
-      /*methodBody=*/"return $_op.getSource();"
-      /*defaultImplementation=*/
+      /*methodName=*/"getSource",
+      /*args=*/(ins)
     >,
     InterfaceMethod<
-      /*desc=*/"Return the vector operand or result.",
+      /*desc=*/[{
+        Return the vector that this operation operates on. In case of a "read",
+        that's the vector OpResult. In case of a "write", that's the vector
+        operand value that is written by the op.
+      }],
       /*retTy=*/"::mlir::Value",
-      /*methodName=*/"vector",
-      /*args=*/(ins),
-      /*methodBody=*/"return $_op.getVector();"
-      /*defaultImplementation=*/
+      /*methodName=*/"getVector",
+      /*args=*/(ins)
     >,
     InterfaceMethod<
-      /*desc=*/"Return the indices operands.",
-      /*retTy=*/"::mlir::ValueRange",
-      /*methodName=*/"indices",
-      /*args=*/(ins),
-      /*methodBody=*/"return $_op.getIndices();"
-      /*defaultImplementation=*/
+      /*desc=*/[{
+        Return the indices that specify the starting offsets into the source
+        operand. The indices are guaranteed to be in-bounds.
+      }],
+      /*retTy=*/"::mlir::OperandRange",
+      /*methodName=*/"getIndices",
+      /*args=*/(ins)
     >,
     InterfaceMethod<
-      /*desc=*/"Return the permutation map.",
+      /*desc=*/[{
+        Return the permutation map that describes the mapping of vector
+        dimensions to source dimensions, as well as broadcast dimensions.
+
+        The permutation result has one result per vector transfer dimension.
+        Each result is either a dim expression, indicating the corresponding
+        dimension in the source operand, or a constant "0" expression,
+        indicating a broadcast dimension.
+
+        Note: Nested dimensions of flattened vector do are not accounted for in
+        the permutation map. E.g.:
+        ```
+        // Vector type has rank 4, but permutation map has only 2 results. That
+        // is because there are only 2 transfer dimensions.
+        %0 = vector.transfer_read %arg1[%c3, %c3], %vf0
+            {permutation_map = affine_map<(d0, d1) -> (d0, d1)>}
+            : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
----------------
dcaballe wrote:

I think vector element types are needed in general because they enforce the alignment of the elements to the vector size boundary. This is similar to allocating a vector type in LLVM. Part of the complexity comes from how we apply the permutation maps to these memrefs so perhaps we could just be more restrictive there. We can limit permutation to simple cases that are well defined and see how things go...

https://github.com/llvm/llvm-project/pull/72353