[mlir][linalg] Use ub.poison in data layout propagation if a packed operand requires padding. #159467

hanhanW · 2025-09-17T22:21:38Z

In the past, it was hard to set padding values because we did not have ub.poison. It is not always correct if we set zeros as padding values. Now we can use ub.poison in this case. The revision adds the support for setting padding value using ub.poison when padding is required in the propagation. Otherwise, it creates an invalid pack op.

Additionally the revision adds a control option for allowing padding in the pattern which is false by default. To correctly do this, a new requirePaddingValueStrict method is added which assumes dynamic dims would mean padding is required.

The revision also removes trailing white space in the lit test file.

Co-authored-by : Nirvedh Meshram nirvedh@gmail.com

Note to downstream users of these patterns, you may need to change to

linalg::populateDataLayoutPropagationPatterns(patterns, control,
                                                /*PoisonPaddingOk=*/true);

The reason is previously we didnt do any shape checking and just made the pack operands without padding, but now if we cant statically verify that padding is not needed the pattern bails out so you have to explictly opt in to say you are ok with padding.

…perand requires padding. In the past, it was hard to set padding values because we did not have ub.poison. It is not always correct if we set zeros as padding values. Now we can use `ub.poison` in this case. The revision adds the support for setting padding value using `ub.poison` when padding is required in the propagation. Otherwise, it creats an invalid pack op. Signed-off-by: hanhanW <hanhan0912@gmail.com>

llvmbot · 2025-09-17T22:22:11Z

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Han-Chung Wang (hanhanW)

Changes

In the past, it was hard to set padding values because we did not have ub.poison. It is not always correct if we set zeros as padding values. Now we can use ub.poison in this case. The revision adds the support for setting padding value using ub.poison when padding is required in the propagation. Otherwise, it creates an invalid pack op.

The revision also removes trailing white space in the lit test file.

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

2 Files Affected:

(modified) mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp (+6-3)
(modified) mlir/test/Dialect/Linalg/data-layout-propagation.mlir (+31-4)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index 6c17c3c2d0cab..2d075d92017f2 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -14,6 +14,7 @@
 #include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
@@ -289,9 +290,11 @@ getOrCreatePackedViewOfOperand(OpBuilder &b, Location loc, PackInfo packInfo,
 
   auto empty = linalg::PackOp::createDestinationTensor(
       b, loc, opOperand->get(), innerTileSizes, innerDimsPos, outerDimsPerm);
-  auto packedOperand = linalg::PackOp::create(
-      b, loc, opOperand->get(), empty, innerDimsPos, innerTileSizes,
-      /*padding=*/std::nullopt, outerDimsPerm);
+  auto poison = ub::PoisonOp::create(
+      b, loc, getElementTypeOrSelf(opOperand->get().getType()));
+  auto packedOperand =
+      linalg::PackOp::create(b, loc, opOperand->get(), empty, innerDimsPos,
+                             innerTileSizes, poison, outerDimsPerm);
   return std::make_tuple(packedOperand, indexingMap);
 }
 
