[Mlir-commits] [mlir] [mlir][linalg][nfc] Update pack-dynamic-inner-tile.mlir (PR #116788)

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/116788

Following on from #116373, updates "pack-dynamic-inner-tile.mlir" to use
TD Ops for all transformations except for lowering to LLVM.

This is an intermediate step before introducing vectorization.


>From f7b0b8a4f9cac3f480b7f6ac062d484f056dad30 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 19 Nov 2024 11:22:11 +0000
Subject: [PATCH] [mlir][linalg][nfc] Update pack-dynamic-inner-tile.mlir

Following on from #116373, updates "pack-dynamic-inner-tile.mlir" to use
TD Ops for all transformations except for lowering to LLVM.

This is an intermediate step before introducing vectorization.
---
 .../Linalg/CPU/pack-dynamic-inner-tile.mlir   | 26 +++++++++++++++----
 1 file changed, 21 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 0428ada86041da..e214847d17c61c 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
@@ -2,9 +2,7 @@
 // DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule |\
 // DEFINE:  mlir-opt --test-linalg-transform-patterns="test-decompose-tensor-pack"\
 // DEFINE:    --test-transform-dialect-erase-schedule \
-// DEFINE:    -one-shot-bufferize="bufferize-function-boundaries" \
-// DEFINE:    -buffer-deallocation-pipeline="private-function-dynamic-ownership" \
-// DEFINE:    -cse -canonicalize -test-lower-to-llvm -o %t
+// DEFINE:    -test-lower-to-llvm -o %t
 // DEFINE: %{entry_point} = main
 // DEFINE: %{run} = mlir-cpu-runner %t -e %{entry_point} -entry-point-result=void \
 // DEFINE:    -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
@@ -84,12 +82,30 @@ func.func private @pack(%A: tensor<7x16xi32>) {
 }
 
 module @transforms attributes { transform.with_named_sequence } {
-  transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+  transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consume}) {
     %pack = transform.structured.match ops{["tensor.pack"]} in %module : (!transform.any_op) -> !transform.any_op
 
-    %tiled_linalg_op_p, %loops:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]
+    // 1. Tile so that we can decompose tensor.pack into tensor.pad,
+    // linalg.transpose, etc (see step 2).
+    %tiled_pack_op_p, %loops:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]
        : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 
+    // 2. Decompose the tiled Op into tensor.pad etc
+    %func_1 = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %func_1 {
+      transform.apply_patterns.linalg.decompose_pack_unpack
+    } : !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
+    %func_2 = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">
+    transform.apply_patterns to %func_2 {
+      transform.apply_patterns.canonicalization
+    } : !transform.op<"func.func">
+
     transform.yield
   }
 }