From 42a82b1ac6c834b7ce9766db22d169bd81415d52 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me@m-sp.org>
Date: Wed, 21 Jun 2023 09:17:32 +0200
Subject: [PATCH] [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
---
 .../Dialect/Linalg/transform-op-vectorize.mlir     | 45 ++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir b/mlir/test/Dialect/Linalg/transform-op-vectorize.mlir
index ecc077ae..d30db93 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)>
 
-- 
2.7.4