Fix a corner case in vector.shape_cast when the trailing dimensions are of size 1.

[whchung]

Summary

Fix a corner case in vector.shape_cast when the trailing dimensions are of size 1.

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

1633

I find the * followed by / hard to read. Can we
(1) add comments on what we are doing here
(2) find better logic that tests before doing the * and canceling /?

[whchung]

Address code review comments.

[whchung]

mlir/lib/Dialect/Vector/VectorOps.cpp

1633

@aartbik Thanks. I revised the logic and added comments.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorOps.cpp

1636

Hmm looks like this could fail if both a and b have trailing ones since only j makes progress and the late i == rankA check will later fail?

How about explicitly checking for llvm::all_of(a... ) and llvm::all_of(b... ) are 1 ?

[whchung]

Address review comments by improve logic.

[whchung]

mlir/lib/Dialect/Vector/VectorOps.cpp

1636

@nicholas Thanks for catching this. Revised the logic.

[whchung]

Update commit message.

[whchung]

Remove an empty line.

mlir/lib/Dialect/Vector/VectorOps.cpp

static bool isValidShapeCast(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
 
   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;
 }
 
 static LogicalResult verifyVectorShapeCast(Operation *op,
                                            VectorType sourceVectorType,
                                            VectorType resultVectorType) {

mlir/test/Dialect/Vector/ops.mlir

func @reshape(%arg0 : vector<3x2x4xf32>) -> (vector<2x3x4xf32>) {
   %1 = vector.reshape %arg0, [%c3, %c6], [%c2, %c9], [4]
     : vector<3x2x4xf32> to vector<2x3x4xf32>
 
   return %1 : vector<2x3x4xf32>
 }
 
 // CHECK-LABEL: @shape_cast
 func @shape_cast(%arg0 : vector<5x1x3x2xf32>,
-                 %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>>)
-  -> (vector<15x2xf32>, tuple<vector<20x2xf32>, vector<12x2xf32>>) {
+                 %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>>,
+                 %arg2 : vector<8x1xf32>,
+                 %arg3 : vector<16x1x1xf32>)
+  -> (vector<15x2xf32>, tuple<vector<20x2xf32>, vector<12x2xf32>>, vector<8xf32>, vector<16xf32>, vector<16x1xf32>) {
 
   // CHECK: vector.shape_cast %{{.*}} : vector<5x1x3x2xf32> to vector<15x2xf32>
   %0 = vector.shape_cast %arg0 : vector<5x1x3x2xf32> to vector<15x2xf32>
 
   // CHECK-NEXT: vector.shape_cast %{{.*}} : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to tuple<vector<20x2xf32>, vector<12x2xf32>>
   %1 = vector.shape_cast %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to
                                  tuple<vector<20x2xf32>, vector<12x2xf32>>
 
-  return %0, %1 : vector<15x2xf32>, tuple<vector<20x2xf32>, vector<12x2xf32>>
+  // CHECK-NEXT: vector.shape_cast %{{.*}} : vector<8x1xf32> to vector<8xf32>
+  %2 = vector.shape_cast %arg2 : vector<8x1xf32> to vector<8xf32>
+
+  // CHECK-NEXT: vector.shape_cast %{{.*}} : vector<16x1x1xf32> to vector<16xf32>
+  %3 = vector.shape_cast %arg3 : vector<16x1x1xf32> to vector<16xf32>
+
+  // CHECK-NEXT: vector.shape_cast %{{.*}} : vector<16x1x1xf32> to vector<16x1xf32>
+  %4 = vector.shape_cast %arg3 : vector<16x1x1xf32> to vector<16x1xf32>
+
+  return %0, %1, %2, %3, %4 : vector<15x2xf32>, tuple<vector<20x2xf32>, vector<12x2xf32>>, vector<8xf32>, vector<16xf32>, vector<16x1xf32>
 }
 
 // CHECK-LABEL: @vector_fma
 func @vector_fma(%a: vector<8xf32>, %b: vector<8x4xf32>) {
   // CHECK: vector.fma %{{.*}} : vector<8xf32>
   vector.fma %a, %a, %a : vector<8xf32>
   // CHECK: vector.fma %{{.*}} : vector<8x4xf32>
   vector.fma %b, %b, %b : vector<8x4xf32>