[Mlir-commits] [mlir] [mlir][Linalg][Tensor] Preserve attrs on `tensor.pad` when lowering to dst-style (PR #182064)
Artem Gindinson
llvmlistbot at llvm.org
Wed Feb 18 09:04:08 PST 2026
https://github.com/AGindinson created https://github.com/llvm/llvm-project/pull/182064
When canonicalizing to generic ops within `EliminateEmptyTensors`, we should take care to preserve the attributes. For example, this attribute mechanism is employed within IREE's SPIRV pipeline to pass on tiling configurations together with the ops.
Additionally, unify handling of invariant and constant values when lowering to `linalg.fill`, so as to minimize discardable attr boilerplate.
>From b9b00c9e63ac1d73450304250b0b2d78d15a7183 Mon Sep 17 00:00:00 2001
From: Artem Gindinson <gindinson at roofline.ai>
Date: Fri, 14 Nov 2025 14:20:59 +0100
Subject: [PATCH] [mlir][Linalg][Tensor] Preserve attributes on `tensor.pad`
when lowering to dst-style
When canonicalizing to generic ops within `EliminateEmptyTensors`, we should take care to preserve the attributes.
For example, this attribute mechanism is employed within IREE's SPIRV pipeline to pass on tiling configurations together with the ops.
Additionally, unify handling of invariant and constant values when lowering to `linalg.fill`, so as to minimize discardable attr boilerplate.
Signed-off-by: Artem Gindinson <gindinson at roofline.ai>
---
.../Transforms/ConvertToDestinationStyle.cpp | 48 +++++++++++--------
...-rewrite-in-destination-passing-style.mlir | 12 ++---
2 files changed, 33 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
index 2ff7f4617512b..0c37bf8bb1297 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
@@ -20,6 +20,7 @@
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "llvm/ADT/STLExtras.h"
@@ -96,39 +97,44 @@ static Operation *movePaddingToFillOrGenericOp(RewriterBase &rewriter,
OpBuilder::InsertionGuard g(rewriter);
RankedTensorType resultType = padOp.getResultType();
- // Examine the yielded value to decide if a linalg.generic is neede or a
+ // Collect user/dialect attributes from the pad op to preserve on the newly
+ // created ops.
+ SmallVector<NamedAttribute> preservedAttrs =
+ getPrunedAttributeList(padOp, PadOp::getAttributeNames());
+
+ // Examine the yielded value to decide if a linalg.generic is needed or a
// linalg.fill is sufficient.
Value yieldedValue =
cast<tensor::YieldOp>(padOp.getBody()->getTerminator()).getValue();
Attribute constYieldedValue;
+ const bool isConstYieldedValue =
+ matchPattern(yieldedValue, m_Constant(&constYieldedValue));
// Is the yielded value a bbArg defined outside of the PadOp?
- bool outsideBbArg =
+ const bool isOutsideBbArg =
isa<BlockArgument>(yieldedValue) &&
cast<BlockArgument>(yieldedValue).getOwner()->getParentOp() !=
padOp.getOperation();
// Is the yielded value an OpResult defined outside of the PadOp?
- bool outsideOpResult =
+ const bool isOutsideOpResult =
isa<OpResult>(yieldedValue) &&
yieldedValue.getDefiningOp()->getParentOp() != padOp.getOperation();
- bool invariantYieldedValue = outsideBbArg || outsideOpResult;
- if (matchPattern(yieldedValue, m_Constant(&constYieldedValue))) {
- // Padding with a constant: Create linalg.fill.
- Dialect *arithDialect =
- rewriter.getContext()->getLoadedDialect<arith::ArithDialect>();
- Value fillValue =
- arithDialect
- ->materializeConstant(rewriter, constYieldedValue,
- yieldedValue.getType(), yieldedValue.getLoc())
- ->getResult(0);
- auto fillOp = linalg::FillOp::create(rewriter, loc, ValueRange(fillValue),
- ValueRange(dest));
- return fillOp;
- }
-
- if (invariantYieldedValue) {
- // Padding with an invariant value.
+ const bool isInvariantYieldedValue =
+ isConstYieldedValue || isOutsideBbArg || isOutsideOpResult;
+ if (isInvariantYieldedValue) {
+ // Padding with an invariant value: Create linalg.fill.
+ if (isConstYieldedValue) {
+ // Padding with a constant value: Materialize the arith constant.
+ Dialect *arithDialect =
+ rewriter.getContext()->getLoadedDialect<arith::ArithDialect>();
+ yieldedValue = arithDialect
+ ->materializeConstant(rewriter, constYieldedValue,
+ yieldedValue.getType(),
+ yieldedValue.getLoc())
+ ->getResult(0);
+ }
auto fillOp = linalg::FillOp::create(
rewriter, loc, ValueRange(yieldedValue), ValueRange(dest));
+ fillOp->setDiscardableAttrs(preservedAttrs);
return fillOp;
}
@@ -140,7 +146,7 @@ static Operation *movePaddingToFillOrGenericOp(RewriterBase &rewriter,
auto genericOp = linalg::GenericOp::create(
rewriter, loc, resultType, /*inputs=*/ValueRange(),
/*outputs=*/ValueRange{dest}, /*indexingMaps=*/
- indexingMaps, iteratorTypes);
+ indexingMaps, iteratorTypes, /*bodyBuild=*/nullptr, preservedAttrs);
Block *body = rewriter.createBlock(&genericOp->getRegion(0), {},
resultType.getElementType(), loc);
rewriter.setInsertionPointToStart(body);
diff --git a/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir b/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir
index 63c9f1f27517b..9b74b5f5cf8d5 100644
--- a/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir
@@ -124,7 +124,7 @@ module attributes {transform.with_named_sequence} {
// CHECK: %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex>
// CHECK: %[[generic:.*]] = linalg.generic
// CHECK-SAME: {indexing_maps = [#[[$map2]]], iterator_types = ["parallel", "parallel"]}
-// CHECK-SAME: outs(%[[empty]] : tensor<?x?xindex>) {
+// CHECK-SAME: outs(%[[empty]] : tensor<?x?xindex>) attrs = {test.discardable.attr = 1 : i64} {
// CHECK: %[[i0:.*]] = linalg.index 0
// CHECK: %[[i1:.*]] = linalg.index 1
// CHECK: %[[mul:.*]] = arith.muli %[[i0]], %[[i1]]
@@ -139,7 +139,7 @@ func.func @tensor_pad(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
^bb0(%arg0: index, %arg1: index):
%m = arith.muli %arg0, %arg1 : index
tensor.yield %m : index
- } : tensor<?x10xindex> to tensor<?x?xindex>
+ } { test.discardable.attr = 1 : i64 } : tensor<?x10xindex> to tensor<?x?xindex>
return %0 : tensor<?x?xindex>
}
@@ -165,7 +165,7 @@ module attributes {transform.with_named_sequence} {
// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[h1]]]
// CHECK-DAG: %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]]
// CHECK: %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex>
-// CHECK: %[[filled:.*]] = linalg.fill ins(%[[c50]] : index) outs(%[[empty]] : tensor<?x?xindex>)
+// CHECK: %[[filled:.*]] = linalg.fill {test.discardable.attr = 1 : i64} ins(%[[c50]] : index) outs(%[[empty]] : tensor<?x?xindex>)
// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]
// CHECK: %[[inserted:.*]] = tensor.insert_slice %[[t1]] into %[[filled]][5, %[[l2]]] [%[[dim0]], 10] [1, 1] : tensor<?x10xindex> into tensor<?x?xindex>
// CHECK: return %[[inserted]]
@@ -175,7 +175,7 @@ func.func @tensor_pad_constant(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
^bb0(%arg0: index, %arg1: index):
%c = arith.constant 50 : index
tensor.yield %c : index
- } : tensor<?x10xindex> to tensor<?x?xindex>
+ } { test.discardable.attr = 1 : i64 } : tensor<?x10xindex> to tensor<?x?xindex>
return %0 : tensor<?x?xindex>
}
@@ -200,7 +200,7 @@ module attributes {transform.with_named_sequence} {
// CHECK-DAG: %[[size0:.*]] = affine.apply #[[$map]]()[%[[dim0]], %[[h1]]]
// CHECK-DAG: %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]]
// CHECK: %[[empty:.*]] = tensor.empty(%[[size0]], %[[size1]]) : tensor<?x?xindex>
-// CHECK: %[[filled:.*]] = linalg.fill ins(%[[padding]] : index) outs(%[[empty]] : tensor<?x?xindex>)
+// CHECK: %[[filled:.*]] = linalg.fill {test.discardable.attr = 1 : i64} ins(%[[padding]] : index) outs(%[[empty]] : tensor<?x?xindex>)
// CHECK-DAG: %[[dim0:.*]] = tensor.dim %[[t1]], %[[c0]]
// CHECK: %[[inserted:.*]] = tensor.insert_slice %[[t1]] into %[[filled]][5, %[[l2]]] [%[[dim0]], 10] [1, 1] : tensor<?x10xindex> into tensor<?x?xindex>
// CHECK: return %[[inserted]]
@@ -209,7 +209,7 @@ func.func @tensor_pad_invariant(%t1: tensor<?x10xindex>, %l2: index, %h1: index,
%0 = tensor.pad %t1 low[5, %l2] high[%h1, %h2] {
^bb0(%arg0: index, %arg1: index):
tensor.yield %padding : index
- } : tensor<?x10xindex> to tensor<?x?xindex>
+ } { test.discardable.attr = 1 : i64 } : tensor<?x10xindex> to tensor<?x?xindex>
return %0 : tensor<?x?xindex>
}
More information about the Mlir-commits
mailing list