[Mlir-commits] [mlir] [mlir][linalg] Add vectorization to the e2e test for tensor.unpack (PR #123032)

[Andrzej Warzyński]

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

>From af1768b649ceadd35439a1abbf38f700e3132f60 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Thu, 9 Jan 2025 18:43:57 +0000
Subject: [PATCH 1/2] [mlir][linalg] Add vectorization to the e2e test for
 `tensor.unpack`

Following on from #122927 + #123031 that added support for masked
vectorization of `tensor.insert_slice`, this PR extends the e2e test for
`tensor.unpack` to leverage the new functionality.
---
 .../Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir  |  4 ++--
 .../Linalg/CPU/unpack-dynamic-inner-tile.mlir        | 12 +++++++++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir
index 15edae8b6d3f85e..9a8708067c48f47 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir
@@ -121,11 +121,11 @@ module @transforms attributes { transform.with_named_sequence } {
       transform.apply_patterns.canonicalization
     } : !transform.op<"func.func">
 
-    // 3. Bufferize before lowering to LLVM
+    // 4. Bufferize before lowering to LLVM
     %bufferize = transform.bufferization.one_shot_bufferize %module
       {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
-    // 4. Canonicalize
+    // 5. Canonicalize
     %func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">
     transform.apply_patterns to %func_op_bufferized {
       transform.apply_patterns.canonicalization
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir
index 4395dfe74914e6b..90ecf4abfd1d9c5 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir
@@ -1,5 +1,6 @@
 // DEFINE: %{compile} =  mlir-opt %s \
-// DEFINE:  -transform-interpreter -test-transform-dialect-erase-schedule |\
+// DEFINE:  -transform-interpreter -test-transform-dialect-erase-schedule \
+// DEFINE:  --lower-vector-mask |\
 // DEFINE: mlir-opt \
 // DEFINE:  -test-lower-to-llvm -o %t
 // DEFINE: %{entry_point} = main
@@ -90,14 +91,19 @@ module @transforms attributes { transform.with_named_sequence } {
     %func_op = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">
     transform.apply_patterns to %func_op {
       transform.apply_patterns.linalg.decompose_pack_unpack
-      transform.apply_patterns.linalg.decompose_pad
+      transform.apply_patterns.canonicalization
     } : !transform.op<"func.func">
 
+    // 3. Vectorize tensor.insert_slice
+    // Vector sizes match the inner tiles in the payload IR.
+    %slice = transform.structured.match ops{["tensor.insert_slice"]} in %func_op : (!transform.op<"func.func">) -> !transform.any_op
+    transform.structured.vectorize %slice vector_sizes [8, 1] : !transform.any_op
+
     // 3. Bufferize before lowering to LLVM
     %bufferize = transform.bufferization.one_shot_bufferize %module
       {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
-   // 4. Canonicalize
+    // 4. Canonicalize
     %func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">
     transform.apply_patterns to %func_op_bufferized {
       transform.apply_patterns.canonicalization

>From eab99bb6bd912ef1c20db78ec7d23c24ab8f2d3f Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 5 Feb 2025 17:43:32 +0000
Subject: [PATCH 2/2] fixup! [mlir][linalg] Add vectorization to the e2e test
 for `tensor.unpack`

Fix numbering
---
 .../Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir         | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir
index 90ecf4abfd1d9c5..a97238e5eeea439 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/unpack-dynamic-inner-tile.mlir
@@ -99,11 +99,11 @@ module @transforms attributes { transform.with_named_sequence } {
     %slice = transform.structured.match ops{["tensor.insert_slice"]} in %func_op : (!transform.op<"func.func">) -> !transform.any_op
     transform.structured.vectorize %slice vector_sizes [8, 1] : !transform.any_op
 
-    // 3. Bufferize before lowering to LLVM
+    // 4. Bufferize before lowering to LLVM
     %bufferize = transform.bufferization.one_shot_bufferize %module
       {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
-    // 4. Canonicalize
+    // 5. Canonicalize
     %func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">
     transform.apply_patterns to %func_op_bufferized {
       transform.apply_patterns.canonicalization