[Mlir-commits] [mlir] [mlir][vector] Add tests for xfer-flatten patterns (PR #102210)

[Andrzej Warzyński]

https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/102210

>From 371f81112767017adb2b9a3e87580a6f6f1dcacf Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 6 Aug 2024 21:10:49 +0100
Subject: [PATCH] [mlir][vector] Add tests for xfer-flatten patterns

Adds tests for scalable vectors in:
  * vector-transfer-flatten.mlir

This is rather straightfoward as the tested patterns (*) do not support
scalable vectors.

(*) `FlattenContiguousRowMajorTransferReadPattern` and
`FlattenContiguousRowMajorTransferWritePattern`
---
 .../Vector/vector-transfer-flatten.mlir       | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
index 9d16aa46a9f2ad..4edc0aa728de8a 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
@@ -6,6 +6,8 @@
 ///----------------------------------------------------------------------------------------
 /// vector.transfer_read
 /// [Pattern: FlattenContiguousRowMajorTransferReadPattern]
+///
+/// NOTE: Scalable vectors are not supported
 ///----------------------------------------------------------------------------------------
 
 func.func @transfer_read_dims_match_contiguous(
@@ -28,6 +30,22 @@ func.func @transfer_read_dims_match_contiguous(
 // CHECK-128B-LABEL: func @transfer_read_dims_match_contiguous
 //       CHECK-128B:   memref.collapse_shape
 
+func.func @transfer_read_dims_match_contiguous_scalable(
+    %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<5x4x3x[2]xi8> {
+
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0 : i8
+  %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :
+    memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<5x4x3x[2]xi8>
+  return %res : vector<5x4x3x[2]xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_dims_match_contiguous_scalable
+// CHECK-NOT: memref.collapse_shape
+
+// CHECK-128B-LABEL: func @transfer_read_dims_match_contiguous_scalable
+//   CHECK-128B-NOT:   memref.collapse_shape
+
 // -----
 
 func.func @transfer_read_dims_match_contiguous_empty_stride(
@@ -259,6 +277,8 @@ func.func @transfer_read_non_contiguous_src(
 ///----------------------------------------------------------------------------------------
 /// vector.transfer_write
 /// [Pattern: FlattenContiguousRowMajorTransferWritePattern]
+///
+/// NOTE: Scalable vectors are not supported
 ///----------------------------------------------------------------------------------------
 
 func.func @transfer_write_dims_match_contiguous(
@@ -281,6 +301,22 @@ func.func @transfer_write_dims_match_contiguous(
 // CHECK-128B-LABEL: func @transfer_write_dims_match_contiguous(
 //       CHECK-128B:   memref.collapse_shape
 
+func.func @transfer_write_dims_match_contiguous_scalable(
+    %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,
+    %vec : vector<5x4x3x[2]xi8>) {
+
+  %c0 = arith.constant 0 : index
+  vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] :
+    vector<5x4x3x[2]xi8>, memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>
+  return
+}
+
+// CHECK-LABEL: func @transfer_write_dims_match_contiguous_scalable(
+// CHECK-NOT:   memref.collapse_shape
+
+// CHECK-128B-LABEL: func @transfer_write_dims_match_contiguous_scalable
+//   CHECK-128B-NOT:   memref.collapse_shape
+
 // -----
 
 func.func @transfer_write_dims_match_contiguous_empty_stride(
@@ -504,7 +540,11 @@ func.func @transfer_write_non_contiguous_src(
 
 ///----------------------------------------------------------------------------------------
 /// [Pattern: DropUnitDimFromElementwiseOps]
+///
 /// TODO: Move to a dedicated file - there's no "flattening" in the following tests
+/// TODO: Potential duplication with tests from:
+///   * "vector-dropleadunitdim-transforms.mlir" 
+///   * "vector-transfer-drop-unit-dims-patterns.mlir"
 ///----------------------------------------------------------------------------------------
 
 func.func @fold_unit_dim_add_basic(%vec : vector<1x8xi32>) -> vector<1x8xi32> {