[Mlir-commits] [mlir] [TOSA] Convert tosa.elementwise to linalg.elementwise (PR #224344)
Renato Golin
llvmlistbot at llvm.org
Thu Sep 24 05:08:58 PDT 2026
https://github.com/rengolin updated https://github.com/llvm/llvm-project/pull/224344
>From 12e4626be1d098f3dacf03afb0652a361cd7ccdd Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Wed, 16 Sep 2026 07:30:07 -0700
Subject: [PATCH 1/4] [MLIR][TOSA] Convert elementwise to linalg.elementwise
Add a new converter for elementwise TOSA operations to Linalg
elementwise operations. Most match directly, some have slightly
different semantics that need adjusting.
Fixes #221267
---
.../Conversion/TosaToLinalg/TosaToLinalg.h | 4 +
.../TosaToLinalg/TosaToLinalgNamed.cpp | 247 ++++++++++++++++++
.../TosaToLinalg/TosaToLinalgNamedPass.cpp | 11 +
.../TosaToLinalg/tosa-to-linalg-invalid.mlir | 13 +
.../TosaToLinalg/tosa-to-linalg-named.mlir | 40 +++
.../TosaToLinalg/tosa-to-linalg.mlir | 74 +++++-
6 files changed, 377 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h b/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
index c1d28528a2f90d..825b4366c5b25b 100644
--- a/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
+++ b/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
@@ -54,6 +54,10 @@ void populateTosaToLinalgNamedConversionPatterns(
const TypeConverter &converter, RewritePatternSet *patterns,
const TosaToLinalgNamedOptions &options);
+/// Returns true if `op` is a TOSA elementwise operation that the
+/// TosaToLinalgNamed conversion can lower to a `linalg.elementwise` named op.
+bool isConvertibleToLinalgElementwise(Operation *op);
+
} // namespace tosa
} // namespace mlir
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 4a11a927090bf6..97353d2dc4d674 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -18,15 +18,161 @@
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
+#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include <optional>
#include <type_traits>
using namespace mlir;
using namespace mlir::tosa;
+// Number of leading data operands the op forwards to linalg.elementwise,
+// matching its unary/binary/ternary arity group. Trailing metadata operands
+// (e.g. MulOp's shift, NegateOp's zero points) are excluded.
+static unsigned getArity(Operation *op) {
+ return llvm::TypeSwitch<Operation *, unsigned>(op)
+ .Case<tosa::AbsOp, tosa::CeilOp, tosa::FloorOp, tosa::ExpOp, tosa::LogOp,
+ tosa::RsqrtOp, tosa::SinOp, tosa::CosOp, tosa::TanhOp, tosa::ErfOp,
+ tosa::ReciprocalOp, tosa::NegateOp>([](Operation *) { return 1u; })
+ .Case<tosa::AddOp, tosa::SubOp, tosa::IntDivOp, tosa::MulOp, tosa::PowOp,
+ tosa::MaximumOp, tosa::MinimumOp>([](Operation *) { return 2u; })
+ .Case<tosa::SelectOp>([](Operation *) { return 3u; })
+ .Default([](Operation *op) { return op->getNumOperands(); });
+}
+
+// Return the linalg.elementwise kind that is semantically equivalent to the
+// given TOSA elementwise op for the provided element type, or std::nullopt if
+// the op has no direct linalg.elementwise counterpart. Only ops whose TOSA
+// lowering maps to a single linalg.elementwise function are handled; the more
+// involved lowerings (e.g. quantized negate, mul with shift, casts) are left to
+// the generic TosaToLinalg conversion.
+static std::optional<linalg::ElementwiseKind>
+getElementwiseKind(Operation *op, Type elementTy) {
+ using linalg::ElementwiseKind;
+ const bool isFloat = isa<FloatType>(elementTy);
+ const bool isSignlessInt =
+ elementTy.isSignlessInteger() && !elementTy.isInteger(1);
+
+ // Unary ops. linalg.elementwise only defines these on floating point types.
+ if (isFloat) {
+ if (isa<tosa::AbsOp>(op))
+ return ElementwiseKind::abs;
+ if (isa<tosa::CeilOp>(op))
+ return ElementwiseKind::ceil;
+ if (isa<tosa::FloorOp>(op))
+ return ElementwiseKind::floor;
+ if (isa<tosa::ExpOp>(op))
+ return ElementwiseKind::exp;
+ if (isa<tosa::LogOp>(op))
+ return ElementwiseKind::log;
+ if (isa<tosa::RsqrtOp>(op))
+ return ElementwiseKind::rsqrt;
+ if (isa<tosa::SinOp>(op))
+ return ElementwiseKind::sin;
+ if (isa<tosa::CosOp>(op))
+ return ElementwiseKind::cos;
+ if (isa<tosa::TanhOp>(op))
+ return ElementwiseKind::tanh;
+ if (isa<tosa::ErfOp>(op))
+ return ElementwiseKind::erf;
+ if (isa<tosa::ReciprocalOp>(op))
+ return ElementwiseKind::reciprocal;
+ if (isa<tosa::PowOp>(op))
+ return ElementwiseKind::powf;
+ }
+
+ // Binary ops defined on both floating point and (non-boolean) integers.
+ if (isFloat || isSignlessInt) {
+ if (isa<tosa::AddOp>(op))
+ return ElementwiseKind::add;
+ if (isa<tosa::SubOp>(op))
+ return ElementwiseKind::sub;
+ }
+
+ // Signed integer division.
+ if (isSignlessInt && isa<tosa::IntDivOp>(op))
+ return ElementwiseKind::div;
+
+ // Multiply maps to a plain elementwise multiply only when no rescaling shift
+ // is requested (the shift must be a constant zero) and the inputs already
+ // share the result element type. The shifted/rescaled integer variant is left
+ // to the generic conversion.
+ if (auto mulOp = dyn_cast<tosa::MulOp>(op)) {
+ DenseElementsAttr shiftAttr;
+ const bool zeroShift =
+ matchPattern(mulOp.getShift(), m_Constant(&shiftAttr)) &&
+ shiftAttr.getValues<APInt>()[0].isZero();
+ const bool matchingTypes =
+ cast<ShapedType>(mulOp.getInput1().getType()).getElementType() ==
+ elementTy &&
+ cast<ShapedType>(mulOp.getInput2().getType()).getElementType() ==
+ elementTy;
+ if ((isFloat || isSignlessInt) && zeroShift && matchingTypes)
+ return ElementwiseKind::mul;
+ }
+
+ // Negate maps to floating-point negation only when the zero points are zero.
+ // The quantized integer variant is left to the generic conversion.
+ if (auto negateOp = dyn_cast<tosa::NegateOp>(op)) {
+ if (isFloat) {
+ FailureOr<int64_t> inZp = negateOp.getInput1ZeroPoint();
+ FailureOr<int64_t> outZp = negateOp.getOutputZeroPoint();
+ if (succeeded(inZp) && succeeded(outZp) && *inZp == 0 && *outZp == 0)
+ return ElementwiseKind::negf;
+ }
+ }
+
+ // Maximum/minimum. The NaN "IGNORE" mode requires an additional compare and
+ // select that linalg.elementwise cannot express, so only the default
+ // "PROPAGATE" mode is handled for floating point types.
+ auto nanModeIsPropagate = [](auto tosaOp) {
+ return tosaOp.getNanMode() == NanPropagationMode::PROPAGATE;
+ };
+ if (isFloat || isSignlessInt) {
+ if (auto maxOp = dyn_cast<tosa::MaximumOp>(op)) {
+ if (!isFloat || nanModeIsPropagate(maxOp))
+ return ElementwiseKind::max_signed;
+ }
+ if (auto minOp = dyn_cast<tosa::MinimumOp>(op)) {
+ if (!isFloat || nanModeIsPropagate(minOp))
+ return ElementwiseKind::min_signed;
+ }
+ }
+
+ // Select (ternary) is valid for any element type.
+ if (isa<tosa::SelectOp>(op))
+ return ElementwiseKind::select;
+
+ return std::nullopt;
+}
+
+bool mlir::tosa::isConvertibleToLinalgElementwise(Operation *op) {
+ if (op->getNumResults() != 1)
+ return false;
+
+ auto resultTy = dyn_cast<ShapedType>(op->getResult(0).getType());
+ if (!resultTy || !resultTy.hasRank())
+ return false;
+
+ if (!getElementwiseKind(op, resultTy.getElementType()))
+ return false;
+
+ // linalg.elementwise broadcasts operands through their indexing maps, so
+ // operands only need to be ranked shaped types (tensors or memrefs, possibly
+ // dynamic); their rank need not match the result's.
+ const unsigned numInputs = getArity(op);
+ for (unsigned i = 0; i < numInputs; ++i) {
+ auto operandTy = dyn_cast<ShapedType>(op->getOperand(i).getType());
+ if (!operandTy || !operandTy.hasRank())
+ return false;
+ }
+ return true;
+}
+
static mlir::Value applyPad(Location loc, Value input, ArrayRef<int64_t> pad,
TypedAttr padAttr, OpBuilder &rewriter) {
// Input should be padded only if necessary.
@@ -603,6 +749,84 @@ class DepthwiseConvConverter
}
};
+// Lower TOSA elementwise ops that have a direct linalg.elementwise named op
+// counterpart. Unlike linalg.generic, linalg.elementwise dispatches on a single
+// function kind and infers the body, so the conversion is a handful of checks
+// (kind, ranks, static shapes) followed by construction of the indexing maps
+// and the named op.
+template <typename TosaOp>
+class ElementwiseConverter : public OpConversionPattern<TosaOp> {
+public:
+ using OpConversionPattern<TosaOp>::OpConversionPattern;
+ using typename OpConversionPattern<TosaOp>::OpAdaptor;
+
+ LogicalResult
+ matchAndRewrite(TosaOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const final {
+ Operation *operation = op;
+ if (!isConvertibleToLinalgElementwise(operation))
+ return rewriter.notifyMatchFailure(
+ op, "op has no linalg.elementwise counterpart");
+
+ Location loc = op.getLoc();
+ auto resultTy = cast<RankedTensorType>(op.getType());
+ const int64_t rank = resultTy.getRank();
+
+ std::optional<linalg::ElementwiseKind> kind =
+ getElementwiseKind(operation, resultTy.getElementType());
+
+ // Trailing metadata operands (e.g. MulOp's shift, NegateOp's zero points)
+ // are not forwarded to linalg.elementwise.
+ ValueRange inputs = adaptor.getOperands().take_front(getArity(operation));
+
+ // Build one indexing map per input, broadcasting static size-1 dimensions,
+ // followed by the identity map for the output.
+ SmallVector<AffineMap> indexingMaps;
+ for (Value operand : inputs) {
+ ArrayRef<int64_t> shape = cast<ShapedType>(operand.getType()).getShape();
+ SmallVector<AffineExpr> exprs;
+ for (const auto &dim : llvm::enumerate(shape)) {
+ bool requiresBroadcast =
+ dim.value() == 1 && resultTy.getDimSize(dim.index()) != 1;
+ exprs.push_back(requiresBroadcast
+ ? rewriter.getAffineConstantExpr(0)
+ : rewriter.getAffineDimExpr(dim.index()));
+ }
+ indexingMaps.push_back(
+ AffineMap::get(rank, /*symbolCount=*/0, exprs, rewriter.getContext()));
+ }
+ indexingMaps.push_back(rewriter.getMultiDimIdentityMap(rank));
+
+ // Compute the output sizes, taking each dynamic result dimension from the
+ // first input that is not broadcast (statically 1) along that dimension.
+ SmallVector<OpFoldResult> sizes;
+ for (int64_t i = 0; i < rank; ++i) {
+ if (!resultTy.isDynamicDim(i)) {
+ sizes.push_back(rewriter.getIndexAttr(resultTy.getDimSize(i)));
+ continue;
+ }
+ Value dimSource = inputs.front();
+ for (Value operand : inputs) {
+ auto operandTy = cast<ShapedType>(operand.getType());
+ if (operandTy.getRank() == rank && operandTy.getDimSize(i) != 1) {
+ dimSource = operand;
+ break;
+ }
+ }
+ sizes.push_back(
+ tensor::DimOp::create(rewriter, loc, dimSource, i).getResult());
+ }
+
+ Value emptyTensor = tensor::EmptyOp::create(rewriter, loc, sizes,
+ resultTy.getElementType());
+
+ rewriter.replaceOpWithNewOp<linalg::ElementwiseOp>(
+ op, inputs, ValueRange{emptyTensor}, *kind,
+ rewriter.getAffineMapArrayAttr(indexingMaps));
+ return success();
+ }
+};
+
class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
public:
using OpConversionPattern<tosa::MatMulOp>::OpConversionPattern;
@@ -1138,6 +1362,29 @@ void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
TransposeConverter
>(patterns->getContext());
+ patterns->add<
+ ElementwiseConverter<tosa::AbsOp>,
+ ElementwiseConverter<tosa::CeilOp>,
+ ElementwiseConverter<tosa::FloorOp>,
+ ElementwiseConverter<tosa::ExpOp>,
+ ElementwiseConverter<tosa::LogOp>,
+ ElementwiseConverter<tosa::RsqrtOp>,
+ ElementwiseConverter<tosa::SinOp>,
+ ElementwiseConverter<tosa::CosOp>,
+ ElementwiseConverter<tosa::TanhOp>,
+ ElementwiseConverter<tosa::ErfOp>,
+ ElementwiseConverter<tosa::ReciprocalOp>,
+ ElementwiseConverter<tosa::PowOp>,
+ ElementwiseConverter<tosa::AddOp>,
+ ElementwiseConverter<tosa::SubOp>,
+ ElementwiseConverter<tosa::IntDivOp>,
+ ElementwiseConverter<tosa::MulOp>,
+ ElementwiseConverter<tosa::NegateOp>,
+ ElementwiseConverter<tosa::MaximumOp>,
+ ElementwiseConverter<tosa::MinimumOp>,
+ ElementwiseConverter<tosa::SelectOp>
+ >(patterns->getContext());
+
patterns->add<
MaxPool2dConverter
>(converter, patterns->getContext());
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
index df066a5cac2612..008951120c58d5 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -62,6 +62,17 @@ struct TosaToLinalgNamed
target.addIllegalOp<tosa::MatMulOp>();
target.addIllegalOp<tosa::TransposeOp>();
+ // TOSA elementwise ops are legal unless they can be lowered to a
+ // linalg.elementwise named op.
+ target.addDynamicallyLegalOp<
+ tosa::AbsOp, tosa::CeilOp, tosa::FloorOp, tosa::ExpOp, tosa::LogOp,
+ tosa::RsqrtOp, tosa::SinOp, tosa::CosOp, tosa::TanhOp, tosa::ErfOp,
+ tosa::ReciprocalOp, tosa::PowOp, tosa::AddOp, tosa::SubOp,
+ tosa::IntDivOp, tosa::MulOp, tosa::NegateOp, tosa::MaximumOp,
+ tosa::MinimumOp, tosa::SelectOp>([](Operation *op) {
+ return !tosa::isConvertibleToLinalgElementwise(op);
+ });
+
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
FunctionOpInterface func = getOperation();
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
index ecfd953f9ef1f4..e8a249d84dc490 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -73,3 +73,16 @@ func.func @unranked_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>)
%0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xf32>, tensor<13x26xi32>) -> tensor<*xf32>
return %0 : tensor<*xf32>
}
+
+// -----
+
+// CHECK-LABEL: @unsupportedRescaleInexactRound
+func.func @unsupportedRescaleInexactRound(%arg0 : tensor<2xi8>) -> (tensor<2xi8>) {
+ %multiplier = "tosa.const"() {values = dense<19689> : tensor<1xi32> } : () -> tensor<1xi32>
+ %shift = "tosa.const"() {values = dense<33> : tensor<1xi8> } : () -> tensor<1xi8>
+ %input_zp = "tosa.const"() {values = dense<0> : tensor<1xi8>} : () -> tensor<1xi8>
+ %output_zp = "tosa.const"() {values = dense<0> : tensor<1xi8>} : () -> tensor<1xi8>
+ // expected-error at +1 {{failed to legalize operation 'tosa.rescale'}}
+ %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {scale32 = true, rounding_mode = INEXACT_ROUND, per_channel = false, input_unsigned = false, output_unsigned = false} : (tensor<2xi8>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2xi8>
+ return %0 : tensor<2xi8>
+}
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index 801144c01829a9..0d14f47374859c 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -1078,3 +1078,43 @@ func.func @max_pool2d_nan_ignore(%arg0: tensor<1x6x34x62xf32>) -> (tensor<1x4x32
%0 = tosa.max_pool2d %arg0 {pad = array<i64: 0, 0, 0, 0>, kernel = array<i64: 3, 3>, stride = array<i64: 1, 1>, nan_mode = IGNORE} : (tensor<1x6x34x62xf32>) -> tensor<1x4x32x62xf32>
return %0: tensor<1x4x32x62xf32>
}
+
+// -----
+
+// CHECK-LABEL: @elementwise_unary_dyn
+func.func @elementwise_unary_dyn(%arg0: tensor<?x3xf32>) -> tensor<?x3xf32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]] : tensor<?x3xf32>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x3xf32>
+ // CHECK: linalg.elementwise <abs> ins(%arg0 : tensor<?x3xf32>) outs(%[[EMPTY]] : tensor<?x3xf32>) -> tensor<?x3xf32>
+ %0 = tosa.abs %arg0 : (tensor<?x3xf32>) -> tensor<?x3xf32>
+ return %0 : tensor<?x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @elementwise_binary_dyn
+func.func @elementwise_binary_dyn(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ // CHECK: %[[DIM0:.+]] = tensor.dim %arg0, %[[C0]] : tensor<?x?xf32>
+ // CHECK: %[[C1:.+]] = arith.constant 1 : index
+ // CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]] : tensor<?x?xf32>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor<?x?xf32>
+ // CHECK: linalg.elementwise <add> ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[EMPTY]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %0 = tosa.add %arg0, %arg1 : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @elementwise_ternary_select_dyn
+func.func @elementwise_ternary_select_dyn(%arg0: tensor<?x?xi1>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ // CHECK: %[[DIM0:.+]] = tensor.dim %arg0, %[[C0]] : tensor<?x?xi1>
+ // CHECK: %[[C1:.+]] = arith.constant 1 : index
+ // CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]] : tensor<?x?xi1>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor<?x?xf32>
+ // CHECK: linalg.elementwise <select> ins(%arg0, %arg1, %arg2 : tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[EMPTY]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %0 = tosa.select %arg0, %arg1, %arg2 : (tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index a803ee7d99153c..7af876b96e51fc 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -1,9 +1,13 @@
// RUN: mlir-opt --split-input-file -pass-pipeline="builtin.module(func.func(tosa-to-linalg))" %s -verify-diagnostics -o -| FileCheck %s
+// RUN: mlir-opt --split-input-file -pass-pipeline="builtin.module(func.func(tosa-to-linalg-named))" %s -verify-diagnostics -o -| FileCheck %s --check-prefix=NAMED
// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
// CHECK-LABEL: @test_abs_scalar
// CHECK-SAME: ([[ARG0:%[0-9a-zA-Z_]*]]
+
+// NAMED-LABEL: @test_abs_scalar
+// NAMED-SAME: ([[ARG0:%[0-9a-zA-Z_]*]]
func.func @test_abs_scalar(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT:%.+]] = tensor.empty() : tensor<f32>
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = []} ins([[ARG0]] : tensor<f32>) outs([[INIT]] : tensor<f32>) {
@@ -11,6 +15,8 @@ func.func @test_abs_scalar(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: [[ELEMENT:%.*]] = math.absf [[ARG1]] : f32
// CHECK: linalg.yield [[ELEMENT]] : f32
// CHECK: } -> tensor<f32>
+ // NAMED: linalg.elementwise <abs>
+ // NAMED-SAME: ins([[ARG0]]
%0 = tosa.abs %arg0 : (tensor<f32>) -> tensor<f32>
// CHECK: return [[GENERIC]] : tensor<f32>
@@ -84,6 +90,8 @@ func.func @test_abs_1d_dynamic(%arg0: tensor<?xf32>) -> tensor<?xf32> {
// CHECK-LABEL: @test_add_0d
// CHECK-SAME: [[ARG0:%[0-9a-zA-Z_]*]]:
// CHECK-SAME: [[ARG1:%[0-9a-zA-Z_]*]]:
+// NAMED-LABEL: @test_add_0d
+// NAMED: linalg.elementwise <add>
func.func @test_add_0d(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
// CHECK: [[EMPTY:%.+]] = tensor.empty() : tensor<f32>
@@ -103,6 +111,8 @@ func.func @test_add_0d(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
// CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (0, d1)>
+// NAMED: #[[$MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// NAMED: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (0, d1)>
// CHECK-LABEL: func.func @test_add_2d_broadcast(
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x1xf32>,
@@ -115,6 +125,9 @@ func.func @test_add_0d(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
// CHECK: } -> tensor<2x1xf32>
// CHECK: return %[[RESULT]] : tensor<2x1xf32>
// CHECK: }
+// NAMED-LABEL: func.func @test_add_2d_broadcast
+// NAMED: linalg.elementwise <add>
+// NAMED-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]]
func.func @test_add_2d_broadcast(%arg0: tensor<2x1xf32>, %arg1: tensor<1x1xf32>) -> tensor<2x1xf32> {
// tosa element-wise operators now require operands of equal ranks
%0 = tosa.add %arg0, %arg1 : (tensor<2x1xf32>, tensor<1x1xf32>) -> tensor<2x1xf32>
@@ -211,9 +224,14 @@ func.func @test_add_1d_broadcast_dynamic_to_static(%arg0: tensor<5xf32>, %arg1:
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (0)>
// CHECK: #[[$MAP1:.+]] = affine_map<(d0) -> (d0)>
+// NAMED: #[[$MAP0:.+]] = affine_map<(d0) -> (0)>
+// NAMED: #[[$MAP1:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_add_1d_broadcast_static_to_dynamic
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
+// NAMED-LABEL: @test_add_1d_broadcast_static_to_dynamic
+// NAMED: linalg.elementwise <add>
+// NAMED-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]]
func.func @test_add_1d_broadcast_static_to_dynamic(%arg0: tensor<1xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: %[[CONST0:.*]] = arith.constant 0 : index
@@ -234,9 +252,14 @@ func.func @test_add_1d_broadcast_static_to_dynamic(%arg0: tensor<1xf32>, %arg1:
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (0)>
// CHECK: #[[$MAP1:.+]] = affine_map<(d0) -> (d0)>
+// NAMED: #[[$MAP0:.+]] = affine_map<(d0) -> (0)>
+// NAMED: #[[$MAP1:.+]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_add_1d_broadcast_static_to_static
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
+// NAMED-LABEL: @test_add_1d_broadcast_static_to_static
+// NAMED: linalg.elementwise <add>
+// NAMED-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]]
func.func @test_add_1d_broadcast_static_to_static(%arg0: tensor<1xf32>, %arg1: tensor<3xf32>) -> tensor<3xf32> {
// CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<3xf32>
@@ -257,6 +280,8 @@ func.func @test_add_1d_broadcast_static_to_static(%arg0: tensor<1xf32>, %arg1: t
// CHECK-LABEL: @test_add_1d_matching_no_broadcast
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
+// NAMED-LABEL: @test_add_1d_matching_no_broadcast
+// NAMED: linalg.elementwise <add>
func.func @test_add_1d_matching_no_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
// CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<1xf32>
@@ -277,6 +302,8 @@ func.func @test_add_1d_matching_no_broadcast(%arg0: tensor<1xf32>, %arg1: tensor
// CHECK-LABEL: @test_add_1d_matching_static
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
+// NAMED-LABEL: @test_add_1d_matching_static
+// NAMED: linalg.elementwise <add>
func.func @test_add_1d_matching_static(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xf32> {
// CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<3xf32>
@@ -453,6 +480,26 @@ func.func @test_select_2d_one_dynamic(%arg0: tensor<2x?xi1>, %arg1: tensor<2x?xf
// -----
// CHECK-LABEL: @test_simple_f32
+// NAMED-LABEL: @test_simple_f32
+// NAMED: linalg.elementwise <tanh>
+// NAMED: linalg.elementwise <abs>
+// NAMED: linalg.elementwise <add>
+// NAMED: linalg.elementwise <sub>
+// NAMED: linalg.elementwise <mul>
+// NAMED: linalg.elementwise <negf>
+// NAMED: linalg.elementwise <powf>
+// NAMED: linalg.elementwise <rsqrt>
+// NAMED: linalg.elementwise <log>
+// NAMED: linalg.elementwise <exp>
+// NAMED: linalg.elementwise <select>
+// NAMED: linalg.elementwise <max_signed>
+// NAMED: linalg.elementwise <min_signed>
+// NAMED: linalg.elementwise <ceil>
+// NAMED: linalg.elementwise <floor>
+// NAMED: linalg.elementwise <reciprocal>
+// NAMED: linalg.elementwise <erf>
+// NAMED: linalg.elementwise <sin>
+// NAMED: linalg.elementwise <cos>
func.func @test_simple_f32(%arg0: tensor<1xf32>) -> () {
// CHECK: linalg.generic
// CHECK: tanh
@@ -639,6 +686,13 @@ func.func @test_simple_ui8(%arg0: tensor<1xui8>) -> () {
// -----
// CHECK-LABEL: @test_simple_i32
+// NAMED-LABEL: @test_simple_i32
+// NAMED: linalg.elementwise <add>
+// NAMED: linalg.elementwise <sub>
+// NAMED: linalg.elementwise <div>
+// NAMED: linalg.elementwise <select>
+// NAMED: linalg.elementwise <max_signed>
+// NAMED: linalg.elementwise <min_signed>
func.func @test_simple_i32(%arg0: tensor<1xi32>, %unsigned: tensor<1xui32>, %unsigned64: tensor<1xui64>) -> () {
// CHECK: linalg.generic
// CHECK: arith.addi
@@ -1537,18 +1591,6 @@ func.func @rescaleUnnecessaryDoubleRound(%arg0 : tensor<2xi8>) -> (tensor<2xi8>)
// -----
-func.func @unsupportedRescaleInexactRound(%arg0 : tensor<2xi8>) -> (tensor<2xi8>) {
- %multiplier = "tosa.const"() {values = dense<19689> : tensor<1xi32> } : () -> tensor<1xi32>
- %shift = "tosa.const"() {values = dense<33> : tensor<1xi8> } : () -> tensor<1xi8>
- %input_zp = "tosa.const"() {values = dense<0> : tensor<1xi8>} : () -> tensor<1xi8>
- %output_zp = "tosa.const"() {values = dense<0> : tensor<1xi8>} : () -> tensor<1xi8>
- // expected-error at +1 {{failed to legalize operation 'tosa.rescale'}}
- %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {scale32 = true, rounding_mode = INEXACT_ROUND, per_channel = false, input_unsigned = false, output_unsigned = false} : (tensor<2xi8>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2xi8>
- return %0 : tensor<2xi8>
-}
-
-// -----
-
// CHECK: #[[$MAP0:.+]] = affine_map<(d0) -> (d0)>
// CHECK: #[[$MAP1:.+]] = affine_map<(d0) -> ()>
// CHECK-LABEL: @rescale_no_const
@@ -2464,6 +2506,8 @@ func.func @reduce_max_nan_ignore(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>)
// -----
// CHECK-LABEL: @minimum_nan_propagate
+// NAMED-LABEL: @minimum_nan_propagate
+// NAMED: linalg.elementwise <min_signed>
func.func @minimum_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
// CHECK: linalg.generic
// CHECK: arith.minimumf
@@ -2477,6 +2521,8 @@ func.func @minimum_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>)
// -----
// CHECK-LABEL: @maximum_nan_propagate
+// NAMED-LABEL: @maximum_nan_propagate
+// NAMED: linalg.elementwise <max_signed>
func.func @maximum_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>) -> () {
// CHECK: linalg.generic
// CHECK: arith.maximumf
@@ -2490,6 +2536,8 @@ func.func @maximum_nan_propagate(%arg0: tensor<5x4xf32>, %arg1: tensor<5x4xf32>)
// -----
// CHECK-LABEL: @minimum_nan_ignore_int
+// NAMED-LABEL: @minimum_nan_ignore_int
+// NAMED: linalg.elementwise <min_signed>
func.func @minimum_nan_ignore_int(%arg0: tensor<5x4xi8>, %arg1: tensor<5x4xi8>) -> () {
// CHECK: linalg.generic
// CHECK: arith.minsi
@@ -2503,6 +2551,8 @@ func.func @minimum_nan_ignore_int(%arg0: tensor<5x4xi8>, %arg1: tensor<5x4xi8>)
// -----
// CHECK-LABEL: @maximum_nan_ignore_int
+// NAMED-LABEL: @maximum_nan_ignore_int
+// NAMED: linalg.elementwise <max_signed>
func.func @maximum_nan_ignore_int(%arg0: tensor<5x4xi8>, %arg1: tensor<5x4xi8>) -> () {
// CHECK: linalg.generic
// CHECK: arith.maxsi
>From 037c09ee3e568a210d7a9923dd92099af6b8b976 Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Thu, 17 Sep 2026 08:42:44 -0700
Subject: [PATCH 2/4] format
---
mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 97353d2dc4d674..bef528e8aa3551 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -792,8 +792,8 @@ class ElementwiseConverter : public OpConversionPattern<TosaOp> {
? rewriter.getAffineConstantExpr(0)
: rewriter.getAffineDimExpr(dim.index()));
}
- indexingMaps.push_back(
- AffineMap::get(rank, /*symbolCount=*/0, exprs, rewriter.getContext()));
+ indexingMaps.push_back(AffineMap::get(rank, /*symbolCount=*/0, exprs,
+ rewriter.getContext()));
}
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(rank));
>From 116a5688991f577243ea977a4f346d13c7facc9a Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Mon, 21 Sep 2026 08:23:31 -0700
Subject: [PATCH 3/4] unreachable for getArity
---
mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index bef528e8aa3551..0d4cb9f174d3d8 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -41,7 +41,7 @@ static unsigned getArity(Operation *op) {
.Case<tosa::AddOp, tosa::SubOp, tosa::IntDivOp, tosa::MulOp, tosa::PowOp,
tosa::MaximumOp, tosa::MinimumOp>([](Operation *) { return 2u; })
.Case<tosa::SelectOp>([](Operation *) { return 3u; })
- .Default([](Operation *op) { return op->getNumOperands(); });
+ .DefaultUnreachable("Invalid elementwise operation");
}
// Return the linalg.elementwise kind that is semantically equivalent to the
>From ce26bb706d4656d1da12cbf943eac5214cda01dd Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Thu, 24 Sep 2026 05:08:37 -0700
Subject: [PATCH 4/4] remove getArity helper
---
.../TosaToLinalg/TosaToLinalgNamed.cpp | 30 +++++++------------
1 file changed, 11 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 0d4cb9f174d3d8..40a7835766a5f6 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -30,20 +30,6 @@
using namespace mlir;
using namespace mlir::tosa;
-// Number of leading data operands the op forwards to linalg.elementwise,
-// matching its unary/binary/ternary arity group. Trailing metadata operands
-// (e.g. MulOp's shift, NegateOp's zero points) are excluded.
-static unsigned getArity(Operation *op) {
- return llvm::TypeSwitch<Operation *, unsigned>(op)
- .Case<tosa::AbsOp, tosa::CeilOp, tosa::FloorOp, tosa::ExpOp, tosa::LogOp,
- tosa::RsqrtOp, tosa::SinOp, tosa::CosOp, tosa::TanhOp, tosa::ErfOp,
- tosa::ReciprocalOp, tosa::NegateOp>([](Operation *) { return 1u; })
- .Case<tosa::AddOp, tosa::SubOp, tosa::IntDivOp, tosa::MulOp, tosa::PowOp,
- tosa::MaximumOp, tosa::MinimumOp>([](Operation *) { return 2u; })
- .Case<tosa::SelectOp>([](Operation *) { return 3u; })
- .DefaultUnreachable("Invalid elementwise operation");
-}
-
// Return the linalg.elementwise kind that is semantically equivalent to the
// given TOSA elementwise op for the provided element type, or std::nullopt if
// the op has no direct linalg.elementwise counterpart. Only ops whose TOSA
@@ -158,14 +144,18 @@ bool mlir::tosa::isConvertibleToLinalgElementwise(Operation *op) {
if (!resultTy || !resultTy.hasRank())
return false;
- if (!getElementwiseKind(op, resultTy.getElementType()))
+ std::optional<linalg::ElementwiseKind> kind =
+ getElementwiseKind(op, resultTy.getElementType());
+ if (!kind)
return false;
// linalg.elementwise broadcasts operands through their indexing maps, so
// operands only need to be ranked shaped types (tensors or memrefs, possibly
- // dynamic); their rank need not match the result's.
- const unsigned numInputs = getArity(op);
- for (unsigned i = 0; i < numInputs; ++i) {
+ // dynamic); their rank need not match the result's. The arity group enum
+ // (Unary=1, Binary=2, Ternary=3) gives the number of forwarded operands.
+ const unsigned arity =
+ llvm::to_underlying(linalg::getArityGroupAndKind(*kind).arityGroup);
+ for (unsigned i = 0; i < arity; ++i) {
auto operandTy = dyn_cast<ShapedType>(op->getOperand(i).getType());
if (!operandTy || !operandTy.hasRank())
return false;
@@ -774,10 +764,12 @@ class ElementwiseConverter : public OpConversionPattern<TosaOp> {
std::optional<linalg::ElementwiseKind> kind =
getElementwiseKind(operation, resultTy.getElementType());
+ const unsigned arity =
+ llvm::to_underlying(linalg::getArityGroupAndKind(*kind).arityGroup);
// Trailing metadata operands (e.g. MulOp's shift, NegateOp's zero points)
// are not forwarded to linalg.elementwise.
- ValueRange inputs = adaptor.getOperands().take_front(getArity(operation));
+ ValueRange inputs = adaptor.getOperands().take_front(arity);
// Build one indexing map per input, broadcasting static size-1 dimensions,
// followed by the identity map for the output.
More information about the Mlir-commits
mailing list