[Mlir-commits] [mlir] [mlir] [linalg] Fold broadcast/transpose into linalg.generic (PR #212415)
Chuanqi Xu
llvmlistbot at llvm.org
Tue Jul 28 18:14:57 PDT 2026
https://github.com/ChuanqiXu9 updated https://github.com/llvm/llvm-project/pull/212415
>From f193948726377cd22d80b5f591c4580100aba7db Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Tue, 28 Jul 2026 14:23:58 +0800
Subject: [PATCH 1/4] [mlir] [linalg] Fold broadcast/transpose into
linalg.generic
Currently we are able to fold broadcast/transpose into
linalg.elementwise. This patch extends the ability to fold
broadcast/transpose into linalg.generic.
For example,
```
%empty = tensor.empty() : tensor<8x16xf32>
%broadcasted = linalg.broadcast ins(%A : tensor<8xf32>) outs(%empty : tensor<8x16xf32>) dimensions = [1]
%result = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%broadcasted : tensor<8x16xf32>) outs(%B : tensor<8x16xf32>) {
^bb0(%in: f32, %out: f32):
%v = arith.addf %in, %in : f32
linalg.yield %v : f32
} -> tensor<8x16xf32>
```
we can fold the broadcast into:
```
%result = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0)>,
affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%A: tensor<8xf32>) outs(%B : tensor<8x16xf32>) {
^bb0(%in: f32, %out: f32):
%v = arith.addf %in, %in : f32
linalg.yield %v : f32
} -> tensor<8x16xf32>
```
For simplicity, we only consider all parallel linalg.generic right now.
AI assisted.
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 12 +-
.../Linalg/Transforms/FoldIntoElementwise.cpp | 83 +++++++--
.../test/Dialect/Linalg/elementwise/fold.mlir | 165 ++++++++++++++++++
3 files changed, 235 insertions(+), 25 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 3a43af9ca1855..42e68f82d0e1b 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -170,15 +170,15 @@ def LinalgInlineScalarOperandsPass : Pass<"linalg-inline-scalar-operands"> {
}
def LinalgFoldIntoElementwisePass : Pass<"linalg-fold-into-elementwise"> {
- let summary = "Fold transpose and broadcast ops into elementwise";
+ let summary = "Fold transpose and broadcast ops into elementwise consumers";
let dependentDialects = ["linalg::LinalgDialect"];
let description = [{
- Fold transpose or broadcast op that feeds a `linalg.elementwise` into the
- elementwise op. `linalg.transpose` and `linalg.broadcast` producers whose
- consumer indexing map is a projected permutation can be absorbed into the
- indexing map of the `linalg.elementwise` by composing the producer's map
- into the elementwise op's indexing map. Other operands remain untouched.
+ Fold a transpose or broadcast that feeds a `linalg.elementwise` or an
+ all-parallel `linalg.generic` into its consumer. `linalg.transpose` and
+ `linalg.broadcast` producers whose consumer indexing map is a projected
+ permutation can be absorbed into the consumer's indexing map by composing
+ the producer's map into it. Other operands remain untouched.
}];
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
index 0be128c3b5e87..eda818fdc5f77 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
@@ -7,7 +7,7 @@
//===----------------------------------------------------------------------===//
//
// This file implements folding ops such as transpose and broadcast into the
-// affine maps of the elementwise op.
+// affine maps of elementwise consumers.
//
//===----------------------------------------------------------------------===//
@@ -31,8 +31,7 @@ using namespace mlir::linalg;
namespace {
template <typename ProducerOpTy>
struct ElementwiseOpFolder {
- // Helper function to fold broadcast etc into elementwise op.
- // Producer in this context is `broadcast op` etc, consumer is elwise operand.
+ // Helper function to fold broadcast etc. into a consumer operand.
static bool fold(OpOperand *elwiseOperand, AffineMap elwiseMap,
SmallVector<Value> &newIns,
SmallVector<AffineMap> &newMaps) {
@@ -48,29 +47,34 @@ struct ElementwiseOpFolder {
}
};
+template <typename ConsumerOpTy, typename... ProducerOps>
+static bool foldInputOperands(ConsumerOpTy op, SmallVector<Value> &newIns,
+ SmallVector<AffineMap> &newMaps) {
+ bool changed = false;
+ for (OpOperand *operand : op.getDpsInputOperands()) {
+ AffineMap consumerMap = op.getMatchingIndexingMap(operand);
+ const bool folded = (ElementwiseOpFolder<ProducerOps>::fold(
+ operand, consumerMap, newIns, newMaps) ||
+ ...);
+ if (folded) {
+ changed = true;
+ } else {
+ newIns.push_back(operand->get());
+ newMaps.push_back(consumerMap);
+ }
+ }
+ return changed;
+}
+
template <typename... ProducerOps>
struct FoldIntoElementwisePattern : public OpRewritePattern<ElementwiseOp> {
using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
LogicalResult matchAndRewrite(ElementwiseOp op,
PatternRewriter &rewriter) const override {
- bool changed = false;
SmallVector<Value> newIns;
SmallVector<AffineMap> newMaps;
- for (OpOperand *operand : op.getDpsInputOperands()) {
- AffineMap consumerMap = op.getMatchingIndexingMap(operand);
- const bool folded = (ElementwiseOpFolder<ProducerOps>::fold(
- operand, consumerMap, newIns, newMaps) ||
- ...);
- if (folded) {
- changed = true;
- } else {
- // push in original operand and its map.
- newIns.push_back(operand->get());
- newMaps.push_back(consumerMap);
- }
- }
- if (!changed)
+ if (!foldInputOperands<ElementwiseOp, ProducerOps...>(op, newIns, newMaps))
return failure();
newMaps.push_back(op.getIndexingMapsArray().back());
@@ -81,6 +85,46 @@ struct FoldIntoElementwisePattern : public OpRewritePattern<ElementwiseOp> {
}
};
+template <typename... ProducerOps>
+struct FoldIntoGenericPattern : public OpRewritePattern<GenericOp> {
+ using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(GenericOp op,
+ PatternRewriter &rewriter) const override {
+ // Restrict this pattern to elementwise-like generic ops.
+ // It may be safe to do so reduction dimensions in some cases. But we try
+ // to focus on simple cases here.
+ if (!op.isAllParallelLoops())
+ return failure();
+
+ SmallVector<Value> newIns;
+ SmallVector<AffineMap> newMaps;
+ if (!foldInputOperands<GenericOp, ProducerOps...>(op, newIns, newMaps))
+ return failure();
+
+ // Keep all output operands and their maps unchanged. The body is cloned
+ // so that the block arguments continue to correspond to the new operand
+ // list.
+ SmallVector<AffineMap> allMaps = op.getIndexingMapsArray();
+ newMaps.append(allMaps.begin() + op.getNumDpsInputs(), allMaps.end());
+ // The maps of the rewritten op must still determine bounds for every loop
+ // dimension. Folding a broadcast can otherwise drop the only map result
+ // that covers a dimension.
+ // See `generic_broadcast_not_folded_non_invertible` in
+ // mlir/test/Dialect/Linalg/elementwise/fold.mlir for an example.
+ if (!inversePermutation(concatAffineMaps(newMaps, op.getContext())))
+ return failure();
+ auto newOp =
+ GenericOp::create(rewriter, op.getLoc(), op.getResultTypes(), newIns,
+ op.getDpsInits(), newMaps, op.getIteratorTypesArray(),
+ /*bodyBuild=*/nullptr, getPrunedAttributeList(op));
+ rewriter.cloneRegionBefore(op.getRegion(), newOp.getRegion(),
+ newOp.getRegion().begin());
+ rewriter.replaceOp(op, newOp->getResults());
+ return success();
+ }
+};
+
struct LinalgFoldIntoElementwisePass
: public impl::LinalgFoldIntoElementwisePassBase<
LinalgFoldIntoElementwisePass> {
@@ -100,6 +144,7 @@ struct LinalgFoldIntoElementwisePass
void mlir::linalg::populateLinalgFoldIntoElementwisePatterns(
RewritePatternSet &patterns) {
- patterns.add<FoldIntoElementwisePattern<TransposeOp, BroadcastOp>>(
+ patterns.add<FoldIntoElementwisePattern<TransposeOp, BroadcastOp>,
+ FoldIntoGenericPattern<TransposeOp, BroadcastOp>>(
patterns.getContext());
}
diff --git a/mlir/test/Dialect/Linalg/elementwise/fold.mlir b/mlir/test/Dialect/Linalg/elementwise/fold.mlir
index 80fd90f3d4dbe..90dd89b408a03 100644
--- a/mlir/test/Dialect/Linalg/elementwise/fold.mlir
+++ b/mlir/test/Dialect/Linalg/elementwise/fold.mlir
@@ -245,3 +245,168 @@ func.func @fold_failed_constant_map(%A: tensor<16xf32>, %B: tensor<16x32xf32>, %
ins(%A, %transposed_B : tensor<16xf32>, tensor<32x16xf32>) outs(%C : tensor<16xf32>) -> tensor<16xf32>
return %result : tensor<16xf32>
}
+
+// -----
+
+// CHECK-DAG: #[[GENERIC_IDENTITY:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[GENERIC_BROADCASTED:.+]] = affine_map<(d0, d1) -> (d0)>
+// CHECK: func.func @generic_broadcast
+// CHECK-NOT: linalg.broadcast
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[GENERIC_BROADCASTED]], #[[GENERIC_IDENTITY]]]
+// CHECK-SAME: ins(%{{.*}} : tensor<8xf32>) outs(%{{.*}} : tensor<8x16xf32>)
+// CHECK: linalg.yield
+//
+#identity_generic = affine_map<(d0, d1) -> (d0, d1)>
+
+func.func @generic_broadcast(%A: tensor<8xf32>, %B: tensor<8x16xf32>) -> tensor<8x16xf32> {
+ %empty = tensor.empty() : tensor<8x16xf32>
+ %broadcasted = linalg.broadcast ins(%A : tensor<8xf32>) outs(%empty : tensor<8x16xf32>) dimensions = [1]
+ %result = linalg.generic {
+ indexing_maps = [#identity_generic, #identity_generic],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%broadcasted : tensor<8x16xf32>) outs(%B : tensor<8x16xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.addf %in, %in : f32
+ linalg.yield %v : f32
+ } -> tensor<8x16xf32>
+ return %result : tensor<8x16xf32>
+}
+
+// -----
+
+// CHECK-DAG: #[[GENERIC_IDENTITY:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[GENERIC_TRANSPOSED:.+]] = affine_map<(d0, d1) -> (d1, d0)>
+// CHECK: func.func @generic_transpose
+// CHECK-NOT: linalg.transpose
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[GENERIC_TRANSPOSED]], #[[GENERIC_IDENTITY]]]
+// CHECK-SAME: ins(%{{.*}} : tensor<16x8xf32>) outs(%{{.*}} : tensor<8x16xf32>)
+//
+func.func @generic_transpose(%A: tensor<16x8xf32>, %B: tensor<8x16xf32>) -> tensor<8x16xf32> {
+ %empty = tensor.empty() : tensor<8x16xf32>
+ %transposed = linalg.transpose
+ ins(%A : tensor<16x8xf32>) outs(%empty : tensor<8x16xf32>) permutation = [1, 0]
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%transposed : tensor<8x16xf32>) outs(%B : tensor<8x16xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.addf %in, %in : f32
+ linalg.yield %v : f32
+ } -> tensor<8x16xf32>
+ return %result : tensor<8x16xf32>
+}
+
+// -----
+
+// CHECK-DAG: #[[GENERIC_IDENTITY:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[GENERIC_BROADCASTED:.+]] = affine_map<(d0, d1) -> (d0)>
+// CHECK-DAG: #[[GENERIC_TRANSPOSED:.+]] = affine_map<(d0, d1) -> (d1, d0)>
+// CHECK: func.func @generic_broadcast_and_transpose
+// CHECK-NOT: linalg.broadcast
+// CHECK-NOT: linalg.transpose
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[GENERIC_BROADCASTED]], #[[GENERIC_TRANSPOSED]], #[[GENERIC_IDENTITY]], #[[GENERIC_IDENTITY]]]
+// CHECK-SAME: ins(%{{.*}}, %{{.*}}, %{{.*}} : tensor<8xf32>, tensor<16x8xf32>, tensor<8x16xf32>) outs(%{{.*}} : tensor<8x16xf32>)
+//
+func.func @generic_broadcast_and_transpose(
+ %A: tensor<8xf32>, %B: tensor<16x8xf32>, %C: tensor<8x16xf32>,
+ %D: tensor<8x16xf32>) -> tensor<8x16xf32> {
+ %broadcast_empty = tensor.empty() : tensor<8x16xf32>
+ %broadcasted = linalg.broadcast
+ ins(%A : tensor<8xf32>) outs(%broadcast_empty : tensor<8x16xf32>) dimensions = [1]
+ %transpose_empty = tensor.empty() : tensor<8x16xf32>
+ %transposed = linalg.transpose
+ ins(%B : tensor<16x8xf32>) outs(%transpose_empty : tensor<8x16xf32>) permutation = [1, 0]
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%broadcasted, %transposed, %C : tensor<8x16xf32>, tensor<8x16xf32>, tensor<8x16xf32>) outs(%D : tensor<8x16xf32>) {
+ ^bb0(%broadcast: f32, %transpose: f32, %input: f32, %out: f32):
+ %sum = arith.addf %broadcast, %transpose : f32
+ %result = arith.addf %sum, %input : f32
+ linalg.yield %result : f32
+ } -> tensor<8x16xf32>
+ return %result : tensor<8x16xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @generic_broadcast_multiple_uses
+// CHECK: %[[BROADCAST:.*]] = linalg.broadcast
+// CHECK: %[[RESULT:.*]] = linalg.generic
+// CHECK-SAME: ins(%{{.*}} : tensor<8xf32>) outs(%{{.*}} : tensor<8x16xf32>)
+// CHECK: return %[[BROADCAST]], %[[RESULT]] : tensor<8x16xf32>, tensor<8x16xf32>
+//
+func.func @generic_broadcast_multiple_uses(%A: tensor<8xf32>, %B: tensor<8x16xf32>)
+ -> (tensor<8x16xf32>, tensor<8x16xf32>) {
+ %empty = tensor.empty() : tensor<8x16xf32>
+ %broadcasted = linalg.broadcast
+ ins(%A : tensor<8xf32>) outs(%empty : tensor<8x16xf32>) dimensions = [1]
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%broadcasted : tensor<8x16xf32>) outs(%B : tensor<8x16xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.addf %in, %in : f32
+ linalg.yield %v : f32
+ } -> tensor<8x16xf32>
+ return %broadcasted, %result : tensor<8x16xf32>, tensor<8x16xf32>
+}
+
+// -----
+
+// This pass currently folds only elementwise-like, all-parallel generic ops.
+// Keep a reduction generic unchanged, even though its input map is foldable.
+// CHECK-LABEL: func.func @generic_reduction_not_folded
+// CHECK: linalg.broadcast
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "reduction"]
+//
+#reduction_map = affine_map<(d0, d1) -> (d0)>
+
+func.func @generic_reduction_not_folded(%A: tensor<8xf32>, %B: tensor<1xf32>) -> tensor<1xf32> {
+ %empty = tensor.empty() : tensor<1x8xf32>
+ %broadcasted = linalg.broadcast ins(%A : tensor<8xf32>) outs(%empty : tensor<1x8xf32>) dimensions = [0]
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, #reduction_map],
+ iterator_types = ["parallel", "reduction"]
+ } ins(%broadcasted : tensor<1x8xf32>) outs(%B : tensor<1xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.addf %in, %out : f32
+ linalg.yield %v : f32
+ } -> tensor<1xf32>
+ return %result : tensor<1xf32>
+}
+
+// -----
+
+// Folding must preserve invertibility of all indexing maps. The broadcast
+// input is the only operand covering d1 before the rewrite.
+// CHECK-LABEL: func.func @generic_broadcast_not_folded_non_invertible
+// CHECK: linalg.broadcast
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%{{.*}} : tensor<8x16xf32>) outs(%{{.*}} : tensor<8xf32>)
+//
+func.func @generic_broadcast_not_folded_non_invertible(
+ %A: tensor<8xf32>, %B: tensor<8xf32>) -> tensor<8xf32> {
+ %empty = tensor.empty() : tensor<8x16xf32>
+ %broadcasted = linalg.broadcast
+ ins(%A : tensor<8xf32>) outs(%empty : tensor<8x16xf32>) dimensions = [1]
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "parallel"]
+ } ins(%broadcasted : tensor<8x16xf32>) outs(%B : tensor<8xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.addf %in, %in : f32
+ linalg.yield %v : f32
+ } -> tensor<8xf32>
+ return %result : tensor<8xf32>
+}
>From 918d9cbcd5ce1147cb9d3ba432232581fc2cc36b Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Tue, 28 Jul 2026 18:27:31 +0800
Subject: [PATCH 2/4] Address comments
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 8 +--
.../Linalg/Transforms/FoldIntoElementwise.cpp | 63 ++++++-------------
.../test/Dialect/Linalg/elementwise/fold.mlir | 4 +-
3 files changed, 26 insertions(+), 49 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 42e68f82d0e1b..9b74ab158200b 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -175,10 +175,10 @@ def LinalgFoldIntoElementwisePass : Pass<"linalg-fold-into-elementwise"> {
let description = [{
Fold a transpose or broadcast that feeds a `linalg.elementwise` or an
- all-parallel `linalg.generic` into its consumer. `linalg.transpose` and
- `linalg.broadcast` producers whose consumer indexing map is a projected
- permutation can be absorbed into the consumer's indexing map by composing
- the producer's map into it. Other operands remain untouched.
+ elementwise-like `linalg.generic` into its consumer. `linalg.transpose`
+ and `linalg.broadcast` producers whose consumer indexing map is a
+ projected permutation can be absorbed into the consumer's indexing map by
+ composing the producer's map into it. Other operands remain untouched.
}];
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
index eda818fdc5f77..5e33721f45839 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
@@ -14,8 +14,10 @@
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
namespace mlir {
@@ -47,8 +49,8 @@ struct ElementwiseOpFolder {
}
};
-template <typename ConsumerOpTy, typename... ProducerOps>
-static bool foldInputOperands(ConsumerOpTy op, SmallVector<Value> &newIns,
+template <typename... ProducerOps>
+static bool foldInputOperands(LinalgOp op, SmallVector<Value> &newIns,
SmallVector<AffineMap> &newMaps) {
bool changed = false;
for (OpOperand *operand : op.getDpsInputOperands()) {
@@ -67,46 +69,23 @@ static bool foldInputOperands(ConsumerOpTy op, SmallVector<Value> &newIns,
}
template <typename... ProducerOps>
-struct FoldIntoElementwisePattern : public OpRewritePattern<ElementwiseOp> {
- using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
+struct FoldIntoElementwisePattern : public OpInterfaceRewritePattern<LinalgOp> {
+ using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
- LogicalResult matchAndRewrite(ElementwiseOp op,
+ LogicalResult matchAndRewrite(LinalgOp op,
PatternRewriter &rewriter) const override {
- SmallVector<Value> newIns;
- SmallVector<AffineMap> newMaps;
- if (!foldInputOperands<ElementwiseOp, ProducerOps...>(op, newIns, newMaps))
- return failure();
- newMaps.push_back(op.getIndexingMapsArray().back());
-
- rewriter.replaceOpWithNewOp<ElementwiseOp>(
- op, newIns, op.getDpsInits()[0], op.getKindAttr(),
- rewriter.getAffineMapArrayAttr(newMaps));
- return success();
- }
-};
-
-template <typename... ProducerOps>
-struct FoldIntoGenericPattern : public OpRewritePattern<GenericOp> {
- using OpRewritePattern<GenericOp>::OpRewritePattern;
-
- LogicalResult matchAndRewrite(GenericOp op,
- PatternRewriter &rewriter) const override {
- // Restrict this pattern to elementwise-like generic ops.
- // It may be safe to do so reduction dimensions in some cases. But we try
- // to focus on simple cases here.
- if (!op.isAllParallelLoops())
+ if (!isa<GenericOp, ElementwiseOp>(op.getOperation()) || !isElementwise(op))
return failure();
SmallVector<Value> newIns;
SmallVector<AffineMap> newMaps;
- if (!foldInputOperands<GenericOp, ProducerOps...>(op, newIns, newMaps))
+ if (!foldInputOperands<ProducerOps...>(op, newIns, newMaps))
return failure();
- // Keep all output operands and their maps unchanged. The body is cloned
- // so that the block arguments continue to correspond to the new operand
- // list.
- SmallVector<AffineMap> allMaps = op.getIndexingMapsArray();
- newMaps.append(allMaps.begin() + op.getNumDpsInputs(), allMaps.end());
+ // Keep all output operands and their maps unchanged.
+ SmallVector<AffineMap> originalMaps = op.getIndexingMapsArray();
+ newMaps.append(originalMaps.begin() + op.getNumDpsInputs(),
+ originalMaps.end());
// The maps of the rewritten op must still determine bounds for every loop
// dimension. Folding a broadcast can otherwise drop the only map result
// that covers a dimension.
@@ -114,13 +93,12 @@ struct FoldIntoGenericPattern : public OpRewritePattern<GenericOp> {
// mlir/test/Dialect/Linalg/elementwise/fold.mlir for an example.
if (!inversePermutation(concatAffineMaps(newMaps, op.getContext())))
return failure();
- auto newOp =
- GenericOp::create(rewriter, op.getLoc(), op.getResultTypes(), newIns,
- op.getDpsInits(), newMaps, op.getIteratorTypesArray(),
- /*bodyBuild=*/nullptr, getPrunedAttributeList(op));
- rewriter.cloneRegionBefore(op.getRegion(), newOp.getRegion(),
- newOp.getRegion().begin());
- rewriter.replaceOp(op, newOp->getResults());
+
+ rewriter.modifyOpInPlace(op, [&] {
+ for (auto [index, operand] : llvm::enumerate(op.getDpsInputOperands()))
+ op->setOperand(operand->getOperandNumber(), newIns[index]);
+ op->setAttr("indexing_maps", rewriter.getAffineMapArrayAttr(newMaps));
+ });
return success();
}
};
@@ -144,7 +122,6 @@ struct LinalgFoldIntoElementwisePass
void mlir::linalg::populateLinalgFoldIntoElementwisePatterns(
RewritePatternSet &patterns) {
- patterns.add<FoldIntoElementwisePattern<TransposeOp, BroadcastOp>,
- FoldIntoGenericPattern<TransposeOp, BroadcastOp>>(
+ patterns.add<FoldIntoElementwisePattern<TransposeOp, BroadcastOp>>(
patterns.getContext());
}
diff --git a/mlir/test/Dialect/Linalg/elementwise/fold.mlir b/mlir/test/Dialect/Linalg/elementwise/fold.mlir
index 90dd89b408a03..3a0bec37d9463 100644
--- a/mlir/test/Dialect/Linalg/elementwise/fold.mlir
+++ b/mlir/test/Dialect/Linalg/elementwise/fold.mlir
@@ -362,8 +362,8 @@ func.func @generic_broadcast_multiple_uses(%A: tensor<8xf32>, %B: tensor<8x16xf3
// -----
-// This pass currently folds only elementwise-like, all-parallel generic ops.
-// Keep a reduction generic unchanged, even though its input map is foldable.
+// This pass currently folds only elementwise-like generic ops. Keep a
+// reduction generic unchanged, even though its input map is foldable.
// CHECK-LABEL: func.func @generic_reduction_not_folded
// CHECK: linalg.broadcast
// CHECK: linalg.generic
>From 2dc3f960a3d918a021696f9581e652a3f11a911e Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Wed, 29 Jul 2026 09:06:42 +0800
Subject: [PATCH 3/4] Address comments
---
.../Linalg/Transforms/FoldIntoElementwise.cpp | 35 ++++++++-----------
1 file changed, 15 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
index 5e33721f45839..25b557c70da9a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
@@ -49,25 +49,6 @@ struct ElementwiseOpFolder {
}
};
-template <typename... ProducerOps>
-static bool foldInputOperands(LinalgOp op, SmallVector<Value> &newIns,
- SmallVector<AffineMap> &newMaps) {
- bool changed = false;
- for (OpOperand *operand : op.getDpsInputOperands()) {
- AffineMap consumerMap = op.getMatchingIndexingMap(operand);
- const bool folded = (ElementwiseOpFolder<ProducerOps>::fold(
- operand, consumerMap, newIns, newMaps) ||
- ...);
- if (folded) {
- changed = true;
- } else {
- newIns.push_back(operand->get());
- newMaps.push_back(consumerMap);
- }
- }
- return changed;
-}
-
template <typename... ProducerOps>
struct FoldIntoElementwisePattern : public OpInterfaceRewritePattern<LinalgOp> {
using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
@@ -77,9 +58,23 @@ struct FoldIntoElementwisePattern : public OpInterfaceRewritePattern<LinalgOp> {
if (!isa<GenericOp, ElementwiseOp>(op.getOperation()) || !isElementwise(op))
return failure();
+ bool changed = false;
SmallVector<Value> newIns;
SmallVector<AffineMap> newMaps;
- if (!foldInputOperands<ProducerOps...>(op, newIns, newMaps))
+ for (OpOperand *operand : op.getDpsInputOperands()) {
+ AffineMap consumerMap = op.getMatchingIndexingMap(operand);
+ const bool folded = (ElementwiseOpFolder<ProducerOps>::fold(
+ operand, consumerMap, newIns, newMaps) ||
+ ...);
+ if (folded) {
+ changed = true;
+ } else {
+ newIns.push_back(operand->get());
+ newMaps.push_back(consumerMap);
+ }
+ }
+
+ if (!changed)
return failure();
// Keep all output operands and their maps unchanged.
>From 8026d69d5c832297f2616b9d07ce8b0a89547db3 Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Wed, 29 Jul 2026 09:14:30 +0800
Subject: [PATCH 4/4] small change
---
mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
index 25b557c70da9a..28e72b2e7058f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
@@ -69,11 +69,11 @@ struct FoldIntoElementwisePattern : public OpInterfaceRewritePattern<LinalgOp> {
if (folded) {
changed = true;
} else {
+ // push in original operand and its map.
newIns.push_back(operand->get());
newMaps.push_back(consumerMap);
}
}
-
if (!changed)
return failure();
More information about the Mlir-commits
mailing list