[mlir][vector] Add support for scalable vectors in `trimLeadingOneDims`

[awarzynski]

Summary

This patch updates one specific hook in "VectorDropLeadUnitDim.cpp" to
make sure that "scalable dims" are handled correctly. While this change
affects multiple patterns, I am only adding one regression tests that
captures one specific case that affects me right now.

I am also adding Vector dialect to the list of dependencies of
-test-vector-to-vector-lowering. Otherwise my test case won't work as
a standalone test.

[benmxwl-arm]

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

25–35

This could still incorrectly trim something like vector<[1]x1xi128> to vector<1xi128> as the drop_while() is not checking if the dimension is scalable.

[awarzynski]

Refine how scalable vectors are treated

[awarzynski]

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

25–35

Great catch, thanks! It's a bit tricky to test this here (many patterns require an update to`tensor.extract` first), but in https://reviews.llvm.org/D158335 I've added a fair tests for cases like this.

[benmxwl-arm]

Very minor nit, otherwise LGTM!

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

25–29

nit: These are constructed a little inconsistently. Standard assignment should work here :)

[awarzynski]

Fix inconsistency - thanks for pointing that out Ben!

[dcaballe]

Thanks!

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

 #define DEBUG_TYPE "vector-drop-unit-dim"
 
 using namespace mlir;
 using namespace mlir::vector;
 
 // Trims leading one dimensions from `oldType` and returns the result type.
 // Returns `vector<1xT>` if `oldType` only has one element.
 static VectorType trimLeadingOneDims(VectorType oldType) {
   ArrayRef<int64_t> oldShape = oldType.getShape();
-  ArrayRef<int64_t> newShape =
-      oldShape.drop_while([](int64_t dim) { return dim == 1; });
+  ArrayRef<int64_t> newShape(oldShape.begin(), oldShape.end());
+
+  ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
+  ArrayRef<bool> newScalableDims(oldScalableDims);
+
+  while (!newShape.empty() && newShape.front() == 1 && !newScalableDims.front()) {
+    newShape = newShape.drop_front(1);
+    newScalableDims = newScalableDims.drop_front(1);
+  }
+
   // Make sure we have at least 1 dimension per vector type requirements.
-  if (newShape.empty())
+  if (newShape.empty()) {
     newShape = oldShape.take_back();
-  return VectorType::get(newShape, oldType.getElementType());
+    newScalableDims = oldType.getScalableDims().take_back();
+  }
+  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
 }
 
 /// Return a smallVector of size `rank` containing all zeros.
 static SmallVector<int64_t> splatZero(int64_t rank) {
   return SmallVector<int64_t>(rank, 0);
 }
 namespace {
 

mlir/test/Dialect/Vector/vector-transforms.mlir

 // CHECK-NEXT: %[[VEC1:.*]] = vector.insert_strided_slice %[[A2]], %[[VEC0]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x2xf32>
 // CHECK-NEXT: return %[[VEC1:.*]] : vector<4x2xf32>
 
 func.func @add4x2(%0: vector<4x2xf32>) -> vector<4x2xf32> {
   %1 = arith.addf %0, %0: vector<4x2xf32>
   return %1: vector<4x2xf32>
 }
 
+// Regression test. Previously, this example would trigger
+// CastAwayElementwiseLeadingOneDim as:
+//    * `vector<2x[4]x1xf32>`, would be reformulated as
+//    * `vector<2x4x1xf32>`.
+// With the updated shape, the conversion pattern would incorrectly assume that
+// some leading dims have been dropped.
+// CHECK-LABEL:   func.func @no_change(
+// CHECK-SAME:      %[[VAL_0:.*]]: vector<2x[4]x1xf32>,
+// CHECK-SAME:      %[[VAL_1:.*]]: vector<2x[4]x1xf32>)
+// CHECK-NEXT:           %[[VAL_2:.*]] = arith.mulf %[[VAL_0]], %[[VAL_1]] : vector<2x[4]x1xf32>
+// CHECK-NEXT:           return %[[VAL_2]]
+func.func @no_change(%arg0: vector<2x[4]x1xf32>, %arg1: vector<2x[4]x1xf32>) -> vector<2x[4]x1xf32> {
+  %1 = arith.mulf %arg0, %arg1 : vector<2x[4]x1xf32>
+  return %1 : vector<2x[4]x1xf32>
+}
+
 // CHECK-LABEL: func @add4x4
 //      CHECK: %[[S1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
 // CHECK-NEXT: %[[S2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
 
 // CHECK-NEXT: %[[A1:.*]] = arith.addf %[[S1]], %[[S2]] : vector<2x2xf32>
 
 // CHECK-NEXT: %[[S3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>
 // CHECK-NEXT: %[[S4:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xf32> to vector<2x2xf32>

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

StringRef getArgument() const final {
     return "test-vector-to-vector-lowering";
   }
   StringRef getDescription() const final {
     return "Test lowering patterns between ops in the vector dialect";
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<affine::AffineDialect>();
+    registry.insert<vector::VectorDialect>();
   }
 
   Option<bool> unroll{*this, "unroll", llvm::cl::desc("Include unrolling"),
                       llvm::cl::init(false)};
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);