[Mlir-commits] [mlir] [MLIR][Linalg] Improve morphism semantics (PR #222959)
Renato Golin
llvmlistbot at llvm.org
Mon Sep 14 03:14:21 PDT 2026
https://github.com/rengolin updated https://github.com/llvm/llvm-project/pull/222959
>From 75d5605da202c0cb0008b51016dd30609c89a1ee Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Thu, 10 Sep 2026 07:22:10 -0700
Subject: [PATCH 01/12] cleanup specializeLinalgElementwise
---
mlir/include/mlir/Dialect/Linalg/IR/Linalg.h | 15 +
.../mlir/Dialect/Linalg/IR/LinalgInterfaces.h | 16 +-
.../Dialect/Linalg/IR/LinalgInterfaces.cpp | 23 +-
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 14 +-
.../Dialect/Linalg/Transforms/Specialize.cpp | 261 ++++++--------
...ic-by-unfolding-projected-permutation.mlir | 14 +-
.../Linalg/specialize-generic-ops.mlir | 321 ++++++++----------
7 files changed, 286 insertions(+), 378 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
index 9d46f9aeb9c15..aaf4a61b19567 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
@@ -147,6 +147,21 @@ template <typename OpTy,
std::is_same_v<OpTy, linalg::UnPackOp>>>
SmallVector<int64_t> getPackedOuterShapeWithoutTransposition(OpTy packOrUnPack);
+/// Elementwise Arity and Kind groups.
+struct ArityGroupAndKind {
+ // The enum class {Unary, Binary, Ternary, ..}
+ ElementwiseArityGroup arityGroup;
+
+ // The kind (e.g. `exp` or `add`) belonging to the arity group.
+ union Kind {
+ UnaryFn unaryFn;
+ BinaryFn binaryFn;
+ TernaryFn ternaryFn;
+ } kind;
+};
+ArityGroupAndKind getArityGroupAndKind(ElementwiseKind kind);
+
+
/// Specialization of `linalg.matmul` op that has a transpose map on A
class MatmulTransposeAOp : public MatmulOp {
/// Create an affine map for a transpose-A matmul. Used only in the builders.
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
index 28039eb8230fb..89b42f1beef7e 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
@@ -149,22 +149,18 @@ isaTransposeOpInterface(GenericOp genericOp);
/// Checks whether a given `genericOp` is semantically equivalent to a single
/// linalg elementwise unary op, e.g.
/// `linalg.elementwise <exp>`.
-/// If `allowNonIdentityMaps` is true, operations with custom indexing maps are
-/// included in the check. Note that these operations can only be represented by
-/// the category op.
+/// Operations with non-identity indexing maps are included in the check, as
+/// they can be represented by the category op.
/// A linalg.generic body could be a series of unary elementwise ops e.g.
/// `exp(neg(x))`, such as formed by linalg op fusion. Here we restrict it to
/// detecting cases where body is a single computation op.
-bool isaElemwiseSingleUnaryOpInterface(GenericOp genericOp,
- bool allowNonIdentityMaps = false);
+bool isaElemwiseSingleUnaryOpInterface(GenericOp genericOp);
/// Checks whether `genericOp` is semantically equivalent to a single linalg
/// elementwise binary op e.g. linalg.sub.
-/// If `allowNonIdentityMaps` is true, operations with custom indexing maps are
-/// included in the check. Note that these operations can only be represented by
-/// the category op.
-bool isaElemwiseSingleBinaryOpInterface(GenericOp genericOp,
- bool allowNonIdentityMaps = false);
+/// Operations with non-identity indexing maps are included in the check, as
+/// they can be represented by the category op.
+bool isaElemwiseSingleBinaryOpInterface(GenericOp genericOp);
/// Checks whether `genericOp` is semantically equivalent to a `linalg.fill`.
/// Supports two patterns:
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index cff613f389669..c8088ce183983 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -228,18 +228,15 @@ linalg::isaTransposeOpInterface(GenericOp op) {
// Elementwise Single Unary/Binary-OpInterface implementation
//===----------------------------------------------------------------------===//
static bool
-isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op, unsigned arity,
- bool allowNonIdentityMaps) {
+isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op,
+ unsigned arity) {
// Check all loops are parallel.
if (!op.isAllParallelLoops() || op.getNumLoops() < 1)
return false;
- // Check there are arity-inputs, 1-output and all are identity-maps (unless
- // requested otherwise).
- if (op.getNumDpsInputs() != arity || op.getNumDpsInits() != 1 ||
- (!allowNonIdentityMaps &&
- !llvm::all_of(op.getIndexingMapsArray(),
- [](AffineMap map) { return map.isIdentity(); })))
+ // Check there are arity-inputs and 1-output. Non-identity indexing maps are
+ // allowed as they can be represented by the category op.
+ if (op.getNumDpsInputs() != arity || op.getNumDpsInits() != 1)
return false;
// Init should not be referenced for elementwise operations.
@@ -267,10 +264,9 @@ isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op, unsigned arity,
yieldOp->getOperand(0).getDefiningOp() != oper);
}
-bool linalg::isaElemwiseSingleUnaryOpInterface(linalg::GenericOp op,
- bool allowNonIdentityMaps) {
+bool linalg::isaElemwiseSingleUnaryOpInterface(linalg::GenericOp op) {
// All basic elemwise checks.
- if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 1, allowNonIdentityMaps))
+ if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 1))
return false;
// Check input is actually used.
@@ -279,10 +275,9 @@ bool linalg::isaElemwiseSingleUnaryOpInterface(linalg::GenericOp op,
return true;
}
-bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op,
- bool allowNonIdentityMaps) {
+bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op) {
// All basic elemwise checks.
- if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 2, allowNonIdentityMaps))
+ if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 2))
return false;
// Check both inputs are used (elementwise).
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 0909732728eca..87b14439597b4 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5081,24 +5081,12 @@ Speculation::Speculatability BatchMatmulOp::getSpeculatability() {
//===----------------------------------------------------------------------===//
//
namespace {
-struct ArityGroupAndKind {
- // The enum class {Unary, Binary, Ternary, ..}
- ElementwiseArityGroup arityGroup;
-
- // The kind (e.g. `exp` or `add`) belonging to the arity group.
- union Kind {
- UnaryFn unaryFn;
- BinaryFn binaryFn;
- TernaryFn ternaryFn;
- } kind;
-};
-
unsigned getArityGroupAsUInt(ElementwiseArityGroup arityGroup) {
return static_cast<unsigned>(arityGroup);
}
} // namespace
-static ArityGroupAndKind getArityGroupAndKind(ElementwiseKind kind) {
+ArityGroupAndKind getArityGroupAndKind(ElementwiseKind kind) {
constexpr int lastUnary = static_cast<int>(ElementwiseCaseLimits::LastUnary);
constexpr int lastBinary =
static_cast<int>(ElementwiseCaseLimits::LastBinary);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index a73c466484c0d..a3a37bc5be242 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -20,6 +20,7 @@
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/TypeSwitch.h"
namespace mlir {
#define GEN_PASS_DEF_LINALGSPECIALIZEGENERICOPSPASS
@@ -90,8 +91,56 @@ static bool findIndexOfScalarOperand(GenericOp genericOp, int &index) {
return false;
}
-// Attempt to specialize unary or binary linalg.generic ops to named elementwise
-// ops or linalg.elementwise.
+// Maps a linalg.generic body operation to its corresponding elementwise kind,
+// if one exists. Handles the ops with a straightforward one-to-one mapping;
+// ops needing extra context (e.g. boolean add/mul, reciprocal, square) are
+// handled by the caller.
+static std::optional<ElementwiseKind> getElementwiseKind(Operation *op) {
+ return llvm::TypeSwitch<Operation *, std::optional<ElementwiseKind>>(op)
+ .Case<arith::AddFOp, arith::AddIOp, complex::AddOp>(
+ [](Operation *) { return ElementwiseKind::add; })
+ .Case<arith::MulIOp, arith::MulFOp, complex::MulOp>(
+ [](Operation *) { return ElementwiseKind::mul; })
+ .Case([](math::ExpOp) { return ElementwiseKind::exp; })
+ .Case([](math::AbsFOp) { return ElementwiseKind::abs; })
+ .Case([](math::CeilOp) { return ElementwiseKind::ceil; })
+ .Case([](math::FloorOp) { return ElementwiseKind::floor; })
+ .Case([](arith::NegFOp) { return ElementwiseKind::negf; })
+ .Case([](math::RoundOp) { return ElementwiseKind::round; })
+ .Case([](math::SqrtOp) { return ElementwiseKind::sqrt; })
+ .Case([](math::RsqrtOp) { return ElementwiseKind::rsqrt; })
+ .Case([](math::TanhOp) { return ElementwiseKind::tanh; })
+ .Case([](math::ErfOp) { return ElementwiseKind::erf; })
+ .Case([](math::SinOp) { return ElementwiseKind::sin; })
+ .Case([](math::CosOp) { return ElementwiseKind::cos; })
+ .Case([](math::TanOp) { return ElementwiseKind::tan; })
+ .Case([](math::AcosOp) { return ElementwiseKind::acos; })
+ .Case([](math::AcoshOp) { return ElementwiseKind::acosh; })
+ .Case([](math::AsinOp) { return ElementwiseKind::asin; })
+ .Case([](math::AsinhOp) { return ElementwiseKind::asinh; })
+ .Case([](math::AtanOp) { return ElementwiseKind::atan; })
+ .Case([](math::AtanhOp) { return ElementwiseKind::atanh; })
+ .Case([](math::LogOp) { return ElementwiseKind::log; })
+ .Case([](math::Log10Op) { return ElementwiseKind::log10; })
+ .Case([](math::Log1pOp) { return ElementwiseKind::log1p; })
+ .Case([](math::Log2Op) { return ElementwiseKind::log2; })
+ .Case<arith::SubIOp, arith::SubFOp, complex::SubOp>(
+ [](Operation *) { return ElementwiseKind::sub; })
+ .Case<arith::DivSIOp, arith::DivFOp, complex::DivOp>(
+ [](Operation *) { return ElementwiseKind::div; })
+ .Case([](arith::DivUIOp) { return ElementwiseKind::div_unsigned; })
+ .Case<arith::MaxSIOp, arith::MaximumFOp>(
+ [](Operation *) { return ElementwiseKind::max_signed; })
+ .Case<arith::MinSIOp, arith::MinimumFOp>(
+ [](Operation *) { return ElementwiseKind::min_signed; })
+ .Case([](math::PowFOp) { return ElementwiseKind::powf; })
+ .Case([](arith::MaxUIOp) { return ElementwiseKind::max_unsigned; })
+ .Case([](arith::MinUIOp) { return ElementwiseKind::min_unsigned; })
+ .Default([](Operation *) { return std::nullopt; });
+}
+
+// Attempt to specialize unary or binary linalg.generic ops
+// to linalg.elementwise.
//
// Example:
// %0 = linalg.generic {
@@ -107,10 +156,7 @@ static bool findIndexOfScalarOperand(GenericOp genericOp, int &index) {
// is specialized to
// linalg.elementwise <exp> ...
//
-// A named op is emitted instead for binary/ternary ops that still have a
-// linalg.* named equivalent (e.g. linalg.add).
-//
-// Only the category op can carry non-identity indexing maps; these are
+// The category op can carry non-identity indexing maps; these are
// transferred verbatim from the `genericOp`.
//
// In addition to the canonical forms used by the generalization path, this
@@ -120,21 +166,11 @@ static bool findIndexOfScalarOperand(GenericOp genericOp, int &index) {
// 2) Unary generic ops with a binary body op (see the
// `findIndexOfScalarOperand` helper)
static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
- GenericOp genericOp,
- bool emitCategoryOp) {
- bool hasNonIdentityMaps =
- !llvm::all_of(genericOp.getIndexingMapsArray(),
- [](AffineMap map) { return map.isIdentity(); });
-
- // Early exit: Named ops cannot carry user-defined maps.
- if (hasNonIdentityMaps && !emitCategoryOp)
- return rewriter.notifyMatchFailure(
- genericOp,
- "non-identity indexing maps prevent specialization to named op");
-
+ GenericOp genericOp) {
// Classify the generic op.
- bool isUnary = genericOp.getNumDpsInputs() == 1;
- bool isBinary = genericOp.getNumDpsInputs() == 2;
+ unsigned arity = genericOp.getNumDpsInputs();
+ bool isUnary = arity == 1;
+ bool isBinary = arity == 2;
// Will inspect the body operation to determine named op or elementwise kind.
Operation *op = &genericOp.getBody()->front();
@@ -148,149 +184,69 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
// Helper to dispatch between named op and `linalg.elementwise`.
// Lambdas with explicit template parameter list are a C++20 feature, hence
// the dummy op object.
- auto replaceOp = [&](auto namedOp, ElementwiseKind kind,
+ auto replaceOp = [&](ElementwiseKind kind,
bool mayHoistScalarOperand = true) -> LinalgOp {
SmallVector<Value> inputs = genericOp.getDpsInputs();
- if (hasSwappedOperands)
+ SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
+ if (hasSwappedOperands) {
std::swap(inputs[0], inputs[1]);
+ std::swap(indexingMaps[0], indexingMaps[1]);
+ }
- LinalgOp newOp;
- using NamedOpTy = decltype(namedOp);
- // A null named op means the op only has a category form; emit
- // `linalg.elementwise` regardless of the requested output form.
- if (!std::is_null_pointer_v<NamedOpTy>) {
- if constexpr (!std::is_null_pointer_v<NamedOpTy>)
- newOp = NamedOpTy::create(rewriter, genericOp.getLoc(), inputs,
- genericOp.getDpsInits(),
- ArrayRef<NamedAttribute>{});
- } else {
- SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
- // Swap indexing maps, too.
- if (hasSwappedOperands)
- std::swap(indexingMaps[0], indexingMaps[1]);
-
- // Represent unary generic op as a binary `linalg.elementwise` with a
- // scalar operand and broadcasting map.
- if (hasScalarOperand && mayHoistScalarOperand) {
- // Adjust inputs and indexing maps accordingly.
- inputs.insert(inputs.begin() + scalarOprIdx,
- op->getOperand(scalarOprIdx));
- auto scalarBroadcastMap =
- AffineMap::get(genericOp.getNumParallelLoops(), /*symbolCount=*/0,
- rewriter.getContext());
- indexingMaps.insert(indexingMaps.begin() + scalarOprIdx,
- scalarBroadcastMap);
- }
- newOp = ElementwiseOp::create(
- rewriter, genericOp.getLoc(), inputs, genericOp.getDpsInits(), kind,
- rewriter.getAffineMapArrayAttr(indexingMaps));
+ if (hasScalarOperand && mayHoistScalarOperand) {
+ // Adjust inputs and indexing maps accordingly.
+ inputs.insert(inputs.begin() + scalarOprIdx,
+ op->getOperand(scalarOprIdx));
+ auto scalarBroadcastMap =
+ AffineMap::get(genericOp.getNumParallelLoops(), /*symbolCount=*/0,
+ rewriter.getContext());
+ indexingMaps.insert(indexingMaps.begin() + scalarOprIdx,
+ scalarBroadcastMap);
}
+ auto newOp = ElementwiseOp::create(
+ rewriter, genericOp.getLoc(), inputs, genericOp.getDpsInits(), kind,
+ rewriter.getAffineMapArrayAttr(indexingMaps));
rewriter.replaceOp(genericOp, newOp);
return newOp;
};
- // There are no named ops for these elementwise operations; can only emit the
- // category form.
- if (emitCategoryOp) {
- if (isa<math::ExpOp>(op))
- return replaceOp(nullptr, ElementwiseKind::exp);
- if (isa<math::AbsFOp>(op))
- return replaceOp(nullptr, ElementwiseKind::abs);
- if (isa<math::CeilOp>(op))
- return replaceOp(nullptr, ElementwiseKind::ceil);
- if (isa<math::FloorOp>(op))
- return replaceOp(nullptr, ElementwiseKind::floor);
- if (isa<arith::NegFOp>(op))
- return replaceOp(nullptr, ElementwiseKind::negf);
- if (auto divOp = dyn_cast<arith::DivFOp>(op)) {
- if (auto constOp = dyn_cast_if_present<arith::ConstantOp>(
- divOp.getLhs().getDefiningOp()))
- if (cast<FloatAttr>(constOp.getValue()).getValue().isExactlyValue(1.0))
- return replaceOp(nullptr, ElementwiseKind::reciprocal,
- /*mayHoistScalarOperand=*/false);
- }
- if (isa<math::RoundOp>(op))
- return replaceOp(nullptr, ElementwiseKind::round);
- if (isa<math::SqrtOp>(op))
- return replaceOp(nullptr, ElementwiseKind::sqrt);
- if (isa<math::RsqrtOp>(op))
- return replaceOp(nullptr, ElementwiseKind::rsqrt);
- if (auto mulOp = dyn_cast<arith::MulFOp>(op);
- mulOp && mulOp.getLhs() == mulOp.getRhs())
- return replaceOp(nullptr, ElementwiseKind::square);
- if (isa<math::TanhOp>(op))
- return replaceOp(nullptr, ElementwiseKind::tanh);
- if (isa<math::ErfOp>(op))
- return replaceOp(nullptr, ElementwiseKind::erf);
- if (isa<math::SinOp>(op))
- return replaceOp(nullptr, ElementwiseKind::sin);
- if (isa<math::CosOp>(op))
- return replaceOp(nullptr, ElementwiseKind::cos);
- if (isa<math::TanOp>(op))
- return replaceOp(nullptr, ElementwiseKind::tan);
- if (isa<math::AcosOp>(op))
- return replaceOp(nullptr, ElementwiseKind::acos);
- if (isa<math::AcoshOp>(op))
- return replaceOp(nullptr, ElementwiseKind::acosh);
- if (isa<math::AsinOp>(op))
- return replaceOp(nullptr, ElementwiseKind::asin);
- if (isa<math::AsinhOp>(op))
- return replaceOp(nullptr, ElementwiseKind::asinh);
- if (isa<math::AtanOp>(op))
- return replaceOp(nullptr, ElementwiseKind::atan);
- if (isa<math::AtanhOp>(op))
- return replaceOp(nullptr, ElementwiseKind::atanh);
- if (isa<math::LogOp>(op))
- return replaceOp(nullptr, ElementwiseKind::log);
- if (isa<math::Log10Op>(op))
- return replaceOp(nullptr, ElementwiseKind::log10);
- if (isa<math::Log1pOp>(op))
- return replaceOp(nullptr, ElementwiseKind::log1p);
- if (isa<math::Log2Op>(op))
- return replaceOp(nullptr, ElementwiseKind::log2);
-
- // The remaining kinds are binary. A single-input generic can only be
- // represented as a binary elementwise if it has a scalar operand to hoist;
- // otherwise (e.g. a body reusing a block argument twice) it has no
- // category form.
- if (isUnary && !hasScalarOperand)
- return rewriter.notifyMatchFailure(
- genericOp, "unary elementwise operation cannot be specialized to a "
- "category op");
-
- // Boolean-typed `linalg.add` and `linalg.mul` require special handling.
- bool allBool = llvm::all_of(
- op->getOperands(), [](Value v) { return v.getType().isInteger(1); });
-
- if (isa<arith::AddFOp, arith::AddIOp, complex::AddOp>(op) ||
- (allBool && isa<arith::OrIOp>(op)))
- return replaceOp(nullptr, ElementwiseKind::add);
- if (isa<arith::SubIOp, arith::SubFOp, complex::SubOp>(op))
- return replaceOp(nullptr, ElementwiseKind::sub);
- if (isa<arith::MulIOp, arith::MulFOp, complex::MulOp>(op) ||
- (allBool && isa<arith::AndIOp>(op)))
- return replaceOp(nullptr, ElementwiseKind::mul);
- if (isa<arith::DivSIOp, arith::DivFOp, complex::DivOp>(op))
- return replaceOp(nullptr, ElementwiseKind::div);
- if (isa<arith::DivUIOp>(op))
- return replaceOp(nullptr, ElementwiseKind::div_unsigned);
- if (isa<arith::MaxSIOp, arith::MaximumFOp>(op))
- return replaceOp(nullptr, ElementwiseKind::max_signed);
- if (isa<arith::MinSIOp, arith::MinimumFOp>(op))
- return replaceOp(nullptr, ElementwiseKind::min_signed);
- if (isa<math::PowFOp>(op))
- return replaceOp(nullptr, ElementwiseKind::powf);
- // No named ops for unsigned maximum/minimum.
- if (isa<arith::MaxUIOp>(op))
- return replaceOp(nullptr, ElementwiseKind::max_unsigned);
- if (isa<arith::MinUIOp>(op))
- return replaceOp(nullptr, ElementwiseKind::min_unsigned);
+ // Reciprocal
+ if (auto divOp = dyn_cast<arith::DivFOp>(op)) {
+ if (auto constOp = dyn_cast_if_present<arith::ConstantOp>(
+ divOp.getLhs().getDefiningOp()))
+ if (cast<FloatAttr>(constOp.getValue()).getValue().isExactlyValue(1.0))
+ return replaceOp(ElementwiseKind::reciprocal,
+ /*mayHoistScalarOperand=*/false);
+ }
+
+ // Square
+ if (auto mulOp = dyn_cast<arith::MulFOp>(op))
+ if (mulOp.getLhs() == mulOp.getRhs())
+ return replaceOp(ElementwiseKind::square);
+
+ // Boolean-typed `add` and `mul`.
+ if (isBinary && llvm::all_of(
+ op->getOperands(), [](Value v) {
+ return v.getType().isInteger(1); })) {
+ if (isa<arith::OrIOp>(op))
+ return replaceOp(ElementwiseKind::add);
+ if (isa<arith::AndIOp>(op))
+ return replaceOp(ElementwiseKind::mul);
+ }
+
+ // Table driven
+ if (std::optional<ElementwiseKind> kind = getElementwiseKind(op)) {
+ auto arityGroupAndKind = linalg::getArityGroupAndKind(*kind);
+ // A hoisted scalar operand adds one input to the elementwise op.
+ unsigned numInputs = arity + (hasScalarOperand ? 1 : 0);
+ if (numInputs == static_cast<unsigned>(arityGroupAndKind.arityGroup))
+ return replaceOp(*kind);
}
return rewriter.notifyMatchFailure(
genericOp,
- "elementwise operation cannot be specialized to named or category op");
+ "elementwise operation cannot be specialized to category op");
}
//===----------------------------------------------------------------------===//
@@ -727,10 +683,9 @@ FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(
RewriterBase &rewriter, GenericOp genericOp,
const GenericOpSpecializationOptions &options) {
// Elementwise - e.g. exp, add
- if (isaElemwiseSingleUnaryOpInterface(genericOp, options.emitCategoryOps) ||
- isaElemwiseSingleBinaryOpInterface(genericOp, options.emitCategoryOps)) {
- return specializeLinalgElementwise(rewriter, genericOp,
- options.emitCategoryOps);
+ if (isaElemwiseSingleUnaryOpInterface(genericOp) ||
+ isaElemwiseSingleBinaryOpInterface(genericOp)) {
+ return specializeLinalgElementwise(rewriter, genericOp);
}
// Contraction - e.g. matmul
diff --git a/mlir/test/Dialect/Linalg/decompose-generic-by-unfolding-projected-permutation.mlir b/mlir/test/Dialect/Linalg/decompose-generic-by-unfolding-projected-permutation.mlir
index 991434b81c699..f7f7c8be90ab4 100644
--- a/mlir/test/Dialect/Linalg/decompose-generic-by-unfolding-projected-permutation.mlir
+++ b/mlir/test/Dialect/Linalg/decompose-generic-by-unfolding-projected-permutation.mlir
@@ -16,11 +16,7 @@ func.func @transpose_and_broadcast(%x : tensor<7x8x9xf32>, %y: tensor<5x9x7x8x1
// CHECK-LABEL: transpose_and_broadcast
// CHECK-SAME: %[[X:.+]]: tensor<7x8x9xf32>, %[[Y:.+]]: tensor<5x9x7x8x10xf32>, %[[Z:.+]]: tensor<5x9x7x8x10xf32>) -> tensor<5x9x7x8x10xf32> {
-// CHECK: %[[E0:.+]] = tensor.empty() : tensor<9x7x8xf32>
-// CHECK: %[[X_trans:.+]] = linalg.transpose ins(%[[X]] : tensor<7x8x9xf32>) outs(%[[E0]] : tensor<9x7x8xf32>) permutation = [2, 0, 1]
-// CHECK: %[[E1:.+]] = tensor.empty() : tensor<5x9x7x8x10xf32>
-// CHECK: %[[X_trans_bc:.+]] = linalg.broadcast ins(%[[X_trans]] : tensor<9x7x8xf32>) outs(%[[E1]] : tensor<5x9x7x8x10xf32>) dimensions = [0, 4]
-// CHECK: {{.*}} = linalg.elementwise <div> ins(%[[X_trans_bc]], %[[Y]] : tensor<5x9x7x8x10xf32>, tensor<5x9x7x8x10xf32>) outs(%[[Z]] : tensor<5x9x7x8x10xf32>) -> tensor<5x9x7x8x10xf32>
+// CHECK: {{.*}} = linalg.elementwise <div> indexing_maps = {{.*}} ins(%[[X]], %[[Y]] : tensor<7x8x9xf32>, tensor<5x9x7x8x10xf32>) outs(%[[Z]] : tensor<5x9x7x8x10xf32>) -> tensor<5x9x7x8x10xf32>
// CHECK-NOT: linalg.generic
// -----
@@ -42,9 +38,7 @@ func.func @transpose_only(%x : tensor<32x2x16xf32>, %y: tensor<2x16x32xf32>, %z
// CHECK-LABEL: transpose_only
// CHECK-SAME: %[[X:.+]]: tensor<32x2x16xf32>, %[[Y:.+]]: tensor<2x16x32xf32>, %[[Z:.+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
-// CHECK: %[[E0:.+]] = tensor.empty() : tensor<2x16x32xf32>
-// CHECK: %[[X_trans:.+]] = linalg.transpose ins(%[[X]] : tensor<32x2x16xf32>) outs(%[[E0]] : tensor<2x16x32xf32>) permutation = [1, 2, 0]
-// CHECK: {{.*}} = linalg.elementwise <div> ins(%[[X_trans]], %[[Y]] : tensor<2x16x32xf32>, tensor<2x16x32xf32>) outs(%[[Z]] : tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
+// CHECK: {{.*}} = linalg.elementwise <div> indexing_maps = {{.*}} ins(%[[X]], %[[Y]] : tensor<32x2x16xf32>, tensor<2x16x32xf32>) outs(%[[Z]] : tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
// CHECK-NOT: linalg.generic
// -----
@@ -65,7 +59,5 @@ func.func @broadcast_only(%x : tensor<2x16x32xf32>, %y: tensor<2x32xf32>, %z :
// CHECK-LABEL: broadcast_only
// CHECK-SAME: %[[X:.+]]: tensor<2x16x32xf32>, %[[Y:.+]]: tensor<2x32xf32>, %[[Z:.+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
-// CHECK: %[[E0:.+]] = tensor.empty() : tensor<2x16x32xf32>
-// CHECK: %[[X_bc:.+]] = linalg.broadcast ins(%[[Y]] : tensor<2x32xf32>) outs(%[[E0]] : tensor<2x16x32xf32>) dimensions = [1]
-// CHECK: {{.*}} = linalg.elementwise <div> ins(%[[X]], %[[X_bc]] : tensor<2x16x32xf32>, tensor<2x16x32xf32>) outs(%arg2 : tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
+// CHECK: {{.*}} = linalg.elementwise <div> indexing_maps = {{.*}} ins(%[[X]], %[[Y]] : tensor<2x16x32xf32>, tensor<2x32xf32>) outs(%arg2 : tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
// CHECK-NOT: linalg.generic
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index 0c6f3228c7d53..1d0fda92355d0 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -160,73 +160,69 @@ func.func @unary_ops(%A: tensor<?x?x?xf32>, %Out: tensor<?x?x?xf32>) -> tensor<?
// ALL-LABEL: unary_ops
// ALL-SAME: %[[A:.+]]: tensor<?x?x?xf32>, %[[OUT:.+]]: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// No unary linalg named ops remain, so generic-to-named leaves them as generics.
-// NAMED-NOT: linalg.elementwise
-// NAMED: linalg.generic
-
-// CATEGORY: %[[RES4:.+]] = linalg.elementwise <floor>
-// CATEGORY-SAME: ins(%[[A]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES5:.+]] = linalg.elementwise <negf>
-// CATEGORY-SAME: ins(%[[RES4]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES6:.+]] = linalg.elementwise <reciprocal>
-// CATEGORY-SAME: ins(%[[RES5]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES7:.+]] = linalg.elementwise <round>
-// CATEGORY-SAME: ins(%[[RES6]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES8:.+]] = linalg.elementwise <sqrt>
-// CATEGORY-SAME: ins(%[[RES7]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES9:.+]] = linalg.elementwise <rsqrt>
-// CATEGORY-SAME: ins(%[[RES8]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES10:.+]] = linalg.elementwise <square>
-// CATEGORY-SAME: ins(%[[RES9]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES11:.+]] = linalg.elementwise <tanh>
-// CATEGORY-SAME: ins(%[[RES10]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES12:.+]] = linalg.elementwise <erf>
-// CATEGORY-SAME: ins(%[[RES11]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES13:.+]] = linalg.elementwise <sin>
-// CATEGORY-SAME: ins(%[[RES12]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES14:.+]] = linalg.elementwise <cos>
-// CATEGORY-SAME: ins(%[[RES13]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES15:.+]] = linalg.elementwise <tan>
-// CATEGORY-SAME: ins(%[[RES14]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES16:.+]] = linalg.elementwise <acos>
-// CATEGORY-SAME: ins(%[[RES15]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES17:.+]] = linalg.elementwise <acosh>
-// CATEGORY-SAME: ins(%[[RES16]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES18:.+]] = linalg.elementwise <asin>
-// CATEGORY-SAME: ins(%[[RES17]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES19:.+]] = linalg.elementwise <asinh>
-// CATEGORY-SAME: ins(%[[RES18]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES20:.+]] = linalg.elementwise <atan>
-// CATEGORY-SAME: ins(%[[RES19]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES21:.+]] = linalg.elementwise <atanh>
-// CATEGORY-SAME: ins(%[[RES20]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES22:.+]] = linalg.elementwise <log10>
-// CATEGORY-SAME: ins(%[[RES21]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES23:.+]] = linalg.elementwise <log1p>
-// CATEGORY-SAME: ins(%[[RES22]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
-// CATEGORY: %[[RES24:.+]] = linalg.elementwise <log2>
-// CATEGORY-SAME: ins(%[[RES23]] : tensor<?x?x?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES4:.+]] = linalg.elementwise <floor>
+// ALL-SAME: ins(%[[A]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES5:.+]] = linalg.elementwise <negf>
+// ALL-SAME: ins(%[[RES4]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES6:.+]] = linalg.elementwise <reciprocal>
+// ALL-SAME: ins(%[[RES5]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES7:.+]] = linalg.elementwise <round>
+// ALL-SAME: ins(%[[RES6]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES8:.+]] = linalg.elementwise <sqrt>
+// ALL-SAME: ins(%[[RES7]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES9:.+]] = linalg.elementwise <rsqrt>
+// ALL-SAME: ins(%[[RES8]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES10:.+]] = linalg.elementwise <square>
+// ALL-SAME: ins(%[[RES9]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES11:.+]] = linalg.elementwise <tanh>
+// ALL-SAME: ins(%[[RES10]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES12:.+]] = linalg.elementwise <erf>
+// ALL-SAME: ins(%[[RES11]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES13:.+]] = linalg.elementwise <sin>
+// ALL-SAME: ins(%[[RES12]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES14:.+]] = linalg.elementwise <cos>
+// ALL-SAME: ins(%[[RES13]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES15:.+]] = linalg.elementwise <tan>
+// ALL-SAME: ins(%[[RES14]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES16:.+]] = linalg.elementwise <acos>
+// ALL-SAME: ins(%[[RES15]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES17:.+]] = linalg.elementwise <acosh>
+// ALL-SAME: ins(%[[RES16]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES18:.+]] = linalg.elementwise <asin>
+// ALL-SAME: ins(%[[RES17]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES19:.+]] = linalg.elementwise <asinh>
+// ALL-SAME: ins(%[[RES18]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES20:.+]] = linalg.elementwise <atan>
+// ALL-SAME: ins(%[[RES19]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES21:.+]] = linalg.elementwise <atanh>
+// ALL-SAME: ins(%[[RES20]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES22:.+]] = linalg.elementwise <log10>
+// ALL-SAME: ins(%[[RES21]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES23:.+]] = linalg.elementwise <log1p>
+// ALL-SAME: ins(%[[RES22]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// ALL: %[[RES24:.+]] = linalg.elementwise <log2>
+// ALL-SAME: ins(%[[RES23]] : tensor<?x?x?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// -----
@@ -249,14 +245,11 @@ func.func @unary_ops_non_identity(%A: tensor<?xf32>, %Out: tensor<?x?xf32>) -> t
// ALL: unary_ops_non_identity
// ALL-SAME: %[[A:.+]]: tensor<?xf32>, %[[OUT:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>
-// Named ops cannot carry user-defined indexing maps -> expect no change.
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.elementwise <exp>
-// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]]]
-// CATEGORY-SAME: ins(%[[A]] : tensor<?xf32>)
-// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+// ALL-NOT: linalg.generic
+// ALL: linalg.elementwise <exp>
+// ALL-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]]]
+// ALL-SAME: ins(%[[A]] : tensor<?xf32>)
+// ALL-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// -----
@@ -324,28 +317,25 @@ func.func @binary_ops_int(%A: tensor<?x?xi32>, %B: tensor<?x?xi32>,
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi32>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-// NAMED-NOT: linalg.elementwise
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: %[[RES1:.+]] = linalg.elementwise <sub>
-// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES2:.+]] = linalg.elementwise <mul>
-// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES3:.+]] = linalg.elementwise <div>
-// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES4:.+]] = linalg.elementwise <div_unsigned>
-// CATEGORY-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES5:.+]] = linalg.elementwise <max_signed>
-// CATEGORY-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES6:.+]] = linalg.elementwise <min_signed>
-// CATEGORY-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL-NOT: linalg.generic
+// ALL: %[[RES1:.+]] = linalg.elementwise <sub>
+// ALL-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES2:.+]] = linalg.elementwise <mul>
+// ALL-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES3:.+]] = linalg.elementwise <div>
+// ALL-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES4:.+]] = linalg.elementwise <div_unsigned>
+// ALL-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES5:.+]] = linalg.elementwise <max_signed>
+// ALL-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES6:.+]] = linalg.elementwise <min_signed>
+// ALL-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// -----
@@ -413,28 +403,25 @@ func.func @binary_ops_float(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-// NAMED-NOT: linalg.elementwise
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: %[[RES1:.+]] = linalg.elementwise <sub>
-// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES2:.+]] = linalg.elementwise <mul>
-// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES3:.+]] = linalg.elementwise <div>
-// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES4:.+]] = linalg.elementwise <max_signed>
-// CATEGORY-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES5:.+]] = linalg.elementwise <min_signed>
-// CATEGORY-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES6:.+]] = linalg.elementwise <powf>
-// CATEGORY-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL-NOT: linalg.generic
+// ALL: %[[RES1:.+]] = linalg.elementwise <sub>
+// ALL-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES2:.+]] = linalg.elementwise <mul>
+// ALL-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES3:.+]] = linalg.elementwise <div>
+// ALL-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES4:.+]] = linalg.elementwise <max_signed>
+// ALL-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES5:.+]] = linalg.elementwise <min_signed>
+// ALL-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES6:.+]] = linalg.elementwise <powf>
+// ALL-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// -----
@@ -477,19 +464,16 @@ func.func @binary_ops_complex(%A: tensor<?x?xcomplex<f32>>,
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xcomplex<f32>>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-// NAMED-NOT: linalg.elementwise
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: %[[RES1:.+]] = linalg.elementwise <sub>
-// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES2:.+]] = linalg.elementwise <mul>
-// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES3:.+]] = linalg.elementwise <div>
-// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL-NOT: linalg.generic
+// ALL: %[[RES1:.+]] = linalg.elementwise <sub>
+// ALL-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES2:.+]] = linalg.elementwise <mul>
+// ALL-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES3:.+]] = linalg.elementwise <div>
+// ALL-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// -----
@@ -512,13 +496,10 @@ func.func @binary_ops_bool(%A: tensor<?x?xi1>, %B: tensor<?x?xi1>,
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi1>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-// NAMED-NOT: linalg.elementwise
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: %[[RES1:.+]] = linalg.elementwise <mul>
-// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL-NOT: linalg.generic
+// ALL: %[[RES1:.+]] = linalg.elementwise <mul>
+// ALL-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// -----
@@ -574,15 +555,11 @@ func.func @binary_ops_non_identity(%A: tensor<?xf32>, %B: tensor<?x?xf32>,
// ALL-SAME: %[[A:.+]]: [[TTY1D:tensor<\?xf32>]], %[[B:.+]]: [[TTY:tensor<\?x\?xf32>]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-// Named ops cannot carry user-defined indexing maps -> expect no change.
-// NAMED-NOT: linalg.sub
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.elementwise <sub>
-// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]], #[[MAP_ID]]]
-// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY1D]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL-NOT: linalg.generic
+// ALL: linalg.elementwise <sub>
+// ALL-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]], #[[MAP_ID]]]
+// ALL-SAME: ins(%[[A]], %[[B]] : [[TTY1D]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// -----
@@ -619,17 +596,14 @@ func.func @binary_ops_swapped(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
// ALL-SAME: %[[C:.+]]: [[TTY1D:tensor<\?xf32>]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-// NAMED-NOT: linalg.elementwise
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: %[[RES0:.+]] = linalg.elementwise <mul>
-// CATEGORY-SAME: ins(%[[B]], %[[A]] : [[TTY]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
-// CATEGORY: %[[RES1:.+]] = linalg.elementwise <sub>
-// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
-// CATEGORY-SAME: ins(%[[C]], %[[RES0]] : [[TTY1D]], [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL-NOT: linalg.generic
+// ALL: %[[RES0:.+]] = linalg.elementwise <mul>
+// ALL-SAME: ins(%[[B]], %[[A]] : [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL: %[[RES1:.+]] = linalg.elementwise <sub>
+// ALL-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
+// ALL-SAME: ins(%[[C]], %[[RES0]] : [[TTY1D]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// -----
@@ -649,22 +623,18 @@ func.func @unary_op_with_scalar(%A: tensor<?xi32>, %Out: tensor<?xi32>)
return %0 : tensor<?xi32>
}
-// CATEGORY-DAG: #[[MAP_ID:.+]] = affine_map<(d0) -> (d0)>
-// CATEGORY-DAG: #[[MAP_BC:.+]] = affine_map<(d0) -> ()>
+// ALL-DAG: #[[MAP_ID:.+]] = affine_map<(d0) -> (d0)>
+// ALL-DAG: #[[MAP_BC:.+]] = affine_map<(d0) -> ()>
// ALL: unary_op_with_scalar
-// CATEGORY-SAME: %[[A:.+]]: [[TTY:tensor<\?xi32>]],
-// CATEGORY-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-
-// Named ops cannot broadcast from a scalar operand -> expect no change.
-// NAMED-NOT: linalg.sub
-// NAMED: linalg.generic
+// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?xi32>]],
+// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: %[[CST:.+]] = arith.constant 123 : i32
-// CATEGORY: linalg.elementwise <sub>
-// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
-// CATEGORY-SAME: ins(%[[CST]], %[[A]] : i32, [[TTY]])
-// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+// ALL-NOT: linalg.generic
+// ALL: %[[CST:.+]] = arith.constant 123 : i32
+// ALL: linalg.elementwise <sub>
+// ALL-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
+// ALL-SAME: ins(%[[CST]], %[[A]] : i32, [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// -----
@@ -686,11 +656,8 @@ func.func @negative_unary_op_using_block_arg_twice(%A: tensor<?xi32>,
// ALL-LABEL: negative_unary_op_using_block_arg_twice
-// Named ops cannot broadcast from a scalar operand -> expect no change.
-// NAMED-NOT: linalg.add
-
// There is no scalar operand to hoist -> expect no change.
-// CATEGORY-NOT: linalg.elementwise <add>
+// ALL-NOT: linalg.elementwise <add>
// ALL: linalg.generic
>From 43d923d9940838b6610f0deb604608ada0683d80 Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Thu, 10 Sep 2026 09:37:33 -0700
Subject: [PATCH 02/12] specialize contractions now stop at contract when there
is no named alternative
---
.../Dialect/Linalg/Transforms/Specialize.cpp | 88 +++++++++++--------
.../Linalg/specialize-generic-ops.mlir | 48 +++-------
2 files changed, 67 insertions(+), 69 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index a3a37bc5be242..ab16952f31b35 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -444,41 +444,13 @@ static bool isSupportedContractionPair(Operation *first, Operation *second) {
return false;
}
-// Converts linalg.generic to named linalg.*matmul* where possible.
-static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
- GenericOp genericOp,
- bool emitCategoryOp) {
- if (genericOp.getNumDpsInputs() != 2 || genericOp.getNumDpsInits() != 1)
- return failure();
-
- // Early exit if not projected permutations.
- auto mapRange = genericOp.getIndexingMapsArray();
- if (llvm::any_of(mapRange,
- [](AffineMap m) { return !m.isProjectedPermutation(); }))
- return failure();
-
- // Only contractions that can be represented by named linalg ops are
- // eligible for specialization:
- // - mul + add (floating-point, integer, complex)
- // - and + or (bool)
- if (!mlir::linalg::detail::isContractionBody(*genericOp.getBlock(),
- isSupportedContractionPair))
- return failure();
-
- // Determine the cast type for the named matmul op, or bail out if casts
- // cannot be represented by the named op.
- std::optional<TypeFn> castTy = getCastTypeForMatmulLikeOp(genericOp);
- if (!castTy)
- return rewriter.notifyMatchFailure(
- genericOp, "contains invalid cast ops for the named matmul op");
-
- // In case of category op, wider range of variants is supported.
- if (emitCategoryOp)
- return replaceWithMatmulVariant<ContractOp>(
- rewriter, genericOp, castTy, genericOp.getIndexingMapsArray());
-
- // Further checks for named variants.
- //
+// Attempts to specialize `genericOp` to a specific named matmul variant
+// (`matmul`, `batch_matmul`, or `mmt4d`). Returns failure without modifying the
+// IR if no named variant matches. `castTy` is the cast type inferred for the
+// contraction body.
+static FailureOr<LinalgOp>
+specializeToNamedContraction(RewriterBase &rewriter, GenericOp genericOp,
+ std::optional<TypeFn> castTy) {
// Linalg generic contraction can be across multiple axis e.g.
// ```
// linalg.generic
@@ -579,6 +551,52 @@ static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
namedOpMaps);
}
+// Converts linalg.generic to named linalg.*matmul* where possible, falling
+// back to the generic `linalg.contract` op otherwise.
+static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
+ GenericOp genericOp,
+ bool emitCategoryOp) {
+ if (genericOp.getNumDpsInputs() != 2 || genericOp.getNumDpsInits() != 1)
+ return failure();
+
+ // Early exit if not projected permutations.
+ auto mapRange = genericOp.getIndexingMapsArray();
+ if (llvm::any_of(mapRange,
+ [](AffineMap m) { return !m.isProjectedPermutation(); }))
+ return failure();
+
+ // Only contractions that can be represented by named linalg ops are
+ // eligible for specialization:
+ // - mul + add (floating-point, integer, complex)
+ // - and + or (bool)
+ if (!mlir::linalg::detail::isContractionBody(*genericOp.getBlock(),
+ isSupportedContractionPair))
+ return failure();
+
+ // Determine the cast type for the named matmul op, or bail out if casts
+ // cannot be represented by the named op.
+ std::optional<TypeFn> castTy = getCastTypeForMatmulLikeOp(genericOp);
+ if (!castTy)
+ return rewriter.notifyMatchFailure(
+ genericOp, "contains invalid cast ops for the named matmul op");
+
+ // TODO: When `emitCategoryOp` is set, skip the named-variant matching below
+ // and go straight to the `linalg.contract` category op.
+
+ // Try to specialize to a specific named matmul variant first.
+ if (!emitCategoryOp) {
+ FailureOr<LinalgOp> namedOp =
+ specializeToNamedContraction(rewriter, genericOp, castTy);
+ if (succeeded(namedOp))
+ return namedOp;
+ }
+
+ // No named variant matched; fall back to the generic `linalg.contract` op,
+ // which supports a wider range of variants.
+ return replaceWithMatmulVariant<ContractOp>(
+ rewriter, genericOp, castTy, genericOp.getIndexingMapsArray());
+}
+
/// Utility to specialize a `genericOp` with a convolution op of type `ConvOpTy`
/// with `dilations` and `strides`.
template <typename ConvOpTy>
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index 1d0fda92355d0..1bc045b4615e0 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -1029,11 +1029,8 @@ func.func @op_multi_reduction(%A: tensor<10x20x30xf32>,
// ALL-LABEL: op_multi_reduction
-// Cannot be lifted to named matrix multiply.
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.contract
+// ALL-NOT: linalg.generic
+// ALL: linalg.contract
// -----
@@ -1059,11 +1056,8 @@ func.func @batch_matmul_non_identity_batch(%A: tensor<4x2x8xf32>, %B: tensor<2x8
// ALL-LABEL: batch_matmul_non_identity_batch
-// Cannot be lifted to named matrix multiply.
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.contract
+// ALL-NOT: linalg.generic
+// ALL: linalg.contract
// -----
@@ -1087,10 +1081,8 @@ func.func @op_matvec(%A: tensor<?x?xf32>, %B: tensor<?xf32>, %Out: tensor<?xf32>
// ALL-LABEL: op_matvec
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.contract
+// ALL-NOT: linalg.generic
+// ALL: linalg.contract
// -----
@@ -1556,11 +1548,8 @@ func.func @op_matmul_broadcast_a(%A: tensor<?xf32>, %B: tensor<?x?xf32>,
// ALL-LABEL: op_matmul_broadcast_a
-// NAMED: linalg.generic
-// NAMED-NOT: linalg.matmul
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.contract
+// ALL-NOT: linalg.generic
+// ALL: linalg.contract
// -----
@@ -1585,11 +1574,8 @@ func.func @op_batch_matmul_broadcast_a(%A: tensor<16x8xf32>, %B: tensor<2x8x16xf
// ALL-LABEL: op_batch_matmul_broadcast_a
-// NAMED: linalg.generic
-// NAMED-NOT: linalg.batch_matmul
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.contract
+// ALL-NOT: linalg.generic
+// ALL: linalg.contract
// -----
@@ -1614,11 +1600,8 @@ func.func @op_batch_matmul_broadcast_b(%A: tensor<2x16x8xf32>, %B: tensor<8xf32>
// ALL-LABEL: op_batch_matmul_broadcast_b
-// NAMED: linalg.generic
-// NAMED-NOT: linalg.batch_matmul
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.contract
+// ALL-NOT: linalg.generic
+// ALL: linalg.contract
// -----
@@ -1808,8 +1791,5 @@ func.func @negative_op_mmt4d(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
// ALL-LABEL: negative_op_mmt4d
-// NAMED-NOT: linalg.mmt4d
-// NAMED: linalg.generic
-
-// CATEGORY-NOT: linalg.generic
-// CATEGORY: linalg.contract
+// ALL-NOT: linalg.generic
+// ALL: linalg.contract
>From a8db5a62ae12194b7b0112572d8bc9269110d709 Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 03:28:43 -0700
Subject: [PATCH 03/12] remove morphism options struct
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 30 +++++++----
.../Dialect/Linalg/Transforms/Transforms.h | 52 ++++++++++++-------
.../TransformOps/LinalgTransformOps.cpp | 4 +-
.../Linalg/Transforms/Generalization.cpp | 3 +-
.../Dialect/Linalg/Transforms/MorphOps.cpp | 22 ++++++--
.../Dialect/Linalg/Transforms/Specialize.cpp | 11 ++--
6 files changed, 81 insertions(+), 41 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 1a9268d514e01..c6a49f37e2879 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -29,9 +29,17 @@ include "mlir/IR/Constraints.td"
// \-> sub
// \-> ...
//
-// Morphisms between representations can happen in the following 6 ways:
-// generic <---> category <---> named
-// \-------------------------/
+// Morphisms between representations can happen in the following ways:
+// * Specialize: Convert `generic` operations to either:
+// * `named`, where the interfaces match and there is such op, OR
+// * `category`, where enough interfaces still match, OR
+// * `category`, when the `emitCategory` flag is set (regardless of named ops)
+// * Generalize: Convert `named` or `category` operations to `generic` operations
+// * Any structured operation can be converted to `generic`
+// * Categorize: Convert both `named` and `generic` operations to `category` operations:
+// * `elementwise`, for non-contracting `generic` operations that match the interfaces, OR
+// * `contract`, for contracting `generic` operations that match the interfaces
+// * `contract`, for contracting `named` operations that match the interfaces
//
// generic subsumes category which subsumes structured named (not softmax,
// convolutions, etc). The generalization path is guaranteed, the
@@ -77,14 +85,18 @@ def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
"convert category ops e.g. `linalg.elementwise` to equivalent named ops"> ];
}
-def LinalgGeneralizeNamedOpsPass : Pass<"linalg-generalize-named-ops">,
- Deprecated<"Use 'linalg-morph-ops' instead."> {
- let summary = "Convert named ops into generic ops";
+def LinalgGeneralizeNamedOpsPass : Pass<"linalg-generalize-named-ops"> {
+ let summary = "Convert `named` and `category` ops into `generic` ops";
+ let dependentDialects = ["linalg::LinalgDialect"];
+}
+
+def LinalgSpecializeGenericOpsPass : Pass<"linalg-specialize-generic-ops"> {
+ let summary = "Convert `generic` ops into `named` or `category` ops";
+ let dependentDialects = ["linalg::LinalgDialect"];
}
-def LinalgSpecializeGenericOpsPass : Pass<"linalg-specialize-generic-ops">,
- Deprecated<"Use 'linalg-morph-ops' instead."> {
- let summary = "Convert generic ops back to named ops";
+def LinalgCategorizeOpsPass : Pass<"linalg-categorize-ops"> {
+ let summary = "Convert `named` and `generic` ops into `category` ops";
let dependentDialects = ["linalg::LinalgDialect"];
}
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index cbc338c9f1b32..77d20247d867a 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -917,21 +917,17 @@ FailureOr<GenericOp> interchangeGenericOp(RewriterBase &rewriter,
GenericOp genericOp,
ArrayRef<unsigned> interchangeVector);
-/// Create a GenericOp from the given named operation `linalgOp` and replace
-/// the given `linalgOp`.
-/// Return failure if `linalgOp` is a GenericOp or misses a region builder.
+/// Create a GenericOp or CategoryOp from the given named operation `linalgOp`
+/// and replace the given `linalgOp`. Return failure if `linalgOp` is a
+/// GenericOp or misses a region builder.
FailureOr<GenericOp> generalizeNamedOp(RewriterBase &rewriter,
- LinalgOp linalgOp);
-
-struct GenericOpSpecializationOptions {
- // Specialize generics to category ops (default: named ops).
- bool emitCategoryOps = false;
-};
+ LinalgOp linalgOp,
+ bool emitCategoryOps = false);
/// Replace the given GenericOp with a namedOp or categoryOp.
-FailureOr<LinalgOp>
-specializeGenericOp(RewriterBase &rewriter, GenericOp genericOp,
- const GenericOpSpecializationOptions &options = {});
+FailureOr<LinalgOp> specializeGenericOp(RewriterBase &rewriter,
+ GenericOp genericOp,
+ bool emitCategoryOps = false);
/// Create a new buffer using the `allocationFn` provided. The size of this
/// buffer is either the original subview size when 'useOriginalSubviewSize' is
@@ -1689,13 +1685,13 @@ struct LinalgGeneralizationPattern
struct LinalgSpecializationPattern : public OpRewritePattern<GenericOp> {
LinalgSpecializationPattern(
- MLIRContext *context, const GenericOpSpecializationOptions &options = {},
+ MLIRContext *context, bool emitCategoryOps = false,
PatternBenefit benefit = 1)
- : OpRewritePattern<GenericOp>(context, benefit), options(options) {}
+ : OpRewritePattern<GenericOp>(context, benefit), emitCategoryOps(emitCategoryOps) {}
FailureOr<GenericOp>
returningMatchAndRewrite(GenericOp op, PatternRewriter &rewriter) const {
- return specializeGenericOp(rewriter, op, options);
+ return specializeGenericOp(rewriter, op, emitCategoryOps);
}
LogicalResult matchAndRewrite(GenericOp op,
@@ -1704,7 +1700,26 @@ struct LinalgSpecializationPattern : public OpRewritePattern<GenericOp> {
}
private:
- GenericOpSpecializationOptions options;
+ bool emitCategoryOps;
+};
+
+struct LinalgCategorizationPattern : public OpInterfaceRewritePattern<LinalgOp> {
+
+ LinalgCategorizationPattern(
+ MLIRContext *context, PatternBenefit benefit = 1)
+ : OpInterfaceRewritePattern<LinalgOp>(context, benefit) {}
+
+ FailureOr<LinalgOp>
+ returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const {
+ if (GenericOp generic = dyn_cast<GenericOp>(*op))
+ return specializeGenericOp(rewriter, generic, true);
+ return generalizeNamedOp(rewriter, op);
+ }
+
+ LogicalResult matchAndRewrite(LinalgOp op,
+ PatternRewriter &rewriter) const override {
+ return returningMatchAndRewrite(op, rewriter);
+ }
};
/// Vectorization pattern for memref::CopyOp.
@@ -1919,9 +1934,8 @@ void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns);
/// e.g. linalg.generic that takes a tensor and computes a polynomial such as:
/// p(x) = an*x^n + ... + a1x + a0
/// There is no equivalent named op to convert to. Many such cases exist.
-void populateLinalgGenericOpsSpecializationPatterns(
- RewritePatternSet &patterns,
- const GenericOpSpecializationOptions &options = {});
+void populateLinalgGenericOpsSpecializationPatterns(RewritePatternSet &patterns,
+ bool emitCategoryOps = false);
/// Populates `patterns` with patterns that fold operations like
/// `linalg.transform` into elementwise op map.
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 74cba3c6ac484..9948d7b12d859 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -1449,10 +1449,8 @@ transform::SpecializeOp::applyToOne(transform::TransformRewriter &rewriter,
return DiagnosedSilenceableFailure::success();
}
rewriter.setInsertionPoint(target);
- GenericOpSpecializationOptions opts;
- opts.emitCategoryOps = getEmitCategory();
FailureOr<LinalgOp> named =
- specializeGenericOp(rewriter, cast<GenericOp>(target), opts);
+ specializeGenericOp(rewriter, cast<GenericOp>(target), getEmitCategory());
if (succeeded(named)) {
results.push_back(named->getOperation());
return DiagnosedSilenceableFailure::success();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index 35abdc027f194..bcaa0992e7895 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -44,7 +44,8 @@ static LogicalResult generalizeNamedOpPrecondition(LinalgOp linalgOp) {
}
FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
- LinalgOp linalgOp) {
+ LinalgOp linalgOp,
+ bool emitCategoryOps) {
if (failed(generalizeNamedOpPrecondition(linalgOp)))
return rewriter.notifyMatchFailure(linalgOp, "preconditions not met");
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
index 02cadf7567de3..eb4b66f2c7315 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -21,6 +21,7 @@
namespace mlir {
#define GEN_PASS_DEF_LINALGMORPHOPSPASS
+#define GEN_PASS_DEF_LINALGCATEGORIZEOPSPASS
#include "mlir/Dialect/Linalg/Passes.h.inc"
} // namespace mlir
@@ -49,12 +50,27 @@ void LinalgMorphOpsPass::runOnOperation() {
// Lifting paths (named <- category <- generic)
if (genericToNamed || genericToCategory) {
- GenericOpSpecializationOptions opts;
- opts.emitCategoryOps = genericToCategory;
- populateLinalgGenericOpsSpecializationPatterns(patterns, opts);
+ populateLinalgGenericOpsSpecializationPatterns(patterns, genericToCategory);
}
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
signalPassFailure();
}
+
+struct LinalgCategorizeOpsPass
+ : public impl::LinalgCategorizeOpsPassBase<LinalgCategorizeOpsPass> {
+
+ using impl::LinalgCategorizeOpsPassBase<
+ LinalgCategorizeOpsPass>::LinalgCategorizeOpsPassBase;
+
+ void runOnOperation() override;
+};
+
+void LinalgCategorizeOpsPass::runOnOperation() {
+ RewritePatternSet patterns(&getContext());
+ patterns.add<LinalgCategorizationPattern>(&getContext());
+
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ signalPassFailure();
+}
} // namespace
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index ab16952f31b35..a30e5face72f0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -699,7 +699,7 @@ static FailureOr<LinalgOp> specializeLinalgConvolutions(RewriterBase &rewriter,
//===----------------------------------------------------------------------===//
FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(
RewriterBase &rewriter, GenericOp genericOp,
- const GenericOpSpecializationOptions &options) {
+ bool emitCategoryOps) {
// Elementwise - e.g. exp, add
if (isaElemwiseSingleUnaryOpInterface(genericOp) ||
isaElemwiseSingleBinaryOpInterface(genericOp)) {
@@ -709,13 +709,13 @@ FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(
// Contraction - e.g. matmul
if (isaContractionOpInterface(genericOp)) {
return specializeLinalgContractions(rewriter, genericOp,
- options.emitCategoryOps);
+ emitCategoryOps);
}
// Early exit in case of category specialization.
// TODO: Remove when matches for other ops account for both named and
// category.
- if (options.emitCategoryOps)
+ if (emitCategoryOps)
return rewriter.notifyMatchFailure(
genericOp, "no matching category op specialization");
@@ -785,7 +785,6 @@ void LinalgSpecializeGenericOpsPass::runOnOperation() {
}
void mlir::linalg::populateLinalgGenericOpsSpecializationPatterns(
- RewritePatternSet &patterns,
- const GenericOpSpecializationOptions &options) {
- patterns.add<LinalgSpecializationPattern>(patterns.getContext(), options);
+ RewritePatternSet &patterns, bool emitCategoryOps) {
+ patterns.add<LinalgSpecializationPattern>(patterns.getContext(), emitCategoryOps);
}
>From 3f508c1f721652f4e219290ece9150b132c876ff Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 06:15:18 -0700
Subject: [PATCH 04/12] add missing support for ternary specialization
---
.../mlir/Dialect/Linalg/IR/LinalgInterfaces.h | 6 +
.../Dialect/Linalg/IR/LinalgInterfaces.cpp | 15 +++
.../Dialect/Linalg/Transforms/Specialize.cpp | 6 +-
.../Linalg/specialize-generic-ops.mlir | 109 ++++++++++++++++++
4 files changed, 134 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
index 89b42f1beef7e..ecb2397a21689 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
@@ -162,6 +162,12 @@ bool isaElemwiseSingleUnaryOpInterface(GenericOp genericOp);
/// they can be represented by the category op.
bool isaElemwiseSingleBinaryOpInterface(GenericOp genericOp);
+/// Checks whether `genericOp` is semantically equivalent to a single linalg
+/// elementwise ternary op e.g. linalg.fma.
+/// Operations with non-identity indexing maps are included in the check, as
+/// they can be represented by the category op.
+bool isaElemwiseSingleTernaryOpInterface(GenericOp genericOp);
+
/// Checks whether `genericOp` is semantically equivalent to a `linalg.fill`.
/// Supports two patterns:
/// 1. External: linalg.generic ins(%scalar) outs(%tensor) { yield %scalar }
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index c8088ce183983..e1152b287330b 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -287,6 +287,21 @@ bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op) {
!op.payloadUsesValueFromOperand(inputOpOperand1));
}
+bool linalg::isaElemwiseSingleTernaryOpInterface(linalg::GenericOp op) {
+ // All basic elemwise checks.
+ if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 3))
+ return false;
+
+ // Check all three inputs are used (elementwise).
+ OpOperand *inputOpOperand0 = op.getDpsInputOperand(0);
+ OpOperand *inputOpOperand1 = op.getDpsInputOperand(1);
+ OpOperand *inputOpOperand2 = op.getDpsInputOperand(2);
+ return !(!op.payloadUsesValueFromOperand(inputOpOperand0) ||
+ !op.payloadUsesValueFromOperand(inputOpOperand1) ||
+ !op.payloadUsesValueFromOperand(inputOpOperand2));
+}
+
+
//===----------------------------------------------------------------------===//
// ContractionOpInterface implementation
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index a30e5face72f0..df11be9b38c75 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -136,6 +136,7 @@ static std::optional<ElementwiseKind> getElementwiseKind(Operation *op) {
.Case([](math::PowFOp) { return ElementwiseKind::powf; })
.Case([](arith::MaxUIOp) { return ElementwiseKind::max_unsigned; })
.Case([](arith::MinUIOp) { return ElementwiseKind::min_unsigned; })
+ .Case([](arith::SelectOp) { return ElementwiseKind::select; })
.Default([](Operation *) { return std::nullopt; });
}
@@ -700,9 +701,10 @@ static FailureOr<LinalgOp> specializeLinalgConvolutions(RewriterBase &rewriter,
FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(
RewriterBase &rewriter, GenericOp genericOp,
bool emitCategoryOps) {
- // Elementwise - e.g. exp, add
+ // Elementwise - e.g. exp, add are always category ops
if (isaElemwiseSingleUnaryOpInterface(genericOp) ||
- isaElemwiseSingleBinaryOpInterface(genericOp)) {
+ isaElemwiseSingleBinaryOpInterface(genericOp) ||
+ isaElemwiseSingleTernaryOpInterface(genericOp)) {
return specializeLinalgElementwise(rewriter, genericOp);
}
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index 1bc045b4615e0..6441f76eb849c 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -663,6 +663,115 @@ func.func @negative_unary_op_using_block_arg_twice(%A: tensor<?xi32>,
// -----
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @ternary_op_select_f32(%C: tensor<?x?xi1>, %T: tensor<?x?xf32>,
+ %F: tensor<?x?xf32>, %Out: tensor<?x?xf32>)
+ -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%C, %T, %F : tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: i1, %in_0: f32, %in_1: f32, %out: f32):
+ %v = arith.select %in, %in_0, %in_1 : f32
+ linalg.yield %v : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// ALL-LABEL: ternary_op_select_f32
+// ALL-SAME: %[[C:.+]]: [[ITY:tensor<\?x\?xi1>]], %[[T:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[F:.+]]: [[TTY]],
+// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
+
+// ALL-NOT: linalg.generic
+// ALL: linalg.elementwise <select>
+// ALL-SAME: ins(%[[C]], %[[T]], %[[F]] : [[ITY]], [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @ternary_op_select_i32(%C: tensor<?x?xi1>, %T: tensor<?x?xi32>,
+ %F: tensor<?x?xi32>, %Out: tensor<?x?xi32>)
+ -> tensor<?x?xi32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%C, %T, %F : tensor<?x?xi1>, tensor<?x?xi32>, tensor<?x?xi32>)
+ outs(%Out : tensor<?x?xi32>) {
+ ^bb0(%in: i1, %in_0: i32, %in_1: i32, %out: i32):
+ %v = arith.select %in, %in_0, %in_1 : i32
+ linalg.yield %v : i32
+ } -> tensor<?x?xi32>
+ return %0 : tensor<?x?xi32>
+}
+
+// ALL-LABEL: ternary_op_select_i32
+// ALL-SAME: %[[C:.+]]: [[ITY:tensor<\?x\?xi1>]], %[[T:.+]]: [[TTY:tensor<\?x\?xi32>]], %[[F:.+]]: [[TTY]],
+// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
+
+// ALL-NOT: linalg.generic
+// ALL: linalg.elementwise <select>
+// ALL-SAME: ins(%[[C]], %[[T]], %[[F]] : [[ITY]], [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @ternary_op_select_i1(%C: tensor<?x?xi1>, %T: tensor<?x?xi1>,
+ %F: tensor<?x?xi1>, %Out: tensor<?x?xi1>)
+ -> tensor<?x?xi1> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%C, %T, %F : tensor<?x?xi1>, tensor<?x?xi1>, tensor<?x?xi1>)
+ outs(%Out : tensor<?x?xi1>) {
+ ^bb0(%in: i1, %in_0: i1, %in_1: i1, %out: i1):
+ %v = arith.select %in, %in_0, %in_1 : i1
+ linalg.yield %v : i1
+ } -> tensor<?x?xi1>
+ return %0 : tensor<?x?xi1>
+}
+
+// ALL-LABEL: ternary_op_select_i1
+// ALL-SAME: %[[C:.+]]: [[TTY:tensor<\?x\?xi1>]], %[[T:.+]]: [[TTY]], %[[F:.+]]: [[TTY]],
+// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
+
+// ALL-NOT: linalg.generic
+// ALL: linalg.elementwise <select>
+// ALL-SAME: ins(%[[C]], %[[T]], %[[F]] : [[TTY]], [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+
+// -----
+
+// Mask comes from outside and is constant
+// this can be elided completely by canonicalization
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @negative_ternary_op_select(%C: tensor<?x?xi1>, %T: tensor<?x?xf32>,
+ %F: tensor<?x?xf32>, %Out: tensor<?x?xf32>)
+ -> tensor<?x?xf32> {
+ %true = arith.constant 1 : i1
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%C, %T, %F : tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: i1, %in_0: f32, %in_1: f32, %out: f32):
+ %v = arith.select %true, %in_0, %in_1 : f32
+ linalg.yield %v : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// ALL-LABEL: negative_ternary_op_select
+// ALL-SAME: %[[C:.+]]: [[ITY:tensor<\?x\?xi1>]], %[[T:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[F:.+]]: [[TTY]],
+// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
+
+// ALL: linalg.generic
+// ALL-NOT: linalg.elementwise <select>
+
+// -----
+
///----------------------------------------------------------------------------------------
/// Tests for linalg.matmul
///----------------------------------------------------------------------------------------
>From ff84a9baef35329a23cbed360c432138dadb0d48 Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 06:46:27 -0700
Subject: [PATCH 05/12] add categorize pass and tests
---
.../Dialect/Linalg/Transforms/Transforms.h | 13 +-
.../Linalg/Transforms/Generalization.cpp | 47 ++-
.../Dialect/Linalg/Transforms/Transforms.cpp | 8 +-
.../Dialect/Linalg/categorize-named-ops.mlir | 303 ++++++++++++++++++
4 files changed, 358 insertions(+), 13 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/categorize-named-ops.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 77d20247d867a..cd2ad82c9a0e6 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -920,9 +920,9 @@ FailureOr<GenericOp> interchangeGenericOp(RewriterBase &rewriter,
/// Create a GenericOp or CategoryOp from the given named operation `linalgOp`
/// and replace the given `linalgOp`. Return failure if `linalgOp` is a
/// GenericOp or misses a region builder.
-FailureOr<GenericOp> generalizeNamedOp(RewriterBase &rewriter,
- LinalgOp linalgOp,
- bool emitCategoryOps = false);
+FailureOr<LinalgOp> generalizeNamedOp(RewriterBase &rewriter,
+ LinalgOp linalgOp,
+ bool emitCategoryOps = false);
/// Replace the given GenericOp with a namedOp or categoryOp.
FailureOr<LinalgOp> specializeGenericOp(RewriterBase &rewriter,
@@ -1671,7 +1671,7 @@ struct LinalgGeneralizationPattern
/// `matchAndRewrite` implementation that returns the significant
/// transformed pieces of IR.
- FailureOr<GenericOp>
+ FailureOr<LinalgOp>
returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const {
return generalizeNamedOp(rewriter, op);
}
@@ -1711,9 +1711,10 @@ struct LinalgCategorizationPattern : public OpInterfaceRewritePattern<LinalgOp>
FailureOr<LinalgOp>
returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const {
+ bool emitCategoryOps = true;
if (GenericOp generic = dyn_cast<GenericOp>(*op))
- return specializeGenericOp(rewriter, generic, true);
- return generalizeNamedOp(rewriter, op);
+ return specializeGenericOp(rewriter, generic, emitCategoryOps);
+ return generalizeNamedOp(rewriter, op, emitCategoryOps);
}
LogicalResult matchAndRewrite(LinalgOp op,
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index bcaa0992e7895..d6b0e05928b78 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -43,12 +43,51 @@ static LogicalResult generalizeNamedOpPrecondition(LinalgOp linalgOp) {
return success();
}
-FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
- LinalgOp linalgOp,
- bool emitCategoryOps) {
+// Converts a named matmul-like op (`matmul`, `batch_matmul`, or
+// `batch_reduce_matmul`) into a `linalg.contract` category op, preserving the
+// operand indexing maps and cast semantics. Returns failure for other ops.
+static FailureOr<LinalgOp> generalizeToContractOp(RewriterBase &rewriter,
+ LinalgOp namedOp) {
+ if (!isa<MatmulOp, BatchMatmulOp, BatchReduceMatmulOp>(
+ namedOp.getOperation()))
+ return failure();
+
+ SmallVector<NamedAttribute> attributes;
+
+ // Preserve operand indexing semantics (transposition, batch/reduction dims)
+ // via the named op's indexing maps.
+ SmallVector<Attribute> indexingMaps = llvm::map_to_vector(
+ namedOp.getIndexingMapsArray(),
+ [](AffineMap map) -> Attribute { return AffineMapAttr::get(map); });
+ attributes.push_back(rewriter.getNamedAttr(
+ "indexing_maps", rewriter.getArrayAttr(indexingMaps)));
+
+ // Only the unsigned cast needs to be explicit; signed is the default.
+ if (auto castAttr = namedOp->getAttrOfType<TypeFnAttr>("cast");
+ castAttr && castAttr.getValue() == TypeFn::cast_unsigned)
+ attributes.push_back(rewriter.getNamedAttr("cast", castAttr));
+
+ LinalgOp contractOp = rewriter.replaceOpWithNewOp<ContractOp>(
+ namedOp,
+ ValueRange{namedOp.getDpsInputs()[0], namedOp.getDpsInputs()[1]},
+ ValueRange{namedOp.getDpsInits()[0]}, attributes);
+ return contractOp;
+}
+
+FailureOr<LinalgOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
+ LinalgOp linalgOp,
+ bool emitCategoryOps) {
if (failed(generalizeNamedOpPrecondition(linalgOp)))
return rewriter.notifyMatchFailure(linalgOp, "preconditions not met");
+ // Emit the `linalg.contract` category op for matmul-like named ops.
+ if (emitCategoryOps) {
+ FailureOr<LinalgOp> contractOp = generalizeToContractOp(rewriter, linalgOp);
+ if (succeeded(contractOp))
+ return contractOp;
+ return rewriter.notifyMatchFailure(linalgOp, "failed to categorize to named op");
+ }
+
SmallVector<Value> inputs = linalgOp.getDpsInputs();
ValueRange outputs = linalgOp.getDpsInits();
SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
@@ -74,7 +113,7 @@ FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
genericOp->setDiscardableAttrs(linalgOp->getDiscardableAttrDictionary());
rewriter.replaceOp(linalgOp, genericOp->getResults());
- return genericOp;
+ return cast<LinalgOp>(genericOp.getOperation());
}
namespace {
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index aa021cbad24a7..a2d16c0a5b7d4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -825,10 +825,12 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
// 2.a. Rewrite as a generic.
auto genericOp = dyn_cast<GenericOp>(linalgOp.getOperation());
if (!genericOp) {
- FailureOr<GenericOp> generalizeResult =
+ FailureOr<LinalgOp> generalizeResult =
generalizeNamedOp(rewriter, linalgOp);
- assert(succeeded(generalizeResult) && "unexpected failure generalizing op");
- genericOp = *generalizeResult;
+ assert(succeeded(generalizeResult) &&
+ isa<GenericOp>(generalizeResult->getOperation()) &&
+ "unexpected failure generalizing op");
+ genericOp = cast<GenericOp>(generalizeResult->getOperation());
}
// 2.b. Interchange to move the dimensions (k, m, n) as most-minor
diff --git a/mlir/test/Dialect/Linalg/categorize-named-ops.mlir b/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
new file mode 100644
index 0000000000000..42d22f84f0c45
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
@@ -0,0 +1,303 @@
+// RUN: mlir-opt %s -split-input-file -linalg-categorize-ops | FileCheck %s
+
+func.func @categorize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) {
+ linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>)
+ outs(%C: memref<16x32xf32>)
+ return
+}
+
+// CHECK-DAG: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @categorize_matmul_buffer
+// CHECK-SAME: %[[A:.+]]: memref<16x8xf32>
+// CHECK-SAME: %[[B:.+]]: memref<8x32xf32>
+// CHECK-SAME: %[[C:.+]]: memref<16x32xf32>
+
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%[[A]], %[[B]] : memref<16x8xf32>, memref<8x32xf32>)
+// CHECK-SAME: outs(%[[C]] : memref<16x32xf32>)
+
+// -----
+
+func.func @categorize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
+ %0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)
+ outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+ return %0: tensor<16x32xf32>
+}
+
+// CHECK-DAG: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @categorize_matmul_tensor
+
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)
+
+// -----
+
+// Unsigned cast is preserved on the category op.
+func.func @categorize_matmul_unsigned_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
+ %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
+ %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi16>)
+ outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
+ return %0: tensor<16x32xi32>
+}
+
+// CHECK-LABEL: func @categorize_matmul_unsigned_cast
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: {cast = #linalg.type_fn<cast_unsigned>}
+
+// -----
+
+func.func @categorize_batch_matmul(%A: tensor<2x3x5xf32>, %B: tensor<2x5x7xf32>, %C: tensor<2x3x7xf32>) -> tensor<2x3x7xf32> {
+ %0 = linalg.batch_matmul ins(%A, %B: tensor<2x3x5xf32>, tensor<2x5x7xf32>)
+ outs(%C: tensor<2x3x7xf32>) -> tensor<2x3x7xf32>
+ return %0 : tensor<2x3x7xf32>
+}
+
+// CHECK-DAG: #[[A_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+// CHECK: func @categorize_batch_matmul
+
+// CHECK-NOT: linalg.batch_matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x3x5xf32>, tensor<2x5x7xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<2x3x7xf32>)
+
+// -----
+
+func.func @categorize_batch_reduce_matmul(%A: memref<7x8x9xf32>, %B: memref<7x9x8xf32>, %C: memref<8x8xf32>) {
+ linalg.batch_reduce_matmul ins(%A, %B: memref<7x8x9xf32>, memref<7x9x8xf32>)
+ outs(%C: memref<8x8xf32>)
+ return
+}
+
+// CHECK-DAG: #[[A_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+// CHECK: func @categorize_batch_reduce_matmul
+
+// CHECK-NOT: linalg.batch_reduce_matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<7x8x9xf32>, memref<7x9x8xf32>)
+// CHECK-SAME: outs(%{{.+}} : memref<8x8xf32>)
+
+// -----
+
+///----------------------------------------------------------------------------------------
+/// Generic contraction ops (from specialize-generic-ops.mlir) also categorize
+/// to linalg.contract.
+///----------------------------------------------------------------------------------------
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @op_matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
+ %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func @op_matmul
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<?x?xf32>)
+
+// -----
+
+// Matmul transpose A: A is accessed as (k, m) instead of (m, k)
+#map_ta = affine_map<(d0, d1, d2) -> (d2, d0)>
+#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @op_matmul_transpose_a(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
+ %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map_ta, #map_b, #map_c],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
+// CHECK-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func @op_matmul_transpose_a
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<?x?xf32>)
+
+// -----
+
+// Matmul transpose B: B is accessed as (n, k) instead of (k, n)
+#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map_tb = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @op_matmul_transpose_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
+ %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map_a, #map_tb, #map_c],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func @op_matmul_transpose_b
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<?x?xf32>)
+
+// -----
+
+#mapbA = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#mapbB = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#mapbC = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+func.func @op_batch_matmul(%A: tensor<2x16x8xf32>, %B: tensor<2x8x16xf32>,
+ %Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapbA, #mapbB, #mapbC],
+ iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
+ outs(%Out : tensor<2x16x16xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<2x16x16xf32>
+ return %0 : tensor<2x16x16xf32>
+}
+
+// CHECK-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+// CHECK-LABEL: func @op_batch_matmul
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<2x16x16xf32>)
+
+// -----
+
+// A multi-reduction contraction (not a named-op matmul, still categorizes).
+#mapA = affine_map<(m, n, k1, k2) -> (m, k1, k2)>
+#mapB = affine_map<(m, n, k1, k2) -> (k2, k1, n)>
+#mapC = affine_map<(m, n, k1, k2) -> (m, n)>
+func.func @op_multi_reduction(%A: tensor<10x20x30xf32>,
+ %B: tensor<30x20x40xf32>,
+ %C: tensor<10x40xf32>) -> tensor<10x40xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapA, #mapB, #mapC],
+ iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
+ ins(%A, %B : tensor<10x20x30xf32>, tensor<30x20x40xf32>)
+ outs(%C : tensor<10x40xf32>) {
+ ^bb0(%a: f32, %b: f32, %c: f32):
+ %1 = arith.mulf %a, %b : f32
+ %2 = arith.addf %c, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<10x40xf32>
+ return %0 : tensor<10x40xf32>
+}
+
+// CHECK-LABEL: func @op_multi_reduction
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.contract
+
+// -----
+
+// TODO: named matvec.
+#mapmvA = affine_map<(d0, d1) -> (d0, d1)>
+#mapmvB = affine_map<(d0, d1) -> (d1)>
+#mapmvC = affine_map<(d0, d1) -> (d0)>
+func.func @op_matvec(%A: tensor<?x?xf32>, %B: tensor<?xf32>, %Out: tensor<?xf32>) -> tensor<?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapmvA, #mapmvB, #mapmvC],
+ iterator_types = ["parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?xf32>, tensor<?xf32>)
+ outs(%Out : tensor<?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?xf32>
+ return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @op_matvec
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.contract
+
+// -----
+
+#mapmA = affine_map<(m, n, k, m0, n0, k0) -> (m, k, m0, k0)>
+#mapmB = affine_map<(m, n, k, m0, n0, k0) -> (n, k, n0, k0)>
+#mapmC = affine_map<(m, n, k, m0, n0, k0) -> (m, n, m0, n0)>
+func.func @op_mmt4d(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
+ %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#mapmA, #mapmB, #mapmC],
+ iterator_types = ["parallel", "parallel", "reduction",
+ "parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
+ outs(%C : tensor<?x?x?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %1 = arith.mulf %in, %in_0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?x?x?xf32>
+ return %0 : tensor<?x?x?x?xf32>
+}
+
+// CHECK-LABEL: func @op_mmt4d
+// CHECK-NOT: linalg.generic
+// CHECK: linalg.contract
>From f1b26042514835c3a6d8a9b0082dee392fc60e05 Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 07:17:43 -0700
Subject: [PATCH 06/12] add categorization support to morph-ops
---
.../Dialect/Linalg/Transforms/Transforms.h | 14 +++-
.../Linalg/Transforms/Generalization.cpp | 4 +-
.../Dialect/Linalg/Transforms/MorphOps.cpp | 14 +++-
.../Dialect/Linalg/categorize-named-ops.mlir | 74 +++++++++++++++++++
.../Dialect/Linalg/generalize-named-ops.mlir | 2 +
.../Linalg/specialize-generic-ops.mlir | 8 +-
6 files changed, 102 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index cd2ad82c9a0e6..699ffa61a8863 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1667,19 +1667,26 @@ FailureOr<LinalgOp> downscaleSizeOneWindowedConvolution(RewriterBase &rewriter,
// returning FailureOr<GenericOp>.
struct LinalgGeneralizationPattern
: public OpInterfaceRewritePattern<LinalgOp> {
- using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
+
+ LinalgGeneralizationPattern(
+ MLIRContext *context, bool emitCategoryOps = false,
+ PatternBenefit benefit = 1)
+ : OpInterfaceRewritePattern<LinalgOp>(context, benefit), emitCategoryOps(emitCategoryOps) {}
/// `matchAndRewrite` implementation that returns the significant
/// transformed pieces of IR.
FailureOr<LinalgOp>
returningMatchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const {
- return generalizeNamedOp(rewriter, op);
+ return generalizeNamedOp(rewriter, op, emitCategoryOps);
}
LogicalResult matchAndRewrite(LinalgOp op,
PatternRewriter &rewriter) const override {
return returningMatchAndRewrite(op, rewriter);
}
+
+private:
+ bool emitCategoryOps;
};
struct LinalgSpecializationPattern : public OpRewritePattern<GenericOp> {
@@ -1927,7 +1934,8 @@ void populateLinalgTilingCanonicalizationPatterns(RewritePatternSet &patterns);
/// Populates `patterns` with patterns to convert spec-generated named ops to
/// linalg.generic ops.
-void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns);
+void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns,
+ bool emitCategoryOps = false);
/// Populates `patterns` with patterns to convert linalg.generic ops to named
/// or category ops where possible. A linalg.generic can represent wide range
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index d6b0e05928b78..dd3a89bc21000 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -135,6 +135,6 @@ void LinalgGeneralizeNamedOpsPass::runOnOperation() {
}
void mlir::linalg::populateLinalgNamedOpsGeneralizationPatterns(
- RewritePatternSet &patterns) {
- patterns.add<LinalgGeneralizationPattern>(patterns.getContext());
+ RewritePatternSet &patterns, bool emitCategoryOps) {
+ patterns.add<LinalgGeneralizationPattern>(patterns.getContext(), emitCategoryOps);
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
index eb4b66f2c7315..ac05fca4621b3 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -46,12 +46,18 @@ void LinalgMorphOpsPass::runOnOperation() {
// Lowering paths (named -> category -> generic)
if (namedToGeneric || categoryToGeneric)
- populateLinalgNamedOpsGeneralizationPatterns(patterns);
+ populateLinalgNamedOpsGeneralizationPatterns(patterns, false);
// Lifting paths (named <- category <- generic)
- if (genericToNamed || genericToCategory) {
- populateLinalgGenericOpsSpecializationPatterns(patterns, genericToCategory);
- }
+ if (genericToNamed || categoryToNamed)
+ populateLinalgGenericOpsSpecializationPatterns(patterns, false);
+
+ // Category paths (named -> category <- generic)
+ if (genericToCategory)
+ populateLinalgGenericOpsSpecializationPatterns(patterns, true);
+ if (namedToCategory)
+ populateLinalgNamedOpsGeneralizationPatterns(patterns, true);
+
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
signalPassFailure();
diff --git a/mlir/test/Dialect/Linalg/categorize-named-ops.mlir b/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
index 42d22f84f0c45..2a0fe72d13cca 100644
--- a/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
@@ -1,4 +1,6 @@
// RUN: mlir-opt %s -split-input-file -linalg-categorize-ops | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=named-to-category | FileCheck %s -check-prefix=NAMED
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-category | FileCheck %s -check-prefix=GENERIC
func.func @categorize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) {
linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>)
@@ -21,6 +23,12 @@ func.func @categorize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>,
// CHECK-SAME: ins(%[[A]], %[[B]] : memref<16x8xf32>, memref<8x32xf32>)
// CHECK-SAME: outs(%[[C]] : memref<16x32xf32>)
+// NAMED-LABEL: func @categorize_matmul_buffer
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_buffer
+// GENERIC-NOT: linalg.contract
+
// -----
func.func @categorize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
@@ -41,6 +49,12 @@ func.func @categorize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>,
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)
+// NAMED-LABEL: func @categorize_matmul_tensor
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_tensor
+// GENERIC-NOT: linalg.contract
+
// -----
// Unsigned cast is preserved on the category op.
@@ -57,6 +71,12 @@ func.func @categorize_matmul_unsigned_cast(%A: tensor<16x8xi16>, %B: tensor<8x32
// CHECK: linalg.contract
// CHECK-SAME: {cast = #linalg.type_fn<cast_unsigned>}
+// NAMED-LABEL: func @categorize_matmul_unsigned_cast
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_unsigned_cast
+// GENERIC-NOT: linalg.contract
+
// -----
func.func @categorize_batch_matmul(%A: tensor<2x3x5xf32>, %B: tensor<2x5x7xf32>, %C: tensor<2x3x7xf32>) -> tensor<2x3x7xf32> {
@@ -77,6 +97,12 @@ func.func @categorize_batch_matmul(%A: tensor<2x3x5xf32>, %B: tensor<2x5x7xf32>,
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x3x5xf32>, tensor<2x5x7xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<2x3x7xf32>)
+// NAMED-LABEL: func @categorize_batch_matmul
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_batch_matmul
+// GENERIC-NOT: linalg.contract
+
// -----
func.func @categorize_batch_reduce_matmul(%A: memref<7x8x9xf32>, %B: memref<7x9x8xf32>, %C: memref<8x8xf32>) {
@@ -97,6 +123,12 @@ func.func @categorize_batch_reduce_matmul(%A: memref<7x8x9xf32>, %B: memref<7x9x
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<7x8x9xf32>, memref<7x9x8xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<8x8xf32>)
+// NAMED-LABEL: func @categorize_batch_reduce_matmul
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_batch_reduce_matmul
+// GENERIC-NOT: linalg.contract
+
// -----
///----------------------------------------------------------------------------------------
@@ -133,6 +165,12 @@ func.func @op_matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?xf32>)
+// GENERIC-LABEL: func @op_matmul
+// GENERIC: linalg.contract
+
+// NAMED-LABEL: func @op_matmul
+// NAMED-NOT: linalg.contract
+
// -----
// Matmul transpose A: A is accessed as (k, m) instead of (m, k)
@@ -164,6 +202,12 @@ func.func @op_matmul_transpose_a(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?xf32>)
+// GENERIC-LABEL: func @op_matmul_transpose_a
+// GENERIC: linalg.contract
+
+// NAMED-LABEL: func @op_matmul_transpose_a
+// NAMED-NOT: linalg.contract
+
// -----
// Matmul transpose B: B is accessed as (n, k) instead of (k, n)
@@ -196,6 +240,12 @@ func.func @op_matmul_transpose_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?xf32>)
+// GENERIC-LABEL: func @op_matmul_transpose_b
+// GENERIC: linalg.contract
+
+// NAMED-LABEL: func @op_matmul_transpose_b
+// NAMED-NOT: linalg.contract
+
// -----
#mapbA = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
@@ -227,6 +277,12 @@ func.func @op_batch_matmul(%A: tensor<2x16x8xf32>, %B: tensor<2x8x16xf32>,
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<2x16x16xf32>)
+// GENERIC-LABEL: func @op_batch_matmul
+// GENERIC: linalg.contract
+
+// NAMED-LABEL: func @op_batch_matmul
+// NAMED-NOT: linalg.contract
+
// -----
// A multi-reduction contraction (not a named-op matmul, still categorizes).
@@ -253,6 +309,12 @@ func.func @op_multi_reduction(%A: tensor<10x20x30xf32>,
// CHECK-NOT: linalg.generic
// CHECK: linalg.contract
+// GENERIC-LABEL: func @op_multi_reduction
+// GENERIC: linalg.contract
+
+// NAMED-LABEL: func @op_multi_reduction
+// NAMED-NOT: linalg.contract
+
// -----
// TODO: named matvec.
@@ -277,6 +339,12 @@ func.func @op_matvec(%A: tensor<?x?xf32>, %B: tensor<?xf32>, %Out: tensor<?xf32>
// CHECK-NOT: linalg.generic
// CHECK: linalg.contract
+// GENERIC-LABEL: func @op_matvec
+// GENERIC: linalg.contract
+
+// NAMED-LABEL: func @op_matvec
+// NAMED-NOT: linalg.contract
+
// -----
#mapmA = affine_map<(m, n, k, m0, n0, k0) -> (m, k, m0, k0)>
@@ -301,3 +369,9 @@ func.func @op_mmt4d(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
// CHECK-LABEL: func @op_mmt4d
// CHECK-NOT: linalg.generic
// CHECK: linalg.contract
+
+// GENERIC-LABEL: func @op_mmt4d
+// GENERIC: linalg.contract
+
+// NAMED-LABEL: func @op_mmt4d
+// NAMED-NOT: linalg.contract
diff --git a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
index 45f00f8ef1619..ef8c25e819716 100644
--- a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
@@ -1,4 +1,6 @@
// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=named-to-generic | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=category-to-generic | FileCheck %s
func.func @generalize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) {
linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>)
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index 6441f76eb849c..d6787084668a8 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -1,8 +1,6 @@
-// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-named \
-// RUN: | FileCheck %s --check-prefix=NAMED,ALL
-
-// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-category \
-// RUN: | FileCheck %s --check-prefix=CATEGORY,ALL
+// RUN: mlir-opt %s -split-input-file -linalg-specialize-generic-ops | FileCheck %s --check-prefix=ALL
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-named | FileCheck %s --check-prefix=NAMED,ALL
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-category | FileCheck %s --check-prefix=CATEGORY,ALL
#umap = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
func.func @unary_ops(%A: tensor<?x?x?xf32>, %Out: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
>From d660b8d9f7061e8585a6c6e82533af085b77790f Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 07:39:30 -0700
Subject: [PATCH 07/12] fixup
---
mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
index ac05fca4621b3..78ded9ccc21c0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -57,7 +57,6 @@ void LinalgMorphOpsPass::runOnOperation() {
populateLinalgGenericOpsSpecializationPatterns(patterns, true);
if (namedToCategory)
populateLinalgNamedOpsGeneralizationPatterns(patterns, true);
-
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
signalPassFailure();
>From be8d56f8bc20fb2a16055a4f6852905a583db0ed Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 07:51:17 -0700
Subject: [PATCH 08/12] clang-format
---
.../Dialect/Linalg/Transforms/Transforms.h | 31 ++++++++++---------
.../Dialect/Linalg/IR/LinalgInterfaces.cpp | 1 -
.../Linalg/Transforms/Generalization.cpp | 9 +++---
.../Dialect/Linalg/Transforms/Specialize.cpp | 16 +++++-----
4 files changed, 29 insertions(+), 28 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 699ffa61a8863..575b8c20b790c 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -920,8 +920,7 @@ FailureOr<GenericOp> interchangeGenericOp(RewriterBase &rewriter,
/// Create a GenericOp or CategoryOp from the given named operation `linalgOp`
/// and replace the given `linalgOp`. Return failure if `linalgOp` is a
/// GenericOp or misses a region builder.
-FailureOr<LinalgOp> generalizeNamedOp(RewriterBase &rewriter,
- LinalgOp linalgOp,
+FailureOr<LinalgOp> generalizeNamedOp(RewriterBase &rewriter, LinalgOp linalgOp,
bool emitCategoryOps = false);
/// Replace the given GenericOp with a namedOp or categoryOp.
@@ -1668,10 +1667,11 @@ FailureOr<LinalgOp> downscaleSizeOneWindowedConvolution(RewriterBase &rewriter,
struct LinalgGeneralizationPattern
: public OpInterfaceRewritePattern<LinalgOp> {
- LinalgGeneralizationPattern(
- MLIRContext *context, bool emitCategoryOps = false,
- PatternBenefit benefit = 1)
- : OpInterfaceRewritePattern<LinalgOp>(context, benefit), emitCategoryOps(emitCategoryOps) {}
+ LinalgGeneralizationPattern(MLIRContext *context,
+ bool emitCategoryOps = false,
+ PatternBenefit benefit = 1)
+ : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
+ emitCategoryOps(emitCategoryOps) {}
/// `matchAndRewrite` implementation that returns the significant
/// transformed pieces of IR.
@@ -1691,10 +1691,11 @@ struct LinalgGeneralizationPattern
struct LinalgSpecializationPattern : public OpRewritePattern<GenericOp> {
- LinalgSpecializationPattern(
- MLIRContext *context, bool emitCategoryOps = false,
- PatternBenefit benefit = 1)
- : OpRewritePattern<GenericOp>(context, benefit), emitCategoryOps(emitCategoryOps) {}
+ LinalgSpecializationPattern(MLIRContext *context,
+ bool emitCategoryOps = false,
+ PatternBenefit benefit = 1)
+ : OpRewritePattern<GenericOp>(context, benefit),
+ emitCategoryOps(emitCategoryOps) {}
FailureOr<GenericOp>
returningMatchAndRewrite(GenericOp op, PatternRewriter &rewriter) const {
@@ -1710,10 +1711,10 @@ struct LinalgSpecializationPattern : public OpRewritePattern<GenericOp> {
bool emitCategoryOps;
};
-struct LinalgCategorizationPattern : public OpInterfaceRewritePattern<LinalgOp> {
+struct LinalgCategorizationPattern
+ : public OpInterfaceRewritePattern<LinalgOp> {
- LinalgCategorizationPattern(
- MLIRContext *context, PatternBenefit benefit = 1)
+ LinalgCategorizationPattern(MLIRContext *context, PatternBenefit benefit = 1)
: OpInterfaceRewritePattern<LinalgOp>(context, benefit) {}
FailureOr<LinalgOp>
@@ -1943,8 +1944,8 @@ void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns,
/// e.g. linalg.generic that takes a tensor and computes a polynomial such as:
/// p(x) = an*x^n + ... + a1x + a0
/// There is no equivalent named op to convert to. Many such cases exist.
-void populateLinalgGenericOpsSpecializationPatterns(RewritePatternSet &patterns,
- bool emitCategoryOps = false);
+void populateLinalgGenericOpsSpecializationPatterns(
+ RewritePatternSet &patterns, bool emitCategoryOps = false);
/// Populates `patterns` with patterns that fold operations like
/// `linalg.transform` into elementwise op map.
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index e1152b287330b..0ee3eaf8e6900 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -301,7 +301,6 @@ bool linalg::isaElemwiseSingleTernaryOpInterface(linalg::GenericOp op) {
!op.payloadUsesValueFromOperand(inputOpOperand2));
}
-
//===----------------------------------------------------------------------===//
// ContractionOpInterface implementation
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index dd3a89bc21000..d9ac5f68a2500 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -68,8 +68,7 @@ static FailureOr<LinalgOp> generalizeToContractOp(RewriterBase &rewriter,
attributes.push_back(rewriter.getNamedAttr("cast", castAttr));
LinalgOp contractOp = rewriter.replaceOpWithNewOp<ContractOp>(
- namedOp,
- ValueRange{namedOp.getDpsInputs()[0], namedOp.getDpsInputs()[1]},
+ namedOp, ValueRange{namedOp.getDpsInputs()[0], namedOp.getDpsInputs()[1]},
ValueRange{namedOp.getDpsInits()[0]}, attributes);
return contractOp;
}
@@ -85,7 +84,8 @@ FailureOr<LinalgOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,
FailureOr<LinalgOp> contractOp = generalizeToContractOp(rewriter, linalgOp);
if (succeeded(contractOp))
return contractOp;
- return rewriter.notifyMatchFailure(linalgOp, "failed to categorize to named op");
+ return rewriter.notifyMatchFailure(linalgOp,
+ "failed to categorize to named op");
}
SmallVector<Value> inputs = linalgOp.getDpsInputs();
@@ -136,5 +136,6 @@ void LinalgGeneralizeNamedOpsPass::runOnOperation() {
void mlir::linalg::populateLinalgNamedOpsGeneralizationPatterns(
RewritePatternSet &patterns, bool emitCategoryOps) {
- patterns.add<LinalgGeneralizationPattern>(patterns.getContext(), emitCategoryOps);
+ patterns.add<LinalgGeneralizationPattern>(patterns.getContext(),
+ emitCategoryOps);
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index df11be9b38c75..d16e1aea5a6c7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -594,8 +594,8 @@ static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
// No named variant matched; fall back to the generic `linalg.contract` op,
// which supports a wider range of variants.
- return replaceWithMatmulVariant<ContractOp>(
- rewriter, genericOp, castTy, genericOp.getIndexingMapsArray());
+ return replaceWithMatmulVariant<ContractOp>(rewriter, genericOp, castTy,
+ genericOp.getIndexingMapsArray());
}
/// Utility to specialize a `genericOp` with a convolution op of type `ConvOpTy`
@@ -698,9 +698,9 @@ static FailureOr<LinalgOp> specializeLinalgConvolutions(RewriterBase &rewriter,
//===----------------------------------------------------------------------===//
// Categorize linalg generic to named op where possible.
//===----------------------------------------------------------------------===//
-FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(
- RewriterBase &rewriter, GenericOp genericOp,
- bool emitCategoryOps) {
+FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(RewriterBase &rewriter,
+ GenericOp genericOp,
+ bool emitCategoryOps) {
// Elementwise - e.g. exp, add are always category ops
if (isaElemwiseSingleUnaryOpInterface(genericOp) ||
isaElemwiseSingleBinaryOpInterface(genericOp) ||
@@ -710,8 +710,7 @@ FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(
// Contraction - e.g. matmul
if (isaContractionOpInterface(genericOp)) {
- return specializeLinalgContractions(rewriter, genericOp,
- emitCategoryOps);
+ return specializeLinalgContractions(rewriter, genericOp, emitCategoryOps);
}
// Early exit in case of category specialization.
@@ -788,5 +787,6 @@ void LinalgSpecializeGenericOpsPass::runOnOperation() {
void mlir::linalg::populateLinalgGenericOpsSpecializationPatterns(
RewritePatternSet &patterns, bool emitCategoryOps) {
- patterns.add<LinalgSpecializationPattern>(patterns.getContext(), emitCategoryOps);
+ patterns.add<LinalgSpecializationPattern>(patterns.getContext(),
+ emitCategoryOps);
}
>From 1b19eb404c30ebc983597e2956d6f8f5bf7d9693 Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 07:56:04 -0700
Subject: [PATCH 09/12] more format
---
mlir/include/mlir/Dialect/Linalg/IR/Linalg.h | 1 -
mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp | 6 +++---
mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp | 13 ++++++-------
3 files changed, 9 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
index aaf4a61b19567..5a7a5d989782e 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
@@ -161,7 +161,6 @@ struct ArityGroupAndKind {
};
ArityGroupAndKind getArityGroupAndKind(ElementwiseKind kind);
-
/// Specialization of `linalg.matmul` op that has a transpose map on A
class MatmulTransposeAOp : public MatmulOp {
/// Create an affine map for a transpose-A matmul. Used only in the builders.
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index 0ee3eaf8e6900..330642e6b0bbc 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -227,9 +227,9 @@ linalg::isaTransposeOpInterface(GenericOp op) {
//===----------------------------------------------------------------------===//
// Elementwise Single Unary/Binary-OpInterface implementation
//===----------------------------------------------------------------------===//
-static bool
-isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op,
- unsigned arity) {
+
+static bool isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op,
+ unsigned arity) {
// Check all loops are parallel.
if (!op.isAllParallelLoops() || op.getNumLoops() < 1)
return false;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index d16e1aea5a6c7..cacdf11652a72 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -200,7 +200,7 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
op->getOperand(scalarOprIdx));
auto scalarBroadcastMap =
AffineMap::get(genericOp.getNumParallelLoops(), /*symbolCount=*/0,
- rewriter.getContext());
+ rewriter.getContext());
indexingMaps.insert(indexingMaps.begin() + scalarOprIdx,
scalarBroadcastMap);
}
@@ -218,7 +218,7 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
divOp.getLhs().getDefiningOp()))
if (cast<FloatAttr>(constOp.getValue()).getValue().isExactlyValue(1.0))
return replaceOp(ElementwiseKind::reciprocal,
- /*mayHoistScalarOperand=*/false);
+ /*mayHoistScalarOperand=*/false);
}
// Square
@@ -227,9 +227,9 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
return replaceOp(ElementwiseKind::square);
// Boolean-typed `add` and `mul`.
- if (isBinary && llvm::all_of(
- op->getOperands(), [](Value v) {
- return v.getType().isInteger(1); })) {
+ if (isBinary && llvm::all_of(op->getOperands(), [](Value v) {
+ return v.getType().isInteger(1);
+ })) {
if (isa<arith::OrIOp>(op))
return replaceOp(ElementwiseKind::add);
if (isa<arith::AndIOp>(op))
@@ -246,8 +246,7 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
}
return rewriter.notifyMatchFailure(
- genericOp,
- "elementwise operation cannot be specialized to category op");
+ genericOp, "elementwise operation cannot be specialized to category op");
}
//===----------------------------------------------------------------------===//
>From e1a4305b37f20140819a1c768cd873f93add61b5 Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Fri, 11 Sep 2026 08:25:16 -0700
Subject: [PATCH 10/12] address comments
---
.../Linalg/Transforms/Generalization.cpp | 3 +
.../Dialect/Linalg/categorize-named-ops.mlir | 93 +++++++++++++++++++
2 files changed, 96 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index d9ac5f68a2500..d313d773d21f0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -48,6 +48,9 @@ static LogicalResult generalizeNamedOpPrecondition(LinalgOp linalgOp) {
// operand indexing maps and cast semantics. Returns failure for other ops.
static FailureOr<LinalgOp> generalizeToContractOp(RewriterBase &rewriter,
LinalgOp namedOp) {
+ // These are the ODS-defined matmul-like operations.
+ // For OpDSL declared contractions, please move them to ODS first,
+ // then add them to the isa<> check below + tests.
if (!isa<MatmulOp, BatchMatmulOp, BatchReduceMatmulOp>(
namedOp.getOperation()))
return failure();
diff --git a/mlir/test/Dialect/Linalg/categorize-named-ops.mlir b/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
index 2a0fe72d13cca..d95a82a75c643 100644
--- a/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/categorize-named-ops.mlir
@@ -57,6 +57,99 @@ func.func @categorize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>,
// -----
+// Transpose A: A is accessed as (k, m) instead of (m, k).
+func.func @categorize_matmul_transpose_a(%A: tensor<8x16xf32>, %B: tensor<8x32xf32>,
+ %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
+ %0 = linalg.matmul
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d0)>,
+ affine_map<(d0, d1, d2) -> (d2, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>]
+ ins(%A, %B: tensor<8x16xf32>, tensor<8x32xf32>)
+ outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+ return %0: tensor<16x32xf32>
+}
+
+// CHECK-DAG: #[[TA_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @categorize_matmul_transpose_a
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[TA_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<8x16xf32>, tensor<8x32xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)
+
+// NAMED-LABEL: func @categorize_matmul_transpose_a
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_transpose_a
+// GENERIC-NOT: linalg.contract
+
+// -----
+
+// Broadcast B: B is accessed as (k) only, broadcasting over n.
+func.func @categorize_matmul_broadcast_b(%A: tensor<16x8xf32>, %B: tensor<8xf32>,
+ %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
+ %0 = linalg.matmul
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
+ affine_map<(d0, d1, d2) -> (d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>]
+ ins(%A, %B: tensor<16x8xf32>, tensor<8xf32>)
+ outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+ return %0: tensor<16x32xf32>
+}
+
+// CHECK-DAG: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[BB_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @categorize_matmul_broadcast_b
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[BB_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)
+
+// NAMED-LABEL: func @categorize_matmul_broadcast_b
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_broadcast_b
+// GENERIC-NOT: linalg.contract
+
+// -----
+
+// Transpose and broadcast A: A is accessed as (k) only, broadcasting over m.
+func.func @categorize_matmul_transpose_broadcast_a(%A: tensor<8xf32>, %B: tensor<8x32xf32>,
+ %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
+ %0 = linalg.matmul
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d2)>,
+ affine_map<(d0, d1, d2) -> (d2, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>]
+ ins(%A, %B: tensor<8xf32>, tensor<8x32xf32>)
+ outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+ return %0: tensor<16x32xf32>
+}
+
+// CHECK-DAG: #[[TBA_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2)>
+// CHECK-DAG: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @categorize_matmul_transpose_broadcast_a
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[TBA_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<8xf32>, tensor<8x32xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)
+
+// NAMED-LABEL: func @categorize_matmul_transpose_broadcast_a
+// NAMED: linalg.contract
+
+// GENERIC-LABEL: func @categorize_matmul_transpose_broadcast_a
+// GENERIC-NOT: linalg.contract
+
+// -----
+
// Unsigned cast is preserved on the category op.
func.func @categorize_matmul_unsigned_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
%C: tensor<16x32xi32>) -> tensor<16x32xi32> {
>From 052bbbb786211919ff6d80e9b3259c642cd88c7b Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Mon, 14 Sep 2026 03:07:29 -0700
Subject: [PATCH 11/12] address comments
- allow swapped ternary (operands 2 & 3)
- add discardable attributes to contract
---
.../Dialect/Linalg/IR/LinalgInterfaces.cpp | 17 ++++++++----
.../Linalg/Transforms/Generalization.cpp | 8 ++++++
.../Dialect/Linalg/Transforms/Specialize.cpp | 24 +++++++++++------
.../Linalg/specialize-generic-ops.mlir | 27 +++++++++++++++++++
4 files changed, 63 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index 330642e6b0bbc..bb3146c87b430 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -228,8 +228,7 @@ linalg::isaTransposeOpInterface(GenericOp op) {
// Elementwise Single Unary/Binary-OpInterface implementation
//===----------------------------------------------------------------------===//
-static bool isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op,
- unsigned arity) {
+static bool isaElemwiseSingleOpInterface(linalg::GenericOp op, unsigned arity) {
// Check all loops are parallel.
if (!op.isAllParallelLoops() || op.getNumLoops() < 1)
return false;
@@ -266,7 +265,7 @@ static bool isaElemwiseSingleUnaryOrBinaryOpInterface(linalg::GenericOp op,
bool linalg::isaElemwiseSingleUnaryOpInterface(linalg::GenericOp op) {
// All basic elemwise checks.
- if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 1))
+ if (!isaElemwiseSingleOpInterface(op, 1))
return false;
// Check input is actually used.
@@ -277,7 +276,7 @@ bool linalg::isaElemwiseSingleUnaryOpInterface(linalg::GenericOp op) {
bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op) {
// All basic elemwise checks.
- if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 2))
+ if (!isaElemwiseSingleOpInterface(op, 2))
return false;
// Check both inputs are used (elementwise).
@@ -289,7 +288,15 @@ bool linalg::isaElemwiseSingleBinaryOpInterface(linalg::GenericOp op) {
bool linalg::isaElemwiseSingleTernaryOpInterface(linalg::GenericOp op) {
// All basic elemwise checks.
- if (!isaElemwiseSingleUnaryOrBinaryOpInterface(op, 3))
+ if (!isaElemwiseSingleOpInterface(op, 3))
+ return false;
+
+ // The only ternary (select) has a boolean argument as its first operand.
+ // If we add more ternaries later, we need to change this check.
+ // But for now, it simplifies other checks, like checking for swapped
+ // operands.
+ if (!getElementTypeOrSelf(op.getDpsInputOperand(0)->get().getType())
+ .isInteger(1))
return false;
// Check all three inputs are used (elementwise).
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
index d313d773d21f0..5a51328b3eee4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp
@@ -73,6 +73,14 @@ static FailureOr<LinalgOp> generalizeToContractOp(RewriterBase &rewriter,
LinalgOp contractOp = rewriter.replaceOpWithNewOp<ContractOp>(
namedOp, ValueRange{namedOp.getDpsInputs()[0], namedOp.getDpsInputs()[1]},
ValueRange{namedOp.getDpsInits()[0]}, attributes);
+
+ // Discardable attributes carry user-defined metadata (e.g., annotations for
+ // downstream passes). Generalization is a semantics-preserving
+ // transformation, so dropping this metadata would be unexpected. This is safe
+ // because discardable attributes are by definition independent of op
+ // semantics.
+ contractOp->setDiscardableAttrs(namedOp->getDiscardableAttrDictionary());
+
return contractOp;
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index cacdf11652a72..b52ec3258aef4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -47,14 +47,17 @@ using namespace mlir::linalg;
// %0 = arith.subf %b, %a : f32
// linalg.yield %0: f32
// Former is linalg.sub(a,b), latter is linalg.sub(b,a).
-static bool areBinOpsSwapped(GenericOp genericOp) {
+static bool areBinOpsSwapped(GenericOp genericOp, bool isTernary) {
Block *body = genericOp.getBody();
Operation *op = &body->front();
bool swapped = false;
- if (op->getOpOperand(0).get() != body->getArgument(0)) {
+ if (op->getOpOperand(0 + isTernary).get() !=
+ body->getArgument(0 + isTernary)) {
swapped = true;
- assert(op->getOpOperand(0).get() == body->getArgument(1) &&
- op->getOpOperand(1).get() == body->getArgument(0) &&
+ assert(op->getOpOperand(0 + isTernary).get() ==
+ body->getArgument(1 + isTernary) &&
+ op->getOpOperand(1 + isTernary).get() ==
+ body->getArgument(0 + isTernary) &&
"binary op uses just one block arg");
}
return swapped;
@@ -140,7 +143,7 @@ static std::optional<ElementwiseKind> getElementwiseKind(Operation *op) {
.Default([](Operation *) { return std::nullopt; });
}
-// Attempt to specialize unary or binary linalg.generic ops
+// Attempt to specialize unary/binary/ternary linalg.generic ops
// to linalg.elementwise.
//
// Example:
@@ -172,12 +175,14 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
unsigned arity = genericOp.getNumDpsInputs();
bool isUnary = arity == 1;
bool isBinary = arity == 2;
+ bool isTernary = arity == 3;
// Will inspect the body operation to determine named op or elementwise kind.
Operation *op = &genericOp.getBody()->front();
// Detect variations from canonical forms.
- bool hasSwappedOperands = isBinary && areBinOpsSwapped(genericOp);
+ bool hasSwappedOperands =
+ (isBinary || isTernary) && areBinOpsSwapped(genericOp, isTernary);
int scalarOprIdx = -1;
bool hasScalarOperand = isUnary && op->getNumOperands() == 2 &&
findIndexOfScalarOperand(genericOp, scalarOprIdx);
@@ -190,8 +195,11 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
SmallVector<Value> inputs = genericOp.getDpsInputs();
SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
if (hasSwappedOperands) {
- std::swap(inputs[0], inputs[1]);
- std::swap(indexingMaps[0], indexingMaps[1]);
+ // Ternary indices are +1, since the first is the boolean mask.
+ // If new ternary with non-booleans as first argument are created,
+ // we may need to calculate all combinations possible.
+ std::swap(inputs[0+isTernary], inputs[1+isTernary]);
+ std::swap(indexingMaps[0+isTernary], indexingMaps[1+isTernary]);
}
if (hasScalarOperand && mayHoistScalarOperand) {
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index d6787084668a8..da51685fc584f 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -688,6 +688,33 @@ func.func @ternary_op_select_f32(%C: tensor<?x?xi1>, %T: tensor<?x?xf32>,
// -----
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @ternary_op_select_f32_swapped(%C: tensor<?x?xi1>, %T: tensor<?x?xf32>,
+ %F: tensor<?x?xf32>, %Out: tensor<?x?xf32>)
+ -> tensor<?x?xf32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map, #map],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%C, %T, %F : tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%Out : tensor<?x?xf32>) {
+ ^bb0(%in: i1, %in_0: f32, %in_1: f32, %out: f32):
+ %v = arith.select %in, %in_1, %in_0 : f32
+ linalg.yield %v : f32
+ } -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// ALL-LABEL: ternary_op_select_f32_swapped
+// ALL-SAME: %[[C:.+]]: [[ITY:tensor<\?x\?xi1>]], %[[T:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[F:.+]]: [[TTY]],
+// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
+
+// ALL-NOT: linalg.generic
+// ALL: linalg.elementwise <select>
+// ALL-SAME: ins(%[[C]], %[[F]], %[[T]] : [[ITY]], [[TTY]], [[TTY]])
+// ALL-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
+
+// -----
+
#map = affine_map<(d0, d1) -> (d0, d1)>
func.func @ternary_op_select_i32(%C: tensor<?x?xi1>, %T: tensor<?x?xi32>,
%F: tensor<?x?xi32>, %Out: tensor<?x?xi32>)
>From ef0c8fc8a1cc76a9fa50e749d715ab99eea8eeba Mon Sep 17 00:00:00 2001
From: rengolin <rengolin at systemcall.eu>
Date: Mon, 14 Sep 2026 03:14:01 -0700
Subject: [PATCH 12/12] format
---
mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index b52ec3258aef4..8fce91acf4b71 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -198,8 +198,8 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
// Ternary indices are +1, since the first is the boolean mask.
// If new ternary with non-booleans as first argument are created,
// we may need to calculate all combinations possible.
- std::swap(inputs[0+isTernary], inputs[1+isTernary]);
- std::swap(indexingMaps[0+isTernary], indexingMaps[1+isTernary]);
+ std::swap(inputs[0 + isTernary], inputs[1 + isTernary]);
+ std::swap(indexingMaps[0 + isTernary], indexingMaps[1 + isTernary]);
}
if (hasScalarOperand && mayHoistScalarOperand) {
More information about the Mlir-commits
mailing list