[Mlir-commits] [mlir] [mlir][linalg] Split elementwise ops with concat inputs (PR #213630)
Chuanqi Xu
llvmlistbot at llvm.org
Thu Aug 20 20:29:53 PDT 2026
https://github.com/ChuanqiXu9 updated https://github.com/llvm/llvm-project/pull/213630
>From 9e40eaca96c223d2d039117d5c9386a381e8de00 Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Mon, 3 Aug 2026 16:49:38 +0800
Subject: [PATCH 1/5] [mlir][linalg] Split elementwise ops with concat inputs
Close https://github.com/llvm/llvm-project/issues/213216
This patch tries to reorder patterns like elementwise(concat(x0, x1),
concat(y0, y1)) into concat(elementwise(x0, y0), elementwise(x1, y1)).
The transformation itself is not an optimization. But it will make
the elementwise op to be closer with the data so that the optimizer
may find more optimization oppotunities. See the above patch for an
example.
For simplicity, this patch only handle cases for all concats have the
same number of inputs and the size of inputs in the concat dimension
to be the same. We also don't handle linalg.index. And if the inputs of
the elementwise op has other inputs than concat, which is not a scalar,
may be rejected too. We can relax these limitations in the future.
To make the implementation more uniform, we only handles elementwise
like linalg.generic. The elementwise op will be transoformed into
linalg.generic after -convert-elementwise-to-linalg. So the
transformation (-linalg-fuse-elementwise-ops) is expected to be run after
-convert-elementwise-to-linalg.
AI assisted.
---
.../Dialect/Linalg/Transforms/Transforms.h | 9 +
.../Linalg/Transforms/ElementwiseOpFusion.cpp | 308 +++++++++++++
.../fusion-elementwise-ops-with-concat.mlir | 411 ++++++++++++++++++
3 files changed, 728 insertions(+)
create mode 100644 mlir/test/Dialect/Linalg/fusion-elementwise-ops-with-concat.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 130f7f6fd2792..7b0a428275068 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -2007,6 +2007,15 @@ void populateElementwiseOpsFusionPatterns(
RewritePatternSet &patterns,
const ControlFusionFn &controlElementwiseOpFusion);
+/// Patterns that split elementwise `linalg.generic` operations at the
+/// boundaries of compatible `tensor.concat` inputs, exposing more producer
+/// operations to elementwise fusion. Tensor elementwise operations are
+/// represented as `linalg.generic` after `-convert-elementwise-to-linalg`, so
+/// the patterns implement the elementwise/concat interchange in the Linalg
+/// fusion pipeline.
+void populateSplitElementwiseOpsWithConcatInputsPatterns(
+ RewritePatternSet &patterns);
+
/// Function type which is used to control propagation of linalg.pack/unpack
/// ops.
using ControlPropagationFn = std::function<bool(OpOperand *opOperand)>;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index db46de75abd1a..7da547ccc7374 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -17,6 +17,7 @@
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/IR/AffineExpr.h"
@@ -502,6 +503,307 @@ class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
private:
ControlFusionFn controlFn;
};
+
+/// Split an elementwise operation at the boundaries of its `tensor.concat`
+/// inputs. This exposes the producers of the concat inputs to the elementwise
+/// fusion patterns.
+///
+/// elementwise(concat(x0, x1), concat(y0, y1))
+///
+/// becomes
+///
+/// concat(elementwise(x0, x1), elementwise(y0, y1))
+///
+/// This pattern is intentionally expressed on `linalg.generic`: tensor
+/// elementwise operations such as `arith.addf` are converted to that form by
+/// `-convert-elementwise-to-linalg`, before this pattern runs as a preamble to
+/// Linalg elementwise fusion.
+///
+/// A partition is one input of a concat, viewed as a contiguous interval of
+/// the concat dimension. All concat inputs must have matching partitions: the
+/// same number of partitions with the same static size at each index.
+///
+/// All concat operands must partition the same iteration-space dimension into
+/// the same statically-sized pieces. Inputs that do not use that iteration
+/// dimension (for example, broadcast inputs) can be shared by all pieces.
+class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
+public:
+ using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(GenericOp genericOp,
+ PatternRewriter &rewriter) const override {
+ if (!genericOp.hasPureTensorSemantics() || !isElementwise(genericOp) ||
+ genericOp.hasIndexSemantics())
+ return failure();
+
+ SmallVector<tensor::ConcatOp> concatOps(genericOp.getNumDpsInputs());
+ std::optional<unsigned> splitLoopDim;
+ // How many inputs the concat ops have. Now we limit the concat ops
+ // to have the same number of inputs for simplicity.
+ // TODO: technically, elementwise(concat(x0, x1), concat(y0, y1, y2)) ->
+ // concat(elementwise(...), elementwise(...), elementwise(...)) may be
+ // fine too. But that may require we create new slices, which might be
+ // more complex.
+ unsigned numPartitions = 0;
+ // The size in the concat dimension of different inputs. For example,
+ //
+ // x0: tensor<2x3xf32>
+ // x1: tensor<2x4xf32>
+ // x: tensor<2x7xf32>
+ // %x = tensor.concat dim(1) %x0, %x1
+ //
+ // The numPartitions in this case is 2 and partitionSizes is [3, 4].
+ // Same as above, we limit the partitionSizes to be the same for different
+ // concat ops.
+ SmallVector<int64_t> partitionSizes;
+
+ for (auto [index, operand] :
+ llvm::enumerate(genericOp.getDpsInputOperands())) {
+ auto operandType = dyn_cast<RankedTensorType>(operand->get().getType());
+ if (!operandType)
+ continue;
+
+ auto concatOp = operand->get().getDefiningOp<tensor::ConcatOp>();
+ if (!concatOp)
+ continue;
+
+ // Rewriting a concat that has other consumers could increase the amount
+ // of live computation instead of just exposing fusion opportunities.
+ if (llvm::any_of(concatOp.getResult().getUses(), [&](OpOperand &use) {
+ return use.getOwner() != genericOp.getOperation();
+ }))
+ return rewriter.notifyMatchFailure(genericOp,
+ "concat input has another consumer");
+
+ AffineMap inputMap = genericOp.getMatchingIndexingMap(operand);
+ auto concatDimExpr =
+ dyn_cast<AffineDimExpr>(inputMap.getResult(concatOp.getDim()));
+ if (!concatDimExpr)
+ return rewriter.notifyMatchFailure(
+ genericOp, "concat dimension does not map to a loop dimension");
+
+ unsigned currentSplitLoopDim = concatDimExpr.getPosition();
+ if (splitLoopDim && *splitLoopDim != currentSplitLoopDim)
+ return rewriter.notifyMatchFailure(
+ genericOp, "concat inputs partition different loop dimensions");
+ splitLoopDim = currentSplitLoopDim;
+
+ if (!numPartitions) {
+ numPartitions = concatOp.getInputs().size();
+ for (Value input : concatOp.getInputs()) {
+ int64_t size = cast<RankedTensorType>(input.getType())
+ .getDimSize(concatOp.getDim());
+ if (ShapedType::isDynamic(size))
+ return rewriter.notifyMatchFailure(
+ genericOp, "concat partition size is dynamic");
+ partitionSizes.push_back(size);
+ }
+ } else {
+ if (concatOp.getInputs().size() != numPartitions)
+ return rewriter.notifyMatchFailure(
+ genericOp, "concat inputs have different partition counts");
+ for (auto [input, expectedSize] :
+ llvm::zip_equal(concatOp.getInputs(), partitionSizes)) {
+ int64_t size = cast<RankedTensorType>(input.getType())
+ .getDimSize(concatOp.getDim());
+ if (size != expectedSize)
+ return rewriter.notifyMatchFailure(
+ genericOp, "concat inputs have different partition sizes");
+ }
+ }
+ concatOps[index] = concatOp;
+ }
+
+ if (!splitLoopDim)
+ return rewriter.notifyMatchFailure(genericOp, "has no concat input");
+
+ // A tensor input that varies along the split dimension must itself be a
+ // compatible concat. Inputs that are invariant along that dimension can be
+ // reused by every split operation.
+ AffineExpr splitDimExpr =
+ getAffineDimExpr(*splitLoopDim, genericOp.getContext());
+ for (auto [index, operand] :
+ llvm::enumerate(genericOp.getDpsInputOperands())) {
+ if (concatOps[index])
+ continue;
+ Type operandType = operand->get().getType();
+ // Scalars do not vary along an iteration-space dimension and can be
+ // reused in every partition. Other shaped inputs must be ranked tensors
+ // so that their indexing map can be checked below.
+ if (!isa<RankedTensorType>(operandType)) {
+ if (!isa<ShapedType>(operandType))
+ continue;
+ return rewriter.notifyMatchFailure(
+ genericOp, "non-concat shaped input is not a ranked tensor");
+ }
+ if (genericOp.getMatchingIndexingMap(operand).getResultPosition(
+ splitDimExpr))
+ return rewriter.notifyMatchFailure(
+ genericOp, "non-concat input varies along the split dimension");
+ }
+
+ // Map the iteration-space split dimension to each output's physical tensor
+ // dimension. `isElementwise` guarantees that the output maps are
+ // permutations, so every output contains this dimension.
+ SmallVector<unsigned> outputConcatDims;
+ outputConcatDims.reserve(genericOp.getNumDpsInits());
+ for (OpOperand &output : genericOp.getDpsInitsMutable()) {
+ std::optional<unsigned> outputDim =
+ genericOp.getMatchingIndexingMap(&output).getResultPosition(
+ splitDimExpr);
+ // Otherwise the generic op is not elementwise.
+ assert(outputDim &&
+ "elementwise output map must contain the split dimension");
+ outputConcatDims.push_back(*outputDim);
+ }
+
+ Location loc = genericOp.getLoc();
+ // Keep each result in a separate list because a generic can have multiple
+ // outputs.
+ SmallVector<SmallVector<Value>> splitResults(genericOp->getNumResults());
+ SmallVector<int64_t> outputOffsets(genericOp.getNumDpsInits(), 0);
+
+ // Build one generic for each aligned concat partition. For
+ // `elementwise(concat(x0, x1), concat(y0, y1))` becomes
+ // `elementwise(x0, y0)` and `elementwise(x1, y1)`.
+ for (auto [partitionIndex, partitionSize] :
+ llvm::enumerate(partitionSizes)) {
+ // For each partition, turn inputs `concat(x0, x1)`, `concat(y0, y1)`, and
+ // `scalar` into the inputs `x0`, `y0`, and `scalar`.
+ SmallVector<Value> inputs =
+ getPartitionInputs(genericOp, concatOps, partitionIndex);
+ // For a `tensor<7xf32>` output and partition size [3, 4], create a
+ // `tensor<3xf32>` output here; if the body reads `%init`, use
+ // `%init[0:3]` as its init instead.
+ PartitionOutputs outputs =
+ createPartitionOutputs(rewriter, loc, genericOp, outputConcatDims,
+ partitionSize, outputOffsets);
+ // For example, clone the body of
+ // `elementwise(concat(x0, x1), concat(y0, y1))` as
+ // `elementwise(x0, y0)`, retaining its indexing maps and iterators.
+ GenericOp splitOp =
+ cloneGenericForPartition(rewriter, loc, genericOp, inputs, outputs);
+ for (auto [resultIndex, result] : llvm::enumerate(splitOp->getResults()))
+ splitResults[resultIndex].push_back(result);
+ }
+
+ // Reassemble every original result from its partition results. The concat
+ // dimension may differ for each output due to its indexing map.
+ SmallVector<Value> replacements;
+ replacements.reserve(genericOp->getNumResults());
+ for (auto [resultIndex, result] :
+ llvm::enumerate(genericOp->getResults())) {
+ replacements.push_back(tensor::ConcatOp::create(
+ rewriter, loc, cast<RankedTensorType>(result.getType()),
+ outputConcatDims[resultIndex], splitResults[resultIndex]));
+ }
+ rewriter.replaceOp(genericOp, replacements);
+ return success();
+ }
+
+private:
+ /// Return the inputs for one concat partition.
+ /// For example, for `elementwise(concat(x0, x1), concat(y0, y1), scalar)`,
+ /// partition 0 uses `(x0, y0, scalar)` and partition 1 uses
+ /// `(x1, y1, scalar)`.
+ static SmallVector<Value>
+ getPartitionInputs(GenericOp genericOp, ArrayRef<tensor::ConcatOp> concatOps,
+ unsigned partitionIndex) {
+ SmallVector<Value> inputs;
+ inputs.reserve(genericOp.getNumDpsInputs());
+ for (auto [index, operand] :
+ llvm::enumerate(genericOp.getDpsInputOperands())) {
+ if (!concatOps[index]) {
+ inputs.push_back(operand->get());
+ continue;
+ }
+ tensor::ConcatOp concatOp = concatOps[index];
+ inputs.push_back(concatOp.getInputs()[partitionIndex]);
+ }
+ return inputs;
+ }
+
+ struct PartitionOutputs {
+ SmallVector<Value> values;
+ SmallVector<Type> resultTypes;
+ };
+
+ /// Create the output operands and result types for one concat partition.
+ /// For an output `tensor<7xf32>` split into sizes `[3, 4]`, this creates a
+ /// `tensor<3xf32>` output for partition 0 and `tensor<4xf32>` for partition
+ /// 1. If the generic body reads its output block argument, the outputs are
+ /// slices of the original init tensor, e.g. `%init[0:3]` and `%init[3:7]`;
+ /// otherwise they are `tensor.empty` values.
+ static PartitionOutputs createPartitionOutputs(
+ PatternRewriter &rewriter, Location loc, GenericOp genericOp,
+ ArrayRef<unsigned> outputConcatDims, int64_t partitionSize,
+ MutableArrayRef<int64_t> outputOffsets) {
+ PartitionOutputs partitionOutputs;
+ partitionOutputs.values.reserve(genericOp.getNumDpsInits());
+ partitionOutputs.resultTypes.reserve(genericOp->getNumResults());
+ for (auto [outputIndex, output] :
+ llvm::enumerate(genericOp.getDpsInitsMutable())) {
+ Value outputValue = output.get();
+ auto outputType = cast<RankedTensorType>(outputValue.getType());
+ unsigned outputConcatDim = outputConcatDims[outputIndex];
+ SmallVector<int64_t> partitionShape(outputType.getShape());
+ partitionShape[outputConcatDim] = partitionSize;
+ auto partitionType =
+ RankedTensorType::get(partitionShape, outputType.getElementType(),
+ outputType.getEncoding());
+
+ SmallVector<OpFoldResult> sizes =
+ tensor::getMixedSizes(rewriter, loc, outputValue);
+ sizes[outputConcatDim] = rewriter.getIndexAttr(partitionSize);
+
+ Value partitionOutput;
+ // A body such as `linalg.yield %in` does not read `%out`, so the init
+ // value is irrelevant and this partition can use `tensor.empty`.
+ if (!genericOp.payloadUsesValueFromOperand(&output)) {
+ partitionOutput = tensor::EmptyOp::create(rewriter, loc, sizes,
+ outputType.getElementType(),
+ outputType.getEncoding());
+ } else {
+ // A body such as `%sum = arith.addf %in, %out` reads the init value.
+ // For a `[3, 4]` partitioning, extract `%init[0:3]` for partition 0
+ // and `%init[3:7]` for partition 1 to preserve that value.
+ SmallVector<OpFoldResult> offsets(outputType.getRank(),
+ rewriter.getIndexAttr(0));
+ SmallVector<OpFoldResult> strides(outputType.getRank(),
+ rewriter.getIndexAttr(1));
+ offsets[outputConcatDim] =
+ rewriter.getIndexAttr(outputOffsets[outputIndex]);
+ partitionOutput = tensor::ExtractSliceOp::create(
+ rewriter, loc, partitionType, outputValue, offsets, sizes, strides);
+ }
+ partitionOutputs.values.push_back(partitionOutput);
+ partitionOutputs.resultTypes.push_back(partitionType);
+ outputOffsets[outputIndex] += partitionSize;
+ }
+ return partitionOutputs;
+ }
+
+ /// Clone `genericOp` for one partition while preserving its computation and
+ /// relevant attributes. For example, this turns the `x0, y0` inputs from
+ /// `getPartitionInputs` into an `elementwise(x0, y0)` generic with the same
+ /// body as the original `elementwise(concat(x0, x1), concat(y0, y1))`.
+ static GenericOp cloneGenericForPartition(PatternRewriter &rewriter,
+ Location loc, GenericOp genericOp,
+ ArrayRef<Value> inputs,
+ const PartitionOutputs &outputs) {
+ GenericOp splitOp = GenericOp::create(
+ rewriter, loc, outputs.resultTypes, inputs, outputs.values,
+ genericOp.getIndexingMapsArray(), genericOp.getIteratorTypesArray());
+ if (StringAttr doc = genericOp.getDocAttr())
+ splitOp->setAttr(splitOp.getDocAttrName(), doc);
+ if (StringAttr libraryCall = genericOp.getLibraryCallAttr())
+ splitOp->setAttr(splitOp.getLibraryCallAttrName(), libraryCall);
+ splitOp->setDiscardableAttrs(genericOp->getDiscardableAttrDictionary());
+ rewriter.cloneRegionBefore(genericOp.getRegion(), splitOp.getRegion(),
+ splitOp.getRegion().begin());
+ return splitOp;
+ }
+};
} // namespace
//===---------------------------------------------------------------------===//
@@ -2461,6 +2763,11 @@ void mlir::linalg::populateElementwiseOpsFusionPatterns(
populateEraseUnusedOperandsAndResultsPatterns(patterns);
}
+void mlir::linalg::populateSplitElementwiseOpsWithConcatInputsPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<SplitElementwiseOpWithConcatInputs>(patterns.getContext());
+}
+
void mlir::linalg::populateCollapseDimensions(
RewritePatternSet &patterns,
const GetCollapsableDimensionsFn &controlCollapseDimensions) {
@@ -2498,6 +2805,7 @@ struct LinalgElementwiseOpFusionPass
};
// Add elementwise op fusion patterns.
+ populateSplitElementwiseOpsWithConcatInputsPatterns(patterns);
populateElementwiseOpsFusionPatterns(patterns, defaultControlFn);
populateFoldReshapeOpsByExpansionPatterns(patterns, defaultControlFn);
tensor::populateBubbleUpExpandShapePatterns(patterns);
diff --git a/mlir/test/Dialect/Linalg/fusion-elementwise-ops-with-concat.mlir b/mlir/test/Dialect/Linalg/fusion-elementwise-ops-with-concat.mlir
new file mode 100644
index 0000000000000..bbb4270d37f9f
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/fusion-elementwise-ops-with-concat.mlir
@@ -0,0 +1,411 @@
+// RUN: mlir-opt %s -linalg-fuse-elementwise-ops -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -convert-elementwise-to-linalg -linalg-fuse-elementwise-ops -split-input-file | FileCheck %s --check-prefix=PIPELINE
+
+#identity = affine_map<(d0, d1) -> (d0, d1)>
+#transpose = affine_map<(d0, d1) -> (d1, d0)>
+
+// CHECK-LABEL: func.func @split_elementwise_with_concat
+// CHECK-SAME: %[[A0:.+]]: tensor<2x3xf32>, %[[A1:.+]]: tensor<2x4xf32>,
+// CHECK-SAME: %[[B0:.+]]: tensor<3x2xf32>, %[[B1:.+]]: tensor<4x2xf32>,
+// CHECK-SAME: %[[S:.+]]: f32
+// CHECK-NOT: tensor.concat
+// CHECK: %[[EMPTY0:.+]] = tensor.empty() : tensor<2x3xf32>
+// CHECK: %[[PART0:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[A0]], %[[B0]], %[[S]] : tensor<2x3xf32>, tensor<3x2xf32>, f32)
+// CHECK-SAME: outs(%[[EMPTY0]] : tensor<2x3xf32>)
+// CHECK: %[[EMPTY1:.+]] = tensor.empty() : tensor<2x4xf32>
+// CHECK: %[[PART1:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[A1]], %[[B1]], %[[S]] : tensor<2x4xf32>, tensor<4x2xf32>, f32)
+// CHECK-SAME: outs(%[[EMPTY1]] : tensor<2x4xf32>)
+// CHECK: %[[RESULT:.+]] = tensor.concat dim(1) %[[PART0]], %[[PART1]]
+// CHECK-SAME: (tensor<2x3xf32>, tensor<2x4xf32>) -> tensor<2x7xf32>
+// CHECK: return %[[RESULT]]
+func.func @split_elementwise_with_concat(
+ %a0: tensor<2x3xf32>, %a1: tensor<2x4xf32>,
+ %b0: tensor<3x2xf32>, %b1: tensor<4x2xf32>, %s: f32)
+ -> tensor<2x7xf32> {
+ %a = tensor.concat dim(1) %a0, %a1
+ : (tensor<2x3xf32>, tensor<2x4xf32>) -> tensor<2x7xf32>
+ %b = tensor.concat dim(0) %b0, %b1
+ : (tensor<3x2xf32>, tensor<4x2xf32>) -> tensor<7x2xf32>
+ %empty = tensor.empty() : tensor<2x7xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #transpose,
+ affine_map<(d0, d1) -> ()>, #identity],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%a, %b, %s : tensor<2x7xf32>, tensor<7x2xf32>, f32)
+ outs(%empty : tensor<2x7xf32>) {
+ ^bb0(%lhs: f32, %rhs: f32, %scalar: f32, %out: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ %scaled = arith.mulf %sum, %scalar : f32
+ linalg.yield %scaled : f32
+ } -> tensor<2x7xf32>
+ return %result : tensor<2x7xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Preserve a destination that is read by slicing it at the same boundaries.
+// CHECK-LABEL: func.func @split_elementwise_with_accumulator
+// CHECK: %[[INIT0:.+]] = tensor.extract_slice %[[INIT:.+]][0] [3] [1]
+// CHECK: %[[PART0:.+]] = linalg.generic
+// CHECK-SAME: outs(%[[INIT0]] : tensor<3xf32>)
+// CHECK: %[[INIT1:.+]] = tensor.extract_slice %[[INIT]][3] [4] [1]
+// CHECK: %[[PART1:.+]] = linalg.generic
+// CHECK-SAME: outs(%[[INIT1]] : tensor<4xf32>)
+// CHECK: tensor.concat dim(0) %[[PART0]], %[[PART1]]
+func.func @split_elementwise_with_accumulator(
+ %a0: tensor<3xf32>, %a1: tensor<4xf32>, %init: tensor<7xf32>)
+ -> tensor<7xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%a : tensor<7xf32>) outs(%init : tensor<7xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %sum = arith.addf %in, %out : f32
+ linalg.yield %sum : f32
+ } -> tensor<7xf32>
+ return %result : tensor<7xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Splitting the consumer makes its elementwise producers directly fusable.
+// CHECK-LABEL: func.func @split_exposes_elementwise_fusion
+// CHECK-COUNT-2: linalg.generic
+// CHECK-NOT: linalg.generic
+// CHECK: tensor.concat dim(0)
+func.func @split_exposes_elementwise_fusion(
+ %x0: tensor<3xf32>, %x1: tensor<4xf32>,
+ %y0: tensor<3xf32>, %y1: tensor<4xf32>) -> tensor<7xf32> {
+ %c = arith.constant 2.0 : f32
+ %x0e = tensor.empty() : tensor<3xf32>
+ %px0 = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%x0 : tensor<3xf32>) outs(%x0e : tensor<3xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.mulf %in, %c : f32
+ linalg.yield %v : f32
+ } -> tensor<3xf32>
+ %x1e = tensor.empty() : tensor<4xf32>
+ %px1 = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%x1 : tensor<4xf32>) outs(%x1e : tensor<4xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.mulf %in, %c : f32
+ linalg.yield %v : f32
+ } -> tensor<4xf32>
+ %y0e = tensor.empty() : tensor<3xf32>
+ %py0 = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%y0 : tensor<3xf32>) outs(%y0e : tensor<3xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.mulf %in, %c : f32
+ linalg.yield %v : f32
+ } -> tensor<3xf32>
+ %y1e = tensor.empty() : tensor<4xf32>
+ %py1 = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%y1 : tensor<4xf32>) outs(%y1e : tensor<4xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = arith.mulf %in, %c : f32
+ linalg.yield %v : f32
+ } -> tensor<4xf32>
+ %x = tensor.concat dim(0) %px0, %px1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %y = tensor.concat dim(0) %py0, %py1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %empty = tensor.empty() : tensor<7xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%x, %y : tensor<7xf32>, tensor<7xf32>)
+ outs(%empty : tensor<7xf32>) {
+ ^bb0(%lhs: f32, %rhs: f32, %out: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ linalg.yield %sum : f32
+ } -> tensor<7xf32>
+ return %result : tensor<7xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Dynamic partition boundaries cannot be proven to line up.
+// CHECK-LABEL: func.func @dynamic_partitions_not_split
+// CHECK-COUNT-2: tensor.concat
+// CHECK: linalg.generic
+func.func @dynamic_partitions_not_split(
+ %a0: tensor<?xf32>, %a1: tensor<?xf32>,
+ %b0: tensor<?xf32>, %b1: tensor<?xf32>) -> tensor<?xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+ %b = tensor.concat dim(0) %b0, %b1
+ : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+ %c0 = arith.constant 0 : index
+ %dim = tensor.dim %a, %c0 : tensor<?xf32>
+ %empty = tensor.empty(%dim) : tensor<?xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%a, %b : tensor<?xf32>, tensor<?xf32>)
+ outs(%empty : tensor<?xf32>) {
+ ^bb0(%lhs: f32, %rhs: f32, %out: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ linalg.yield %sum : f32
+ } -> tensor<?xf32>
+ return %result : tensor<?xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Reductions are not elementwise and reordering them could change numerical
+// behavior.
+// CHECK-LABEL: func.func @reduction_not_split
+// CHECK: tensor.concat
+// CHECK: linalg.generic
+func.func @reduction_not_split(%a0: tensor<3xf32>, %a1: tensor<4xf32>)
+ -> tensor<f32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %empty = tensor.empty() : tensor<f32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, affine_map<(d0) -> ()>],
+ iterator_types = ["reduction"]}
+ ins(%a : tensor<7xf32>) outs(%empty : tensor<f32>) {
+ ^bb0(%in: f32, %out: f32):
+ %sum = arith.addf %in, %out : f32
+ linalg.yield %sum : f32
+ } -> tensor<f32>
+ return %result : tensor<f32>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Splitting would reset linalg.index at every partition boundary.
+// CHECK-LABEL: func.func @index_semantics_not_split
+// CHECK: tensor.concat
+// CHECK: linalg.generic
+// CHECK: linalg.index
+func.func @index_semantics_not_split(%a0: tensor<3xindex>,
+ %a1: tensor<4xindex>) -> tensor<7xindex> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3xindex>, tensor<4xindex>) -> tensor<7xindex>
+ %empty = tensor.empty() : tensor<7xindex>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%a : tensor<7xindex>) outs(%empty : tensor<7xindex>) {
+ ^bb0(%in: index, %out: index):
+ %index = linalg.index 0 : index
+ %sum = arith.addi %in, %index : index
+ linalg.yield %sum : index
+ } -> tensor<7xindex>
+ return %result : tensor<7xindex>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Different concat counts require slicing one input at a boundary introduced
+// by the other concat, which this pattern intentionally does not do.
+// CHECK-LABEL: func.func @different_partition_counts_not_split
+// CHECK-COUNT-2: tensor.concat
+// CHECK: linalg.generic
+func.func @different_partition_counts_not_split(
+ %a0: tensor<3xf32>, %a1: tensor<4xf32>,
+ %b0: tensor<3xf32>, %b1: tensor<2xf32>, %b2: tensor<2xf32>)
+ -> tensor<7xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %b = tensor.concat dim(0) %b0, %b1, %b2
+ : (tensor<3xf32>, tensor<2xf32>, tensor<2xf32>) -> tensor<7xf32>
+ %empty = tensor.empty() : tensor<7xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%a, %b : tensor<7xf32>, tensor<7xf32>)
+ outs(%empty : tensor<7xf32>) {
+ ^bb0(%lhs: f32, %rhs: f32, %out: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ linalg.yield %sum : f32
+ } -> tensor<7xf32>
+ return %result : tensor<7xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Equal concat counts are insufficient when their boundaries differ.
+// CHECK-LABEL: func.func @different_partition_sizes_not_split
+// CHECK-COUNT-2: tensor.concat
+// CHECK: linalg.generic
+func.func @different_partition_sizes_not_split(
+ %a0: tensor<3xf32>, %a1: tensor<4xf32>,
+ %b0: tensor<2xf32>, %b1: tensor<5xf32>) -> tensor<7xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %b = tensor.concat dim(0) %b0, %b1
+ : (tensor<2xf32>, tensor<5xf32>) -> tensor<7xf32>
+ %empty = tensor.empty() : tensor<7xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%a, %b : tensor<7xf32>, tensor<7xf32>)
+ outs(%empty : tensor<7xf32>) {
+ ^bb0(%lhs: f32, %rhs: f32, %out: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ linalg.yield %sum : f32
+ } -> tensor<7xf32>
+ return %result : tensor<7xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Keep a concat with another live consumer instead of duplicating the
+// elementwise computation while retaining the original concat.
+// CHECK-LABEL: func.func @concat_with_another_consumer_not_split
+// CHECK: tensor.concat
+// CHECK: call @consume_tensor
+// CHECK: linalg.generic
+func.func private @consume_tensor(%arg: tensor<7xf32>)
+
+func.func @concat_with_another_consumer_not_split(
+ %a0: tensor<3xf32>, %a1: tensor<4xf32>) -> tensor<7xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ func.call @consume_tensor(%a) : (tensor<7xf32>) -> ()
+ %empty = tensor.empty() : tensor<7xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%a : tensor<7xf32>) outs(%empty : tensor<7xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %twice = arith.addf %in, %in : f32
+ linalg.yield %twice : f32
+ } -> tensor<7xf32>
+ return %result : tensor<7xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0, d1) -> (d0, d1)>
+#broadcast_second_dim = affine_map<(d0, d1) -> (d1)>
+
+// A non-concat tensor that is broadcast along the split dimension can be
+// reused by every partition.
+// CHECK-LABEL: func.func @broadcast_input_is_reused
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%[[A0:.+]], %[[B:.+]] : tensor<3x2xf32>, tensor<2xf32>)
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%[[A1:.+]], %[[B]] : tensor<4x2xf32>, tensor<2xf32>)
+// CHECK: tensor.concat dim(0)
+func.func @broadcast_input_is_reused(
+ %a0: tensor<3x2xf32>, %a1: tensor<4x2xf32>, %b: tensor<2xf32>)
+ -> tensor<7x2xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3x2xf32>, tensor<4x2xf32>) -> tensor<7x2xf32>
+ %empty = tensor.empty() : tensor<7x2xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #broadcast_second_dim, #identity],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%a, %b : tensor<7x2xf32>, tensor<2xf32>)
+ outs(%empty : tensor<7x2xf32>) {
+ ^bb0(%in: f32, %broadcast: f32, %out: f32):
+ %sum = arith.addf %in, %broadcast : f32
+ linalg.yield %sum : f32
+ } -> tensor<7x2xf32>
+ return %result : tensor<7x2xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0, d1) -> (d0, d1)>
+
+// A non-concat tensor that varies along the split dimension would need slices.
+// CHECK-LABEL: func.func @varying_non_concat_input_not_split
+// CHECK: tensor.concat
+// CHECK: linalg.generic
+func.func @varying_non_concat_input_not_split(
+ %a0: tensor<3x2xf32>, %a1: tensor<4x2xf32>, %b: tensor<7x2xf32>)
+ -> tensor<7x2xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3x2xf32>, tensor<4x2xf32>) -> tensor<7x2xf32>
+ %empty = tensor.empty() : tensor<7x2xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity, #identity],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%a, %b : tensor<7x2xf32>, tensor<7x2xf32>)
+ outs(%empty : tensor<7x2xf32>) {
+ ^bb0(%lhs: f32, %rhs: f32, %out: f32):
+ %sum = arith.addf %lhs, %rhs : f32
+ linalg.yield %sum : f32
+ } -> tensor<7x2xf32>
+ return %result : tensor<7x2xf32>
+}
+
+// -----
+
+#identity = affine_map<(d0, d1) -> (d0, d1)>
+#transpose = affine_map<(d0, d1) -> (d1, d0)>
+
+// The output concat dimension comes from the output map, not directly from
+// the split loop dimension: splitting d0 maps to output dim(1) here.
+// CHECK-LABEL: func.func @transpose_output_concat_dimension
+// CHECK: linalg.generic
+// CHECK-SAME: outs({{.*}} : tensor<2x3xf32>)
+// CHECK: linalg.generic
+// CHECK-SAME: outs({{.*}} : tensor<2x4xf32>)
+// CHECK: tensor.concat dim(1)
+func.func @transpose_output_concat_dimension(
+ %a0: tensor<3x2xf32>, %a1: tensor<4x2xf32>) -> tensor<2x7xf32> {
+ %a = tensor.concat dim(0) %a0, %a1
+ : (tensor<3x2xf32>, tensor<4x2xf32>) -> tensor<7x2xf32>
+ %empty = tensor.empty() : tensor<2x7xf32>
+ %result = linalg.generic {
+ indexing_maps = [#identity, #transpose],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%a : tensor<7x2xf32>) outs(%empty : tensor<2x7xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ linalg.yield %in : f32
+ } -> tensor<2x7xf32>
+ return %result : tensor<2x7xf32>
+}
+
+// -----
+
+// The source-level form from the issue is converted to linalg before this
+// rewrite runs.
+// PIPELINE-LABEL: func.func @source_level_elementwise
+// PIPELINE: linalg.generic
+// PIPELINE: linalg.generic
+// PIPELINE: tensor.concat dim(0)
+func.func @source_level_elementwise(
+ %x0: tensor<3xf32>, %x1: tensor<4xf32>,
+ %y0: tensor<3xf32>, %y1: tensor<4xf32>) -> tensor<7xf32> {
+ %x = tensor.concat dim(0) %x0, %x1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %y = tensor.concat dim(0) %y0, %y1
+ : (tensor<3xf32>, tensor<4xf32>) -> tensor<7xf32>
+ %result = arith.addf %x, %y : tensor<7xf32>
+ return %result : tensor<7xf32>
+}
>From ac61be101941927edc2a33f901235bbc71bfb773 Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Wed, 5 Aug 2026 13:55:46 +0800
Subject: [PATCH 2/5] update comment
---
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 7da547ccc7374..348e764b86325 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -512,7 +512,7 @@ class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
///
/// becomes
///
-/// concat(elementwise(x0, x1), elementwise(y0, y1))
+/// concat(elementwise(x0, y0), elementwise(x1, y1))
///
/// This pattern is intentionally expressed on `linalg.generic`: tensor
/// elementwise operations such as `arith.addf` are converted to that form by
>From edbdce6a68c7c719926b05b5a732ab0b0ef17cac Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Thu, 6 Aug 2026 16:30:28 +0800
Subject: [PATCH 3/5] Address comments
---
.../Linalg/Transforms/ElementwiseOpFusion.cpp | 63 +++++++------------
1 file changed, 24 insertions(+), 39 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 348e764b86325..893bb7b66b1cc 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -537,14 +537,14 @@ class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
return failure();
SmallVector<tensor::ConcatOp> concatOps(genericOp.getNumDpsInputs());
+ SmallVector<OpOperand *> nonConcatInputs;
std::optional<unsigned> splitLoopDim;
- // How many inputs the concat ops have. Now we limit the concat ops
- // to have the same number of inputs for simplicity.
+ // Now we limit the concat ops to have the same number of inputs for
+ // simplicity.
// TODO: technically, elementwise(concat(x0, x1), concat(y0, y1, y2)) ->
// concat(elementwise(...), elementwise(...), elementwise(...)) may be
// fine too. But that may require we create new slices, which might be
// more complex.
- unsigned numPartitions = 0;
// The size in the concat dimension of different inputs. For example,
//
// x0: tensor<2x3xf32>
@@ -552,20 +552,17 @@ class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
// x: tensor<2x7xf32>
// %x = tensor.concat dim(1) %x0, %x1
//
- // The numPartitions in this case is 2 and partitionSizes is [3, 4].
- // Same as above, we limit the partitionSizes to be the same for different
- // concat ops.
- SmallVector<int64_t> partitionSizes;
+ // The partition sizes in this case are [3, 4]. Same as above, we limit the
+ // partition sizes to be the same for different concat ops.
+ SmallVector<SmallVector<int64_t>> partitionSizes;
for (auto [index, operand] :
llvm::enumerate(genericOp.getDpsInputOperands())) {
- auto operandType = dyn_cast<RankedTensorType>(operand->get().getType());
- if (!operandType)
- continue;
-
auto concatOp = operand->get().getDefiningOp<tensor::ConcatOp>();
- if (!concatOp)
+ if (!concatOp) {
+ nonConcatInputs.push_back(operand);
continue;
+ }
// Rewriting a concat that has other consumers could increase the amount
// of live computation instead of just exposing fusion opportunities.
@@ -588,44 +585,32 @@ class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
genericOp, "concat inputs partition different loop dimensions");
splitLoopDim = currentSplitLoopDim;
- if (!numPartitions) {
- numPartitions = concatOp.getInputs().size();
- for (Value input : concatOp.getInputs()) {
- int64_t size = cast<RankedTensorType>(input.getType())
- .getDimSize(concatOp.getDim());
- if (ShapedType::isDynamic(size))
- return rewriter.notifyMatchFailure(
- genericOp, "concat partition size is dynamic");
- partitionSizes.push_back(size);
- }
- } else {
- if (concatOp.getInputs().size() != numPartitions)
+ SmallVector<int64_t> currentPartitionSizes;
+ currentPartitionSizes.reserve(concatOp.getInputs().size());
+ for (Value input : concatOp.getInputs()) {
+ int64_t size = cast<RankedTensorType>(input.getType())
+ .getDimSize(concatOp.getDim());
+ if (ShapedType::isDynamic(size))
return rewriter.notifyMatchFailure(
- genericOp, "concat inputs have different partition counts");
- for (auto [input, expectedSize] :
- llvm::zip_equal(concatOp.getInputs(), partitionSizes)) {
- int64_t size = cast<RankedTensorType>(input.getType())
- .getDimSize(concatOp.getDim());
- if (size != expectedSize)
- return rewriter.notifyMatchFailure(
- genericOp, "concat inputs have different partition sizes");
- }
+ genericOp, "concat partition size is dynamic");
+ currentPartitionSizes.push_back(size);
}
+ partitionSizes.push_back(std::move(currentPartitionSizes));
concatOps[index] = concatOp;
}
- if (!splitLoopDim)
+ if (concatOps.empty())
return rewriter.notifyMatchFailure(genericOp, "has no concat input");
+ if (!llvm::all_equal(partitionSizes))
+ return rewriter.notifyMatchFailure(
+ genericOp, "concat inputs have different partition sizes");
// A tensor input that varies along the split dimension must itself be a
// compatible concat. Inputs that are invariant along that dimension can be
// reused by every split operation.
AffineExpr splitDimExpr =
getAffineDimExpr(*splitLoopDim, genericOp.getContext());
- for (auto [index, operand] :
- llvm::enumerate(genericOp.getDpsInputOperands())) {
- if (concatOps[index])
- continue;
+ for (OpOperand *operand : nonConcatInputs) {
Type operandType = operand->get().getType();
// Scalars do not vary along an iteration-space dimension and can be
// reused in every partition. Other shaped inputs must be ranked tensors
@@ -667,7 +652,7 @@ class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
// `elementwise(concat(x0, x1), concat(y0, y1))` becomes
// `elementwise(x0, y0)` and `elementwise(x1, y1)`.
for (auto [partitionIndex, partitionSize] :
- llvm::enumerate(partitionSizes)) {
+ llvm::enumerate(partitionSizes.front())) {
// For each partition, turn inputs `concat(x0, x1)`, `concat(y0, y1)`, and
// `scalar` into the inputs `x0`, `y0`, and `scalar`.
SmallVector<Value> inputs =
>From 300440b536df45d58d3f553da1c84ca85c790d4e Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Thu, 6 Aug 2026 17:39:32 +0800
Subject: [PATCH 4/5] Update
---
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 893bb7b66b1cc..048eea42c8bcd 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -599,7 +599,7 @@ class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
concatOps[index] = concatOp;
}
- if (concatOps.empty())
+ if (!splitLoopDim)
return rewriter.notifyMatchFailure(genericOp, "has no concat input");
if (!llvm::all_equal(partitionSizes))
return rewriter.notifyMatchFailure(
>From 693d86e621250a124c73423a860e427c3737edab Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <chuanqi.xcq at alibaba-inc.com>
Date: Fri, 21 Aug 2026 11:28:47 +0800
Subject: [PATCH 5/5] update
---
.../Linalg/Transforms/ElementwiseOpFusion.cpp | 17 ++++++++---------
1 file changed, 8 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 048eea42c8bcd..556e5d111b3ff 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -566,9 +566,7 @@ class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
// Rewriting a concat that has other consumers could increase the amount
// of live computation instead of just exposing fusion opportunities.
- if (llvm::any_of(concatOp.getResult().getUses(), [&](OpOperand &use) {
- return use.getOwner() != genericOp.getOperation();
- }))
+ if (!concatOp->hasOneUse())
return rewriter.notifyMatchFailure(genericOp,
"concat input has another consumer");
@@ -613,14 +611,15 @@ class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
for (OpOperand *operand : nonConcatInputs) {
Type operandType = operand->get().getType();
// Scalars do not vary along an iteration-space dimension and can be
- // reused in every partition. Other shaped inputs must be ranked tensors
- // so that their indexing map can be checked below.
- if (!isa<RankedTensorType>(operandType)) {
- if (!isa<ShapedType>(operandType))
- continue;
+ // reused in every partition.
+ if (isa<IntegerType, FloatType, IndexType, ComplexType>(operandType))
+ continue;
+
+ // Otherwise we want ranked tensors.
+ if (!isa<RankedTensorType>(operandType))
return rewriter.notifyMatchFailure(
genericOp, "non-concat shaped input is not a ranked tensor");
- }
+
if (genericOp.getMatchingIndexingMap(operand).getResultPosition(
splitDimExpr))
return rewriter.notifyMatchFailure(
More information about the Mlir-commits
mailing list