Restrict vector.type_cast to only cast to aligned multi-dim vectors

[Origami404]

Summary

Apply vector.type_cast to multi-dim memref<a... x b x T> where T is a
scalar type and b is not a power of 2 can produce MLIR vectors that will
be translated to a non-aligned aggregated LLVM vectors in memory. This
can cause confusing behaviors when accessing the element of the vector.

This commit adds checking for such cases and forbid them.

Link: https://discourse.llvm.org/t/restrict-vector-type-cast-to-only-cast-to-alignable-multi-dim-vectors/67811

[ftynse]

I'd avoid using "alignable" in the documentation, this is an obscure concept that wouldn't be properly defined in MLIR. One could argue, e.g., that any bitsize divisible by 256 is alignable on a machine with 256-bit vectors. Just say explicitly that the innermost size must be a power of 2.

[Origami404]

Adjust doc and diag

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

 }
 
 def Vector_TypeCastOp :
   Vector_Op<"type_cast", [Pure, ViewLikeOpInterface]>,
     Arguments<(ins StaticShapeMemRefOf<[AnyType]>:$memref)>,
     Results<(outs AnyMemRef:$result)> {
   let summary = "type_cast op converts a scalar memref to a vector memref";
   let description = [{
-    Performs a conversion from a memref with scalar element to a memref with a
-    *single* vector element, copying the shape of the memref to the vector. This
-    is the minimal viable operation that is required to makeke
-    super-vectorization operational. It can be seen as a special case of the
-    `view` operation but scoped in the super-vectorization context.
+    Performs a conversion from a memref with alignable scalar element or vector 
+    element to a memref with a *single* vector element, copying the shape of the 
+    memref to the vector. This is the minimal viable operation that is required 
+    to makeke super-vectorization operational. It can be seen as a special case 
+    of the `view` operation but scoped in the super-vectorization context.
 
     Syntax:
 
     ```
     operation ::= `vector.type_cast` ssa-use : memref-type to memref-type
     ```
 
     Example:
 
     ```mlir
-    %A  = memref.alloc() : memref<5x4x3xf32>
-    %VA = vector.type_cast %A : memref<5x4x3xf32> to memref<vector<5x4x3xf32>>
+    %A  = memref.alloc() : memref<5x4x4xf32>
+    %VA = vector.type_cast %A : memref<5x4x4xf32> to memref<vector<5x4x4xf32>>
     ```
+
+    ```mlir
+    %B  = memref.alloc() : memref<5x4xvector<3xf32>>
+    %VB = vector.type_cast %B : memref<5x4xvector<3xf32>> to memref<vector<5x4x3xf32>> 
+    ```
+
+    Because current restriction in the lowering of multi-dim vectors, the inner-
+    most dim of memref with scalar element type can only be a power of 2 to meet
+    the alignment requirement for casted vector.
   }];
 
   /// Build the canonical memRefType with a single vector.
   /// E.g. memref<4 x 5 x vector<6 x f32>> -> memref<vector<4 x 5 x 6 x f32>>.
   let builders = [OpBuilder<(ins "Value":$source)>];
 
   let extraClassDeclaration = [{
     MemRefType getMemRefType() {

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

if (getElementTypeOrSelf(getElementTypeOrSelf(sourceType)) !=
       getElementTypeOrSelf(getElementTypeOrSelf(resultType)))
     return emitOpError(
                "expects result and operand with same underlying scalar type: ")
            << resultType;
   if (extractShape(sourceType) != extractShape(resultType))
     return emitOpError(
                "expects concatenated result and operand shapes to be equal: ")
            << resultType;
+
+  if (!getElementTypeOrSelf(sourceType).isa<VectorType>()) {
+    if (auto shape = extractShape(sourceType); !shape.empty()) {
+      auto innerMostDim = shape.back();
+      if (!llvm::isPowerOf2_64(innerMostDim)) {
+        return emitOpError("can only cast to aligned vectors, inner-most dim "
+                           "for scalar type should be power of 2");
+      }
+    }
+  }
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TransposeOp
 //===----------------------------------------------------------------------===//
 
 void vector::TransposeOp::build(OpBuilder &builder, OperationState &result,

mlir/test/Dialect/Vector/invalid.mlir

 
 func.func @type_cast_layout(%arg0: memref<4x3xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + s2)>>) {
   // expected-error@+1 {{expects operand to be a memref with identity layout}}
   %0 = vector.type_cast %arg0: memref<4x3xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + s2)>> to memref<vector<4x3xf32>>
 }
 
 // -----
 
+func.func @type_cast_align(%arg0: memref<5x4x3xf32>) {
+  // expected-error@+1 {{can only cast to aligned vectors}}
+  %0 = vector.type_cast %arg0 : memref<5x4x3xf32> to memref<vector<5x4x3xf32>>
+}
+
+// -----
+
 func.func @store_unsupported_layout(%memref : memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>>,
                                %i : index, %j : index, %value : vector<8xf32>) {
   // expected-error@+1 {{'vector.store' op most minor memref dim must have unit stride}}
   vector.store %value, %memref[%i, %j] : memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>>,
                                          vector<8xf32>
   return
 }