vectorization check fix

nod-ai · Aug 28, 2024 · a8918aa · a8918aa
1 parent e1e3485
commit a8918aa
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/vectorization.mlir b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/vectorization.mlir
@@ -8,12 +8,12 @@
 // is vectorized to a vector.contract operation.
 // CHECK-LABEL: func @matmul_m4_n4_k8_bf16_f32
 func.func @matmul_m4_n4_k8_bf16_f32(%arg0: tensor<4x8xbf16>, %arg1: tensor<8x4xbf16>, %arg2: tensor<4x4xf32>) -> tensor<4x4xf32> {
-  // CHECK-DAG: arith.extf{{.*}} vector<4x8xbf16> to vector<4x8xf32>
-  // CHECK-DAG: arith.extf{{.*}} vector<8x4xbf16> to vector<8x4xf32>
+  // CHECK-DAG: arith.extf{{.*}} vector<4x4x8xbf16> to vector<4x4x8xf32>
+  // CHECK-DAG: arith.extf{{.*}} vector<4x4x8xbf16> to vector<4x4x8xf32>
   // CHECK-DAG: vector.transfer_read{{.*}} tensor<4x8xbf16>, vector<4x8xbf16>
   // CHECK-DAG: vector.transfer_read{{.*}} tensor<8x4xbf16>, vector<8x4xbf16>
   // CHECK-DAG: vector.transfer_read{{.*}} tensor<4x4xf32>, vector<4x4xf32>
-  // CHECK: vector.contract{{.*}}vector<4x8xf32>, vector<8x4xf32> into vector<4x4xf32>
+  // CHECK: vector.contract{{.*}}vector<4x4x8xf32>, vector<4x4x8xf32> into vector<4x4xf32>
   // CHECK: vector.transfer_write{{.*}} vector<4x4xf32>, tensor<4x4xf32>
   %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<4x8xbf16>, tensor<8x4xbf16>) outs(%arg2 : tensor<4x4xf32>) {
   ^bb0(%in: bf16, %in_0: bf16, %out: f32):