[Mlir-commits] [mlir] [mlir][vector] Refactor vector-transfer-to-vector-load-store.mlir (NFC) (PR #105509)

[Andrzej Warzyński]

================
@@ -330,54 +330,62 @@ func.func @transfer_read_permutations(%arg0 : memref<?x?xf32>, %arg1 : memref<?x
          vector<7x14x8x16xf32>, vector<8xf32>
 }
 
-// CHECK-LABEL: func @transfer_write_permutations
-// CHECK-SAME:      %[[ARG0:.*]]: memref<?x?x?x?xf32>
-// CHECK-SAME:      %[[ARG1:.*]]: tensor<?x?x?x?xf32>
-// CHECK-SAME:      %[[ARG2:.*]]: vector<7x14x8x16xf32>
-// CHECK-SAME:      %[[ARG3:.*]]: vector<8x16xf32>
-// CHECK-SAME:      %[[M:.*]]: i1
-func.func @transfer_write_permutations(
-    %arg0 : memref<?x?x?x?xf32>, %arg1 : tensor<?x?x?x?xf32>,
-    %v1 : vector<7x14x8x16xf32>, %v2 : vector<8x16xf32>, %m: i1) -> tensor<?x?x?x?xf32> {
+// CHECK-LABEL: func @transfer_write_permutations_tensor
+// CHECK-SAME:    %[[SRC:.*]]: tensor<?x?x?x?xf32>
+// CHECK-SAME:    %[[VEC:.*]]: vector<7x14x8x16xf32>
+// CHECK-SAME:    %[[M:.*]]: i1
+func.func @transfer_write_permutations_tensor(
+    %src : tensor<?x?x?x?xf32>,
+    %vec : vector<7x14x8x16xf32>, %m: i1) -> tensor<?x?x?x?xf32> {
   // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
   %c0 = arith.constant 0 : index
 
   // CHECK: %[[MASK:.*]] = vector.splat %[[M]] : vector<16x14x7x8xi1>
   %mask0 = vector.splat %m : vector<16x14x7x8xi1>
-  %0 = vector.transfer_write %v1, %arg1[%c0, %c0, %c0, %c0], %mask0 {in_bounds = [true, false, false, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3, d0)>} : vector<7x14x8x16xf32>, tensor<?x?x?x?xf32>
+  %res = vector.transfer_write %vec, %src[%c0, %c0, %c0, %c0], %mask0 {in_bounds = [true, false, false, true], permutation_map = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3, d0)>} : vector<7x14x8x16xf32>, tensor<?x?x?x?xf32>
   // CHECK: %[[NEW_VEC0:.*]] = vector.transpose %{{.*}} [3, 1, 0, 2] : vector<7x14x8x16xf32> to vector<16x14x7x8xf32>
-  // CHECK: %[[NEW_RES0:.*]] = vector.transfer_write %[[NEW_VEC0]], %[[ARG1]][%c0, %c0, %c0, %c0], %[[MASK]] {in_bounds = [true, false, true, false]} : vector<16x14x7x8xf32>, tensor<?x?x?x?xf32>
+  // CHECK: %[[NEW_RES0:.*]] = vector.transfer_write %[[NEW_VEC0]], %[[SRC]][%c0, %c0, %c0, %c0], %[[MASK]] {in_bounds = [true, false, true, false]} : vector<16x14x7x8xf32>, tensor<?x?x?x?xf32>
 
-  vector.transfer_write %v2, %arg0[%c0, %c0, %c0, %c0] {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d2)>} : vector<8x16xf32>, memref<?x?x?x?xf32>
-  // CHECK: %[[NEW_VEC1:.*]] = vector.transpose %{{.*}} [1, 0] : vector<8x16xf32> to vector<16x8xf32>
-  // CHECK: vector.transfer_write %[[NEW_VEC1]], %[[ARG0]][%c0, %c0, %c0, %c0] : vector<16x8xf32>, memref<?x?x?x?xf32>
+  return %res : tensor<?x?x?x?xf32>
+}
+
+// CHECK-LABEL: func @transfer_write_permutations_memref
+// CHECK-SAME:    %[[MEM:.*]]: memref<?x?x?x?xf32>
+// CHECK-SAME:    %[[VEC:.*]]: vector<8x16xf32>
+func.func @transfer_write_permutations_memref(
+    %mem : memref<?x?x?x?xf32>, %vec : vector<8x16xf32>) {
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  %c0 = arith.constant 0 : index
+
+  vector.transfer_write %vec, %mem[%c0, %c0, %c0, %c0] {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d2)>} : vector<8x16xf32>, memref<?x?x?x?xf32>
+  // CHECK: %[[NEW_VEC0:.*]] = vector.transpose %{{.*}} [1, 0] : vector<8x16xf32> to vector<16x8xf32>
+  // CHECK: vector.transfer_write %[[NEW_VEC0]], %[[MEM]][%c0, %c0, %c0, %c0] : vector<16x8xf32>, memref<?x?x?x?xf32>
 
-  return %0 : tensor<?x?x?x?xf32>
+  return
 }
 
 // CHECK-LABEL: func @transfer_write_broadcast_unit_dim
-// CHECK-SAME:      %[[ARG0:.*]]: memref<?x?x?x?xf32>
-// CHECK-SAME:      %[[ARG1:.*]]: tensor<?x?x?x?xf32>
-// CHECK-SAME:      %[[ARG2:.*]]: vector<14x8x16xf32>
-// CHECK-SAME:      %[[ARG3:.*]]: vector<8x16xf32>
-// CHECK-SAME:      %[[M:.*]]: i1
+// CHECK-SAME:    %[[MEM0:.*]]: memref<?x?x?x?xf32>
+// CHECK-SAME:    %[[SRC0:.*]]: tensor<?x?x?x?xf32>
+// CHECK-SAME:    %[[VEC0:.*]]: vector<14x8x16xf32>
+// CHECK-SAME:    %[[VEC1:.*]]: vector<8x16xf32>
 func.func @transfer_write_broadcast_unit_dim(
----------------
banach-space wrote:

+1

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