[Mlir-commits] [mlir] 1a5fc1c - [mlir] [linalg] Fold broadcast/transpose into linalg.generic (#212415)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 29 18:21:34 PDT 2026
Author: Chuanqi Xu
Date: 2026-07-30T01:21:29Z
New Revision: 1a5fc1c6b6bfff561f767d5e2ffe8b0a7d81cb34
URL: https://github.com/llvm/llvm-project/commit/1a5fc1c6b6bfff561f767d5e2ffe8b0a7d81cb34
DIFF: https://github.com/llvm/llvm-project/commit/1a5fc1c6b6bfff561f767d5e2ffe8b0a7d81cb34.diff
LOG: [mlir] [linalg] Fold broadcast/transpose into linalg.generic (#212415)
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.
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/Passes.td
mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
mlir/test/Dialect/Linalg/elementwise/fold.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 5573cf8b9feb7..b94b0ab477c07 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
+ 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 0be128c3b5e87..94db259d662d4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FoldIntoElementwise.cpp
@@ -7,15 +7,17 @@
//===----------------------------------------------------------------------===//
//
// This file implements folding ops such as transpose and broadcast into the
-// affine maps of the elementwise op.
+// affine maps of elementwise consumers.
//
//===----------------------------------------------------------------------===//
#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 {
@@ -31,8 +33,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) {
@@ -49,11 +50,14 @@ struct ElementwiseOpFolder {
};
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 {
+ if (!isa<GenericOp, ElementwiseOp>(op.getOperation()) || !isElementwise(op))
+ return failure();
+
bool changed = false;
SmallVector<Value> newIns;
SmallVector<AffineMap> newMaps;
@@ -72,11 +76,25 @@ struct FoldIntoElementwisePattern : public OpRewritePattern<ElementwiseOp> {
}
if (!changed)
return failure();
- newMaps.push_back(op.getIndexingMapsArray().back());
- rewriter.replaceOpWithNewOp<ElementwiseOp>(
- op, newIns, op.getDpsInits()[0], op.getKindAttr(),
- rewriter.getAffineMapArrayAttr(newMaps));
+ // 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.
+ // 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();
+
+ 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();
}
};
diff --git a/mlir/test/Dialect/Linalg/elementwise/fold.mlir b/mlir/test/Dialect/Linalg/elementwise/fold.mlir
index 80fd90f3d4dbe..3a0bec37d9463 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 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>
+}
More information about the Mlir-commits
mailing list