[mlir][vector] Update `tensor.extract` for scalable vectors

[awarzynski]

Summary

This patch updates ExtractOp::inferReturnTypes so that scalable
dimensions are correctly recognised. This effectively fixes
CastAwayElementwiseLeadingOneDim when using scalable vectors and
that's what the supplemented test verify.

Depends on D157993

[awarzynski]

Closing in favor of https://reviews.llvm.org/D158335

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

ExtractOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
   auto vectorType = llvm::cast<VectorType>(adaptor.getVector().getType());
   if (static_cast<int64_t>(adaptor.getPosition().size()) ==
       vectorType.getRank()) {
     inferredReturnTypes.push_back(vectorType.getElementType());
   } else {
     auto n =
         std::min<size_t>(adaptor.getPosition().size(), vectorType.getRank());
     inferredReturnTypes.push_back(VectorType::get(
-        vectorType.getShape().drop_front(n), vectorType.getElementType()));
+        vectorType.getShape().drop_front(n), vectorType.getElementType(),
+        vectorType.getScalableDims().drop_front(n)));
   }
   return success();
 }
 
 bool ExtractOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   // Allow extracting 1-element vectors instead of scalars.
   auto isCompatible = [](TypeRange l, TypeRange r) {
     auto vectorType = llvm::dyn_cast<VectorType>(l.front());

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

 // 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.func @cast_away_leading_one_dim(
+// CHECK-SAME:      %[[VAL_0:.*]]: vector<1x4x1xf32>,
+// CHECK-SAME:      %[[VAL_1:.*]]: vector<1x4x1xf32>) -> vector<1x4x1xf32> {
+// CHECK:           %[[VAL_2:.*]] = vector.extract %[[VAL_0]][0] : vector<1x4x1xf32>
+// CHECK:           %[[VAL_3:.*]] = vector.extract %[[VAL_1]][0] : vector<1x4x1xf32>
+// CHECK:           %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<4x1xf32>
+// CHECK:           %[[VAL_5:.*]] = vector.broadcast %[[VAL_4]] : vector<4x1xf32> to vector<1x4x1xf32>
+// CHECK:           return %[[VAL_5]] : vector<1x4x1xf32>
+func.func @cast_away_leading_one_dim(%arg0: vector<1x4x1xf32>, %arg1: vector<1x4x1xf32>) -> vector<1x4x1xf32> {
+  %1 = arith.mulf %arg0, %arg1 : vector<1x4x1xf32>
+  return %1: vector<1x4x1xf32>
+}
+
+// CHECK-LABEL:   func.func @scalable_cast_away_leading_one_dim(
+// CHECK-SAME:       %[[VAL_0:.*]]: vector<1x[4]x1xf32>,
+// CHECK-SAME:       %[[VAL_1:.*]]: vector<1x[4]x1xf32>) -> vector<1x[4]x1xf32> {
+// CHECK:           %[[VAL_2:.*]] = vector.extract %[[VAL_0]][0] : vector<1x[4]x1xf32>
+// CHECK:           %[[VAL_3:.*]] = vector.extract %[[VAL_1]][0] : vector<1x[4]x1xf32>
+// CHECK:           %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<[4]x1xf32>
+// CHECK:           %[[VAL_5:.*]] = vector.broadcast %[[VAL_4]] : vector<[4]x1xf32> to vector<1x[4]x1xf32>
+// CHECK:           return %[[VAL_5]] : vector<1x[4]x1xf32>
+func.func @scalable_cast_away_leading_one_dim(%arg0: vector<1x[4]x1xf32>, %arg1: vector<1x[4]x1xf32>) -> vector<1x[4]x1xf32> {
+  %1 = arith.mulf %arg0, %arg1 : vector<1x[4]x1xf32>
+  return %1: vector<1x[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>