[mlir][linalg] Fix bug in elementwise vectorization

[ThomasRaoux]

Summary

Fix a bug causing to pick the wrong vector size to broadcast to when the source vectors have different ranks.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

private:
   static SmallVector<int64_t, 4> getLargestShape(ArrayRef<Value> values) {
     SmallVector<int64_t, 4> largestShape;
     int64_t maxSize = 1;
     for (Value value : values) {
       auto vecType = value.getType().dyn_cast<VectorType>();
       if (!vecType)
         continue;
       if (maxSize < vecType.getNumElements()) {
+        maxSize = vecType.getNumElements();
         largestShape.assign(vecType.getShape().begin(),
                             vecType.getShape().end());
       }
     }
     return largestShape;
   }
 
   // If the value's type doesn't have the given shape broadcast it.

mlir/test/Dialect/Linalg/vectorization.mlir

^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32,
     %6 = addf %arg4, %arg6 : f32
     %7 = cmpf "ogt", %arg3, %arg6 : f32
     %8 = constant 2.0 : f32
     %9 = divf %arg5, %i : f32
     %10 = exp2 %arg5 : f32
     %11 = mulf %arg5, %8 : f32
     %12 = rsqrt %arg5 : f32
     %13 = select %7, %arg5, %arg6 : f32
-    %14 = subf %arg5, %arg6 : f32
+    %14 = subf %arg5, %arg4 : f32
     %15 = tanh %arg5 : f32
     linalg.yield %6, %8, %c1_f32, %9, %10, %11, %12, %13, %14, %15 : f32, f32,
       f32, f32, f32, f32, f32, f32, f32, f32
   }
   return
 }
 
 // CHECK-LABEL: func @generic_vectorize

 //       CHECK:   %[[CMP:.*]] = cmpf "ogt", %[[V2]], %[[V1]] : vector<4x256xf32>
 //       CHECK:   %[[V3:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32> 
 //       CHECK:   %[[ARG3B:.*]] = vector.broadcast %[[ARG3]] : f32 to vector<4x256xf32>
 //       CHECK:   %[[DIV:.*]] = divf %[[V3]], %[[ARG3B]] : vector<4x256xf32>
 //       CHECK:   %[[EXP:.*]] = exp2 %[[V3]] : vector<4x256xf32>
 //       CHECK:   %[[MUL:.*]] = mulf %[[V3]], %[[CST0]] : vector<4x256xf32>
 //       CHECK:   %[[RSQRT:.*]] = rsqrt %[[V3]] : vector<4x256xf32>
 //       CHECK:   %[[SEL:.*]] = select %[[CMP]], %[[V3]], %[[V1]] : vector<4x256xi1>, vector<4x256xf32>
-//       CHECK:   %[[SUB:.*]] = subf %[[V3]], %[[V1]] : vector<4x256xf32>
+//       CHECK:   %[[V0B:.*]] = vector.broadcast %[[V0]] : vector<256xf32> to vector<4x256xf32>
+//       CHECK:   %[[SUB:.*]] = subf %[[V3]], %[[V0B]] : vector<4x256xf32>
 //       CHECK:   %[[TAN:.*]] = tanh %[[V3]] : vector<4x256xf32>
 //       CHECK:   vector.transfer_write %[[ADD]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[CST0]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[CST1]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[DIV]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[EXP]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[MUL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[RSQRT]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[SEL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[SUB]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
 //       CHECK:   vector.transfer_write %[[TAN]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>