[Mlir-commits] [mlir] [mlir][linalg] Hoist pad constants before fusion (PR #236686)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

The pad fusion pattern can use a constant that is still inside the pad region. Move it out before creating `linalg.fill`, so replacing the pad does not leave a dangling use.

Fixes #<!-- -->236378.


---
Full diff: https://github.com/llvm/llvm-project/pull/236686.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp (+3-1) 
- (modified) mlir/test/Dialect/Linalg/pad_fusion.mlir (+28) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp b/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp
index 0ca89040f34538..12c24565cc8d48 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp
@@ -19,7 +19,6 @@
 using namespace mlir;
 
 namespace {
-
 /// A sequence of operations
 ///
 /// ```mlir
@@ -78,6 +77,9 @@ struct FusePadOp : OpRewritePattern<tensor::PadOp> {
     // Fill the tensor with the pad value.
     // TODO: There is an option to fill only the boundaries. For now just
     // filling the whole tensor.
+    // Move the padding value defined inside the PadOp block to outside.
+    if (padValue.getParentBlock() == &padOp.getRegion().front())
+      rewriter.moveOpBefore(padValue.getDefiningOp(), padOp);
     auto fillTensor = linalg::FillOp::create(rewriter, loc, padValue,
                                              emptyTensor.getResult());
 
diff --git a/mlir/test/Dialect/Linalg/pad_fusion.mlir b/mlir/test/Dialect/Linalg/pad_fusion.mlir
index 56951b365b299d..aa7668b3f09595 100644
--- a/mlir/test/Dialect/Linalg/pad_fusion.mlir
+++ b/mlir/test/Dialect/Linalg/pad_fusion.mlir
@@ -51,6 +51,34 @@ func.func @dynamic_pad_fusion(%arg0 : tensor<?x?xf32>, %arg1 : index, %arg2 : in
 
 // -----
 
+func.func @folded_padding_value(%input: tensor<1xf32>,
+                                %init: tensor<1xf32>) -> tensor<2xf32> {
+  %producer = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+      iterator_types = ["parallel"]}
+      ins(%input : tensor<1xf32>) outs(%init : tensor<1xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %neg = arith.negf %in : f32
+      linalg.yield %neg : f32
+  } -> tensor<1xf32>
+  %padded = tensor.pad %producer low[1] high[0] {
+    ^bb0(%i: index):
+      %c4 = arith.constant 4.0 : f32
+      %c8 = arith.addf %c4, %c4 : f32
+      tensor.yield %c8 : f32
+  } : tensor<1xf32> to tensor<2xf32>
+  return %padded : tensor<2xf32>
+}
+
+// CHECK-LABEL: func.func @folded_padding_value
+//       CHECK:   %[[PAD_VALUE:.*]] = arith.constant 8.000000e+00 : f32
+//       CHECK:   %{{.*}} = linalg.fill ins(%[[PAD_VALUE]] : f32)
+//       CHECK:   tensor.insert_slice
+//   CHECK-NOT:   tensor.pad
+//       CHECK:   return
+
+// -----
+
 func.func @mixed_pad_fusion(%arg0 : tensor<?x42xf32>, %arg1 : index, %arg2 : index,
     %arg3 : f32) -> tensor<49x?xf32> {
   %c0 = arith.constant 0 : index

``````````

</details>


https://github.com/llvm/llvm-project/pull/236686