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

[Andrzej Warzyński]

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

>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 1/2] [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
   }
 }

>From 18dc22fe7004c3f3fe84df49ddb11a29b85c1ca7 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 19 Nov 2024 13:14:22 +0000
Subject: [PATCH 2/2] fixup! [mlir][linalg][nfc] Update
 pack-dynamic-inner-tile.mlir

Add comments
---
 .../Linalg/CPU/pack-dynamic-inner-tile.mlir   | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 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 e214847d17c61c..32b7247e60d622 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
@@ -1,8 +1,7 @@
 // DEFINE: %{compile} =  mlir-opt %s \
-// 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:    -test-lower-to-llvm -o %t
+// DEFINE:  -transform-interpreter -test-transform-dialect-erase-schedule |\
+// DEFINE: mlir-opt \
+// 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
@@ -85,12 +84,19 @@ module @transforms attributes { transform.with_named_sequence } {
   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
 
-    // 1. Tile so that we can decompose tensor.pack into tensor.pad,
-    // linalg.transpose, etc (see step 2).
+    // 1. Tile so that we can decompose tensor.pack into tensor.pad and other
+    // Ops (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
+    // 2. Decompose the tiled Op into (trimmed for brevity):
+    //
+    //  %padded = tensor.pad %slice_of_A (..) :
+    //      tensor<?x?xi32> to tensor<8x1xi32>
+    //  %inserted_slice = tensor.insert_slice %padded into %slice_of_A_pack (...) :
+    //      tensor<8x1xi32> into tensor<1x1x?x1xi32>
+    //
+    // NOTE: no tile is transposed, hence no linalg.transpose
     %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