diff --git a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
index a5f8d63a3e912..7a16bc0a4faee 100644
--- a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
+++ b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
@@ -1450,6 +1450,33 @@ func.func @push_unpack_in_padded_domain_out_used(%arg0: tensor<8x8x4x8xf32>, %ar
 
 // -----
 
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @push_unpack_in_padded_domain_multiple_inputs(%arg0: tensor<1x4x16x16xf32>, %arg1: tensor<8x64xf32>, %arg2: tensor<8x64xf32>) -> tensor<8x64xf32> {
+  %0 = tensor.empty() : tensor<8x64xf32>
+  %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %0 : tensor<1x4x16x16xf32> -> tensor<8x64xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg1, %unpack : tensor<8x64xf32>, tensor<8x64xf32>) outs(%arg2 : tensor<8x64xf32>) {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %2 = arith.addf %in, %in_0 : f32
+    linalg.yield %2 : f32
+  } -> tensor<8x64xf32>
+  return %1 : tensor<8x64xf32>
+}
+// CHECK-LABEL: func.func @push_unpack_in_padded_domain_multiple_inputs
+// CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[ARG2:[a-zA-Z0-9]+]]
+// CHECK-DAG:     %[[POISON:.+]] = ub.poison : f32
+// CHECK:         %[[PACK:.+]] = linalg.pack %[[ARG1]] padding_value(%[[POISON]] : f32)
+// CHECK-SAME:       inner_dims_pos = [0, 1] inner_tiles = [16, 16]
+// CHECK:         %[[ELEM:.+]] = linalg.generic
+// CHECK:           ins(%[[PACK]], %[[ARG0]]
+// CHECK:         %[[UNPACK:.+]] = linalg.unpack %[[ELEM]]
+// CHECK-SAME:      inner_dims_pos = [0, 1] inner_tiles = [16, 16]
+// CHECK-SAME:      into %[[ARG2]]
+// CHECK:         return %[[UNPACK]]
+
+// -----
+
 module {
   func.func @push_extract_through_generic(%arg0: tensor<128x7x128xf32>, %arg1: tensor<?x5x3x128xf32>, %arg2: tensor<?x5x128xbf16>, %arg3: index) -> tensor<?x5x128xbf16> {
     %extracted_slice = tensor.extract_slice %arg0[0, 0, %arg3] [128, 7, %arg3] [1, 1, 1] : tensor<128x7x128xf32> to tensor<128x7x?xf32>
@@ -1473,7 +1500,7 @@ module {
 // CHECK:         } : tensor<?x5x3x128xf32> to tensor<?x5x3x128xf32>
 // CHECK:         %[[EMPTY:.+]] = tensor.empty() : tensor<128x5x128xbf16>
 // CHECK:         %[[GENERIC:.+]] = linalg.generic
-// CHECK-SAME:    ins(%[[ARG0]], %[[PADDED]]   
+// CHECK-SAME:    ins(%[[ARG0]], %[[PADDED]]
 // CHECK-SAME:    outs(%[[EMPTY]]
 // CHECK:         %[[EXTRACT:.+]] = tensor.extract_slice %3[%[[ARG3]], 0, 0] [%[[ARG3]], 5, 128] [1, 1, 1] : tensor<128x5x128xbf16> to tensor<?x5x128xbf16>
 // CHECK:         return %[[EXTRACT]]
@@ -1492,7 +1519,7 @@ func.func @nopush_extract_through_generic_nodimexpr1(%arg0: tensor<128x7x128xf32
 
 // CHECK-LABEL: func.func @nopush_extract_through_generic_nodimexpr1
 // CHECK:         %[[GENERIC:.+]] = linalg.generic
-// CHECK:         return %[[GENERIC]]          
+// CHECK:         return %[[GENERIC]]
 
 // -----
 
@@ -1508,7 +1535,7 @@ func.func @nopush_extract_through_generic_nodimexpr2(%arg0: tensor<128x?x128xf32
 
 // CHECK-LABEL: func.func @nopush_extract_through_generic_nodimexpr2
 // CHECK:         %[[GENERIC:.+]] = linalg.generic
-// CHECK:         return %[[GENERIC]]   
+// CHECK:         return %[[GENERIC]]
 
 // -----
 
@@ -1575,7 +1602,7 @@ func.func @push_extract_through_generic_rank0_operand(%arg0: tensor<128x128xf32>
 
 // CHECK-LABEL: func.func @push_extract_through_generic_rank0_operand
 // CHECK:         %[[GENERIC:.+]] = linalg.generic
-// CHECK:         %[[EXTRACT:.+]] = tensor.extract_slice %[[GENERIC]]         
+// CHECK:         %[[EXTRACT:.+]] = tensor.extract_slice %[[GENERIC]]
 // CHECK:         return %[[EXTRACT]]
 
 // -----

nirvedhmeshram

LGTM

hanhanW · 2025-09-17T22:34:59Z

Thanks for the quick review! I'll wait for a day to see if others have feedback.

Max191

This LGTM

nirvedhmeshram · 2025-09-18T22:11:02Z

Just to update here, @hanhanW and I discussed some control options to let the caller decide if it wants to propagate if it would lead to padding, I should be able to work on that tomorrow and will push those changes to this PR

banach-space · 2025-09-19T10:57:51Z

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

+  auto poison = ub::PoisonOp::create(
+      b, loc, getElementTypeOrSelf(opOperand->get().getType()));


Rather than explicitly creating the pad value by the user, why not take approach similar to #146088? (i.e. make the padding value "optional" and make the builder worry about it)

I think it is okay to have a separate builder to do what you shared, and it makes sense. However, we don't do it in this builder, because users can decide to not use padding values in dynamic shapes at risk. I think it is okay to create a new builder that has such behavior, but we should leave the current builder as what it is. See below explanation for some details. (One of the difference is that padding is optional to pack ops, but padding is required by vector.trasnfer_read op.)

It is hard to make the builder to check if padding value is required or not in dynamic shapes. To me, there are soft check and hard check for padding value requirement. E.g., the existing method is a soft check that returns false in dynamic cases. I'm proposing a hard check (in IREE's issue), and use it with the control that @nirvedhmeshram is working on.

(The hard check version returns true in dynamic cases.)

bool PackOp::requirePaddingValue(ArrayRef<int64_t> inputShape, ArrayRef<int64_t> innerDimsPos, ArrayRef<int64_t> outputShape, ArrayRef<int64_t> outerDimsPerm, ArrayRef<OpFoldResult> innerTiles) { SmallVector<int64_t> outputTileSizes( outputShape.take_front(inputShape.size())); if (!outerDimsPerm.empty()) { assert(outerDimsPerm.size() == outputTileSizes.size() && "expected output and outer_dims_perm to have same size"); applyPermutationToVector(outputTileSizes, invertPermutationVector(outerDimsPerm)); } for (auto [pos, tileSize] : llvm::zip_equal(innerDimsPos, innerTiles)) { if (ShapedType::isDynamic(inputShape[pos])) continue; std::optional<int64_t> constantTile = getConstantIntValue(tileSize); if (!constantTile) { if (ShapedType::isStatic(outputTileSizes[pos]) && (inputShape[pos] % outputTileSizes[pos] != 0)) return true; } else if (inputShape[pos] % (*constantTile) != 0) { return true; } } return false; }

For another example -- that I don't see in practice so far, which is just my hypothesis -- you may drop the padding value in this case:

The inner tile size is 4.

You have int range analysis that tells you the packing dimension size is a multiple of 4 (which is a dynamic shape).

adam-smnk

LGTM 👍
Great fix to prevent propagation from failing (as per the new test case).

+1 to adding controls (callback?) but can be done as a separate follow-up

nirvedhmeshram · 2025-09-19T16:25:22Z

LGTM 👍 Great fix to prevent propagation from failing (as per the new test case).

+1 to adding controls (callback?) but can be done as a separate follow-up

IREE has some downstream cases which break with propagating with padding so would like to add control option before landing, should be able to do it by today.

hanhanW · 2025-09-19T17:05:56Z

LGTM 👍 Great fix to prevent propagation from failing (as per the new test case).
+1 to adding controls (callback?) but can be done as a separate follow-up

IREE has some downstream cases which break with propagating with padding so would like to add control option before landing, should be able to do it by today.

@nirvedhmeshram you'll take over this PR; integrate it into your patch, right?

hanhanW · 2025-09-19T17:06:41Z

Just to update here, @hanhanW and I discussed some control options to let the caller decide if it wants to propagate if it would lead to padding, I should be able to work on that tomorrow and will push those changes to this PR

Oh, I missed this. You'll push changes to this PR. Thanks!

nirvedhmeshram · 2025-09-19T17:06:48Z

LGTM 👍 Great fix to prevent propagation from failing (as per the new test case).
+1 to adding controls (callback?) but can be done as a separate follow-up

IREE has some downstream cases which break with propagating with padding so would like to add control option before landing, should be able to do it by today.

@nirvedhmeshram you'll take over this PR; integrate it into your patch, right?

Yes!

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram · 2025-09-19T21:14:53Z

@hanhanW I cant add you as a reviewer since you are the original author who raised the PR but would like your review.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

hanhanW · 2025-09-19T21:18:38Z

Let's wait for an approval from @adam-smnk before we land the PR.

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

github-actions · 2025-09-19T22:33:11Z

✅ With the latest revision this PR passed the C/C++ code formatter.

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

hanhanW

Just some nits. @adam-smnk would you like to take another look? Thanks!

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

hanhanW

LGTM, I'll wait a day, just in case if @adam-smnk wants to take a look.

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram · 2025-09-24T21:26:50Z

Landing this now on basis of prior approval as I havent heard any further feedback on this.

…perand requires padding. (llvm#159467) In the past, it was hard to set padding values because we did not have ub.poison. It is not always correct if we set zeros as padding values. Now we can use `ub.poison` in this case. The revision adds the support for setting padding value using `ub.poison` when padding is required in the propagation. Otherwise, it creates an invalid pack op. Additionally the revision adds a control option for allowing padding in the pattern which is false by default. To correctly do this, a new `requirePaddingValueStrict` method is added which assumes dynamic dims would mean padding is required. The revision also removes trailing white space in the lit test file. Co-authored-by : Nirvedh Meshram <nirvedh@gmail.com> --------- Signed-off-by: hanhanW <hanhan0912@gmail.com> Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com> Co-authored-by: Nirvedh Meshram <nirvedh@gmail.com>

hanhanW requested review from banach-space, adam-smnk and Max191 September 17, 2025 22:21

hanhanW requested a review from nicolasvasilache as a code owner September 17, 2025 22:21

llvmbot added mlir:linalg mlir labels Sep 17, 2025

hanhanW mentioned this pull request Sep 17, 2025

[GPU] linalg.pack op invalid tile factor or output size provided iree-org/iree#22026

Closed

nirvedhmeshram approved these changes Sep 17, 2025

View reviewed changes

Max191 approved these changes Sep 18, 2025

View reviewed changes

banach-space reviewed Sep 19, 2025

View reviewed changes

adam-smnk approved these changes Sep 19, 2025

View reviewed changes

nirvedhmeshram requested review from dcaballe and rengolin as code owners September 19, 2025 20:58

add option to control poison padding

d9f526e

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram requested a review from Max191 September 19, 2025 21:02

nirvedhmeshram added a commit to nirvedhmeshram/iree that referenced this pull request Sep 19, 2025

[Do not Merge] Test llvm/llvm-project#159467

58f0893

nirvedhmeshram mentioned this pull request Sep 19, 2025

[Do not merge] Testing upstream PR iree-org/iree#22052

Closed

hanhanW commented Sep 19, 2025

View reviewed changes

nirvedhmeshram added a commit to nirvedhmeshram/iree that referenced this pull request Sep 19, 2025

[Do not Merge] Test llvm/llvm-project#159467

eb2f55f

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram force-pushed the propagation-with-poison branch from eb94e20 to 5a5c53e Compare September 19, 2025 22:35

improve padding check

5a5c53e

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram force-pushed the propagation-with-poison branch from 1d794b1 to 15c9016 Compare September 23, 2025 21:12

nirvedhmeshram requested a review from adam-smnk September 23, 2025 21:14

reviewer comments

15c9016

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram added a commit to nirvedhmeshram/iree that referenced this pull request Sep 23, 2025

[Do not Merge] Test llvm/llvm-project#159467

8b74f20

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram added a commit to nirvedhmeshram/iree that referenced this pull request Sep 23, 2025

[Do not Merge] Test llvm/llvm-project#159467

0ef3582

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

hanhanW commented Sep 23, 2025

View reviewed changes

further review comments

83c4777

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

hanhanW commented Sep 23, 2025

View reviewed changes

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp Outdated Show resolved Hide resolved

further review comments 2

ada610e

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>

nirvedhmeshram merged commit 6f58c16 into llvm:main Sep 24, 2025
9 checks passed

hanhanW deleted the propagation-with-poison branch September 26, 2025 00:08

		auto poison = ub::PoisonOp::create(
		b, loc, getElementTypeOrSelf(opOperand->get().getType()));

[mlir][linalg] Use ub.poison in data layout propagation if a packed operand requires padding. #159467

[mlir][linalg] Use ub.poison in data layout propagation if a packed operand requires padding. #159467

Uh oh!

Conversation

hanhanW commented Sep 17, 2025 • edited by nirvedhmeshram Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nirvedhmeshram left a comment

Choose a reason for hiding this comment

Uh oh!

hanhanW commented Sep 17, 2025

Uh oh!

Max191 left a comment

Choose a reason for hiding this comment

Uh oh!

nirvedhmeshram commented Sep 18, 2025

Uh oh!

banach-space Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

hanhanW Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

hanhanW Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

adam-smnk left a comment

Choose a reason for hiding this comment

Uh oh!

nirvedhmeshram commented Sep 19, 2025

Uh oh!

hanhanW commented Sep 19, 2025

Uh oh!

hanhanW commented Sep 19, 2025

Uh oh!

nirvedhmeshram commented Sep 19, 2025

Uh oh!

nirvedhmeshram commented Sep 19, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hanhanW commented Sep 19, 2025

Uh oh!

github-actions bot commented Sep 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nirvedhmeshram commented Sep 24, 2025

Uh oh!

Uh oh!

Uh oh!

hanhanW commented Sep 17, 2025 •

edited by nirvedhmeshram

Loading

llvmbot commented Sep 17, 2025 •

edited

Loading

github-actions bot commented Sep 19, 2025 •

edited

Loading