[Mlir-commits] [mlir] 42a82b1 - [mlir][linalg][NFC] Add test case for memref vectorization

[Matthias Springer]

Author: Matthias Springer
Date: 2023-06-21T09:19:17+02:00
New Revision: 42a82b1ac6c834b7ce9766db22d169bd81415d52

URL: https://github.com/llvm/llvm-project/commit/42a82b1ac6c834b7ce9766db22d169bd81415d52
DIFF: https://github.com/llvm/llvm-project/commit/42a82b1ac6c834b7ce9766db22d169bd81415d52.diff

LOG: [mlir][linalg][NFC] Add test case for memref vectorization

Add test cases for vectorizing linalg.matmul and linalg.copy on tensors.

Differential Revision: https://reviews.llvm.org/D153357

Added: 
    

Modified: 
    mlir/test/Dialect/Linalg/transform-op-vectorize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir b/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
index ecc077ae5e9fb..d30db938f2bd1 100644
--- a/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
@@ -25,6 +25,51 @@ transform.sequence failures(propagate) {
 
 // -----
 
+// CHECK-LABEL: @vectorize_matmul_memref
+// CHECK-SAME: %[[A:.*]]: memref<24x12xf32>
+// CHECK-SAME: %[[B:.*]]: memref<12x25xf32>
+// CHECK-SAME: %[[C:.*]]: memref<24x25xf32>
+func.func @vectorize_matmul_memref(%arg0: memref<24x12xf32>,
+                                   %arg1: memref<12x25xf32>,
+                                   %arg2: memref<24x25xf32>) {
+  // CHECK: %[[vA:.+]] = vector.transfer_read %[[A]]
+  // CHECK: %[[vB:.+]] = vector.transfer_read %[[B]]
+  // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
+  // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
+  // CHECK: vector.transfer_write %[[vR]], %[[C]]
+  linalg.matmul ins(%arg0, %arg1 : memref<24x12xf32>, memref<12x25xf32>) outs(%arg2 : memref<24x25xf32>)
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = get_closest_isolated_parent %0 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.vectorize %1 : (!transform.any_op) -> !transform.any_op
+}
+
+// -----
+
+// CHECK-LABEL: @vectorize_copy_memref
+// CHECK-SAME: %[[A:.*]]: memref<100x100xf32>,
+// CHECK-SAME: %[[B:.*]]: memref<100x100xf32>
+func.func @vectorize_copy_memref(%arg0: memref<100x100xf32>,
+                                 %arg1: memref<100x100xf32>) {
+  // CHECK: %[[vA:.+]] = vector.transfer_read %[[A]]
+  // CHECK: vector.transfer_write %[[vA]], %[[B]]
+  linalg.copy ins(%arg0 : memref<100x100xf32>) outs(%arg1 : memref<100x100xf32>)
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = get_closest_isolated_parent %0 : (!transform.any_op) -> !transform.any_op
+  %2 = transform.structured.vectorize %1 : (!transform.any_op) -> !transform.any_op
+}
+
+// -----
+
 #map0 = affine_map<()[s0] -> (-s0 + 12, 7)>
 #map1 = affine_map<()[s0] -> (-s0 + 7)>