[Mlir-commits] [mlir] [mlir][linalg] Split elementwise ops with concat inputs (PR #213630)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 3 02:23:35 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Chuanqi Xu (ChuanqiXu9)
<details>
<summary>Changes</summary>
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.
---
Patch is 32.74 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/213630.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h (+9)
- (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+308)
- (added) mlir/test/Dialect/Linalg/fusion-elementwise-ops-with-concat.mlir (+411)
``````````diff
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) %[[PAR...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/213630
More information about the Mlir-commits
mailing list