[mlir][Vector] Add 0-d vector support to 'vector.shape_cast`

[dcaballe]

Summary

This patch adds support to shape cast a vector<1x1x1...1xElemenType> to
a vector<ElementType> and the other way around.

[awarzynski]

Nice, thank you! LGTM

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

4620

[nit] The comment above applies to

if (rankA == 0 && llvm::all_of(b, isOne))
  return true;

rather than the lambda (i.e. isOne). Could you move it before merging? Ta!

[dcaballe]

Thanks!

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

let assemblyFormat =
       "$base `[` $indices `]` `,` $mask `,` $valueToStore attr-dict `:` "
       "type($base) `,` type($mask) `,` type($valueToStore)";
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def Vector_ShapeCastOp :
   Vector_Op<"shape_cast", [Pure]>,
-    Arguments<(ins AnyVector:$source)>,
-    Results<(outs AnyVector:$result)> {
+    Arguments<(ins AnyVectorOfAnyRank:$source)>,
+    Results<(outs AnyVectorOfAnyRank:$result)> {
   let summary = "shape_cast casts between vector shapes";
   let description = [{
     The shape_cast operation casts between an n-D source vector shape and
     a k-D result vector shape (the element type remains the same).
 
     If reducing rank (n > k), result dimension sizes must be a product
     of contiguous source dimension sizes.
     If expanding rank (n < k), source dimensions must factor into a

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

 
 /// Returns true if each element of 'a' is equal to the product of a contiguous
 /// sequence of the elements of 'b'. Returns false otherwise.
 static bool isValidShapeCast(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
   unsigned rankA = a.size();
   unsigned rankB = b.size();
   assert(rankA < rankB);
 
+  auto isOne = [](int64_t v) { return v == 1; };
+
+  // Special-case for n-D to 0-d shape cast. 'b' must be all ones to be shape
+  // casted to a 0-d vector.
+  if (rankA == 0 && llvm::all_of(b, isOne))
+    return true;
+
   unsigned i = 0;
   unsigned j = 0;
   while (i < rankA && j < rankB) {
     int64_t dimA = a[i];
     int64_t dimB = 1;
     while (dimB < dimA && j < rankB)
       dimB *= b[j++];
     if (dimA != dimB)
       break;
     ++i;
 
     // Handle the case when trailing dimensions are of size 1.
     // Include them into the contiguous sequence.
-    auto isOne = [](int64_t v) { return v == 1; };
     if (i < rankA && llvm::all_of(a.slice(i), isOne))
       i = rankA;
     if (j < rankB && llvm::all_of(b.slice(j), isOne))
       j = rankB;
   }
 
   return i == rankA && j == rankB;
 }

mlir/test/Dialect/Vector/ops.mlir

func.func @shape_cast(%arg0 : vector<5x1x3x2xf32>,
   %2 = vector.shape_cast %arg2 : vector<16x1x1xf32> to vector<16xf32>
 
   // CHECK-NEXT: vector.shape_cast %{{.*}} : vector<16x1x1xf32> to vector<16x1xf32>
   %3 = vector.shape_cast %arg2 : vector<16x1x1xf32> to vector<16x1xf32>
 
   return %0, %1, %2, %3 : vector<15x2xf32>, vector<8xf32>, vector<16xf32>, vector<16x1xf32>
 }
 
+// CHECK-LABEL: @shape_cast_0d
+func.func @shape_cast_0d(%arg0 : vector<1x1x1x1xf32>) -> (vector<1x1x1x1xf32>) {
+
+  // CHECK: vector.shape_cast %{{.*}} : vector<1x1x1x1xf32> to vector<f32>
+  %0 = vector.shape_cast %arg0 : vector<1x1x1x1xf32> to vector<f32>
+
+  // CHECK: vector.shape_cast %{{.*}} : vector<f32> to vector<1x1x1x1xf32>
+  %1 = vector.shape_cast %0 : vector<f32> to vector<1x1x1x1xf32>
+
+  return %1 : vector<1x1x1x1xf32>
+}
+
 // CHECK-LABEL: @bitcast
 func.func @bitcast(%arg0 : vector<5x1x3x2xf32>,
                  %arg1 : vector<8x1xi32>,
                  %arg2 : vector<16x1x8xi8>,
                  %arg3 : vector<8x2x1xindex>,
                  %arg4 : vector<f32>)
   -> (vector<5x1x3x4xf16>, vector<5x1x3x8xi8>, vector<8x4xi8>, vector<8x1xf32>, vector<16x1x2xi32>, vector<16x1x4xi16>, vector<16x1x1xindex>, vector<8x2x2xf32>, vector<i32>) {