[Mlir-commits] [mlir] da6b0ba - [mlir][linalg] Split elementwise ops with concat inputs (#213630)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 24 20:02:51 PDT 2026
Author: Chuanqi Xu
Date: 2026-08-25T11:02:46+08:00
New Revision: da6b0ba7ef6ffb7652531a23f98e18f73d7dce37
URL: https://github.com/llvm/llvm-project/commit/da6b0ba7ef6ffb7652531a23f98e18f73d7dce37
DIFF: https://github.com/llvm/llvm-project/commit/da6b0ba7ef6ffb7652531a23f98e18f73d7dce37.diff
LOG: [mlir][linalg] Split elementwise ops with concat inputs (#213630)
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.
---------
Co-authored-by: yedeng.yd <yedeng.yd at alibaba-inc.com>
Added:
mlir/test/Dialect/Linalg/fusion-elementwise-ops-with-concat.mlir
Modified:
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index bb7edd934ca21..9ab50f1136ac1 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -2004,6 +2004,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 37386bee2def4..8add883163fe9 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,291 @@ 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, 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
+/// `-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());
+ SmallVector<OpOperand *> nonConcatInputs;
+ std::optional<unsigned> splitLoopDim;
+ // 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.
+ // The size in the concat dimension of
diff erent inputs. For example,
+ //
+ // x0: tensor<2x3xf32>
+ // x1: tensor<2x4xf32>
+ // x: tensor<2x7xf32>
+ // %x = tensor.concat dim(1) %x0, %x1
+ //
+ // The partition sizes in this case are [3, 4]. Same as above, we limit the
+ // partition sizes to be the same for
diff erent concat ops.
+ SmallVector<SmallVector<int64_t>> partitionSizes;
+
+ for (auto [index, operand] :
+ llvm::enumerate(genericOp.getDpsInputOperands())) {
+ auto concatOp = operand->get().getDefiningOp<tensor::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.
+ if (!concatOp->hasOneUse())
+ 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
diff erent loop dimensions");
+ splitLoopDim = currentSplitLoopDim;
+
+ 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 partition size is dynamic");
+ currentPartitionSizes.push_back(size);
+ }
+ partitionSizes.push_back(std::move(currentPartitionSizes));
+ concatOps[index] = concatOp;
+ }
+
+ if (!splitLoopDim)
+ return rewriter.notifyMatchFailure(genericOp, "has no concat input");
+ if (!llvm::all_equal(partitionSizes))
+ return rewriter.notifyMatchFailure(
+ genericOp, "concat inputs have
diff erent 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 (OpOperand *operand : nonConcatInputs) {
+ Type operandType = operand->get().getType();
+ // Scalars do not vary along an iteration-space dimension and can be
+ // 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(
+ 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.front())) {
+ // 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
diff er 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
//===---------------------------------------------------------------------===//
@@ -2482,6 +2768,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) {
@@ -2519,6 +2810,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 @
diff erent_partition_counts_not_split
+// CHECK-COUNT-2: tensor.concat
+// CHECK: linalg.generic
+func.func @
diff erent_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
diff er.
+// CHECK-LABEL: func.func @
diff erent_partition_sizes_not_split
+// CHECK-COUNT-2: tensor.concat
+// CHECK: linalg.generic
+func.func @
diff erent_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>
+}
More information about the Mlir-commits
mailing list