[Mlir-commits] [mlir] [mlir][vector] Add support for scalable vectors to VectorLinearize (PR #86786)

[Cullen Rhodes]

================
@@ -97,3 +97,47 @@ func.func @test_tensor_no_linearize(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf3
 
     return %0, %arg0 : tensor<2x2xf32>, tensor<2x2xf32>
 }
+
+// -----
+
+// ALL-LABEL:   func.func @test_1_scalable_dim(
+// ALL-SAME:    %[[ARG_0:.*]]: vector<2x[4]xf32>) -> vector<2x[4]xf32> {
+func.func @test_1_scalable_dim(%arg0: vector<2x[4]xf32>) -> vector<2x[4]xf32> {
+  // DEFAULT:  %[[SC:.*]] = vector.shape_cast %[[ARG_0]] : vector<2x[4]xf32> to vector<[8]xf32>
+  // DEFAULT:  %[[CST:.*]] = arith.constant dense<3.000000e+00> : vector<[8]xf32>
+  // BW-128:  %[[CST:.*]] = arith.constant dense<3.000000e+00> : vector<2x[4]xf32>
+  // BW-0:  %[[CST:.*]] = arith.constant dense<3.000000e+00> : vector<2x[4]xf32>
+  %0 = arith.constant dense<[[3., 3., 3., 3.], [3., 3., 3., 3.]]> : vector<2x[4]xf32>
+
+  // DEFAULT: %[[SIN:.*]] = math.sin %[[SC]] : vector<[8]xf32>
+  // BW-128: %[[SIN:.*]] = math.sin %[[ARG_0]] : vector<2x[4]xf32>
+  // BW-0: %[[SIN:.*]] = math.sin %[[ARG_0]] : vector<2x[4]xf32>
+  %1 = math.sin %arg0 : vector<2x[4]xf32>
+
+  // DEFAULT: %[[ADDF:.*]] = arith.addf %[[SIN]], %[[CST]] : vector<[8]xf32>
+  // BW-128: %[[RES:.*]] = arith.addf %[[CST]], %[[SIN]] : vector<2x[4]xf32>
+  // BW-0: %[[RES:.*]] = arith.addf %[[CST]], %[[SIN]] : vector<2x[4]xf32>
+  %2 = arith.addf %0, %1 : vector<2x[4]xf32>
+
+  // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[ADDF]] : vector<[8]xf32> to vector<2x[4]xf32>
+  // ALL: return %[[RES]] : vector<2x[4]xf32>
+  return %2 : vector<2x[4]xf32>
+}
+
+// -----
+
+// ALL-LABEL:   func.func @test_2_scalable_dims(
+// ALL-SAME:     %[[VAL_0:.*]]: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {
+func.func @test_2_scalable_dims(%arg0: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {
+  // ALL: %[[CST:.*]] = arith.constant dense<2.000000e+00> : vector<[2]x[2]xf32>
+  %0 = arith.constant dense<[[2., 2.], [2., 2.]]> : vector<[2]x[2]xf32>
----------------
c-rhodes wrote:

add a test to check non-splat constants are linearized? 

https://github.com/llvm/llvm-project/pull/86786