[Mlir-commits] [mlir] [mlir][linalg] Morphism across linalg named, category and generic ops. (PR #148424)
Javed Absar
llvmlistbot at llvm.org
Thu Aug 7 03:41:24 PDT 2025
https://github.com/javedabsar1 updated https://github.com/llvm/llvm-project/pull/148424
>From a005f0c80a8dfc579eec4a3bcb7c84cf1acb07d2 Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Sat, 12 Jul 2025 09:59:56 -0400
Subject: [PATCH 1/5] [mlir][linalg] Convert linalg.named to linalg.elementwise
op.
Convert linalg.named ops which are elementwise (e.g. add/exp)
to `linalg.elementwise`. Currently, named ops have to drop to
linalg.generic (--generalize-named-ops), where one figures
out which generic are elementwise. Also, folding of broadcast
or transpose can occur then only at generic level. Instead,
with this rewrite, these can happen now at linalg.elementwise.
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 5 +
.../Dialect/Linalg/Transforms/Transforms.h | 4 +
.../Dialect/Linalg/Transforms/CMakeLists.txt | 1 +
.../Linalg/Transforms/NamedToElementwise.cpp | 118 ++++++++++++++++++
.../elementwise/named_to_elementwise.mlir | 38 ++++++
5 files changed, 166 insertions(+)
create mode 100644 mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
create mode 100644 mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 373842c9b03de..f2c1b99b138bc 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -99,6 +99,11 @@ def LinalgSpecializeGenericOpsPass : Pass<"linalg-specialize-generic-ops"> {
let dependentDialects = ["linalg::LinalgDialect"];
}
+def LinalgNamedToElementwisePass : Pass<"linalg-named-to-elementwise"> {
+ let summary = "Convert linalg named ops to elementwise where possible";
+ let dependentDialects = ["linalg::LinalgDialect"];
+}
+
def LinalgFoldIntoElementwisePass : Pass<"linalg-fold-into-elementwise"> {
let summary = "Fold transform, broadcast and other ops into elementwise";
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 d4ffe0a91fcfe..1e5b5d46de55f 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1831,6 +1831,10 @@ void populateLinalgNamedOpsGeneralizationPatterns(RewritePatternSet &patterns);
void populateLinalgGenericOpsSpecializationPatterns(
RewritePatternSet &patterns);
+/// Populates `patterns` that convert linalg named ops e.g. `linalg.add`
+/// to equivalent `linalg.elementwise`.
+void populateLinalgNamedToElementwisePatterns(RewritePatternSet &patterns);
+
/// Populates `patterns` with patterns that fold operations like
/// `linalg.transform` into elementwise op map.
void populateLinalgFoldIntoElementwisePatterns(RewritePatternSet &patterns);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
index 70f846e5bbd20..cf375792ae308 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
@@ -26,6 +26,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
TransposeMatmul.cpp
ShardingInterfaceImpl.cpp
NamedOpConversions.cpp
+ NamedToElementwise.cpp
BlockPackMatmul.cpp
PackAndUnpackPatterns.cpp
Padding.cpp
diff --git a/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
new file mode 100644
index 0000000000000..1303b7cb5f6f9
--- /dev/null
+++ b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
@@ -0,0 +1,118 @@
+//===- NamedToElementwise.cpp - convert linalg named op into elementwise --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements rewriting those linalg named ops that are essentially
+// elementwise e.g. `linalg.exp`, to `linalg.elementwise`. This allows further
+// optimization on `linalg.elementwise` such as folding transpose, broadcast.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_LINALGNAMEDTOELEMENTWISEPASS
+#include "mlir/Dialect/Linalg/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+#define DEBUG_TYPE "linalg-named-to-elementwise"
+
+namespace {
+ElementwiseKind getKind(Operation *op) {
+ return llvm::TypeSwitch<Operation *, ElementwiseKind>(op)
+ .Case([](SelectOp) { return ElementwiseKind::select; })
+ .Case([](AddOp) { return ElementwiseKind::add; })
+ .Case([](SubOp) { return ElementwiseKind::sub; })
+ .Case([](MulOp) { return ElementwiseKind::mul; })
+ .Case([](DivOp) { return ElementwiseKind::div; })
+ .Case([](DivUnsignedOp) { return ElementwiseKind::div_unsigned; })
+ .Case([](PowFOp) { return ElementwiseKind::powf; })
+ .Case([](ExpOp) { return ElementwiseKind::exp; })
+ .Case([](LogOp) { return ElementwiseKind::log; })
+ .Case([](AbsOp) { return ElementwiseKind::abs; })
+ .Case([](CeilOp) { return ElementwiseKind::ceil; })
+ .Case([](FloorOp) { return ElementwiseKind::floor; })
+ .Case([](NegFOp) { return ElementwiseKind::negf; })
+ .Case([](ReciprocalOp) { return ElementwiseKind::reciprocal; })
+ .Case([](RoundOp) { return ElementwiseKind::round; })
+ .Case([](SqrtOp) { return ElementwiseKind::sqrt; })
+ .Case([](RsqrtOp) { return ElementwiseKind::rsqrt; })
+ .Case([](SquareOp) { return ElementwiseKind::square; })
+ .Case([](TanhOp) { return ElementwiseKind::tanh; })
+ .Case([](ErfOp) { return ElementwiseKind::erf; })
+ .Default([&](Operation *op) {
+ assert(false && "unexpected op");
+ return ElementwiseKind::sub;
+ });
+}
+
+template <typename NamedOpTy>
+struct NamedToElementwisePattern : public OpRewritePattern<NamedOpTy> {
+ using OpRewritePattern<NamedOpTy>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(NamedOpTy op,
+ PatternRewriter &rewriter) const override {
+ SmallVector<NamedAttribute> attrs;
+ auto kindAttr = ElementwiseKindAttr::get(op.getContext(), getKind(op));
+ attrs.push_back(rewriter.getNamedAttr("kind", kindAttr));
+ attrs.push_back(
+ rewriter.getNamedAttr("indexing_maps", op.getIndexingMaps()));
+
+ rewriter.replaceOpWithNewOp<ElementwiseOp>(op, op.getDpsInputs(),
+ op.getDpsInits(), attrs);
+ return success();
+ }
+};
+
+struct LinalgNamedToElementwisePass
+ : public impl::LinalgNamedToElementwisePassBase<
+ LinalgNamedToElementwisePass> {
+ using impl::LinalgNamedToElementwisePassBase<
+ LinalgNamedToElementwisePass>::LinalgNamedToElementwisePassBase;
+
+ void runOnOperation() override {
+ Operation *op = getOperation();
+ RewritePatternSet patterns(op->getContext());
+ populateLinalgNamedToElementwisePatterns(patterns);
+
+ if (failed(applyPatternsGreedily(op, std::move(patterns))))
+ return signalPassFailure();
+ }
+};
+} // namespace
+
+void mlir::linalg::populateLinalgNamedToElementwisePatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<NamedToElementwisePattern<AddOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<SubOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<MulOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<DivOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<DivUnsignedOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<PowFOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<ExpOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<LogOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<AbsOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<CeilOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<FloorOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<NegFOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<ReciprocalOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<RoundOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<SqrtOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<RsqrtOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<SquareOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<TanhOp>>(patterns.getContext());
+ patterns.add<NamedToElementwisePattern<ErfOp>>(patterns.getContext());
+}
diff --git a/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir b/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
new file mode 100644
index 0000000000000..3dc8275117336
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
@@ -0,0 +1,38 @@
+// RUN: mlir-opt %s -linalg-named-to-elementwise -split-input-file | FileCheck %s
+
+// CHECK: @exp(%[[A:.+]]: tensor<16x8xf32>, %[[B:.+]]: tensor<16x8xf32>) -> tensor<16x8xf32> {
+// CHECK: {{.*}} = linalg.elementwise
+// CHECK-SAME: kind=#linalg.elementwise_kind<exp>
+// CHECK-SAME: ins(%[[A]] : tensor<16x8xf32>)
+// CHECK-SAME: outs(%[[B]] : tensor<16x8xf32>) -> tensor<16x8xf32>
+//
+func.func @exp(%A : tensor<16x8xf32>, %B : tensor<16x8xf32>) -> tensor<16x8xf32> {
+ %exp = linalg.exp ins(%A : tensor<16x8xf32>) outs(%B : tensor<16x8xf32>) -> tensor<16x8xf32>
+ return %exp : tensor<16x8xf32>
+}
+
+// ----
+
+// CHECK: @add(%[[A:.+]]: tensor<16x8xf32>, %[[B:.+]]: tensor<16x8xf32>, %[[C:.+]]: tensor<16x8xf32>) -> tensor<16x8xf32> {
+// CHECK: {{.*}} = linalg.elementwise
+// CHECK-SAME: kind=#linalg.elementwise_kind<add>
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<16x8xf32>, tensor<16x8xf32>)
+// CHECK-SAME: outs(%[[C]] : tensor<16x8xf32>) -> tensor<16x8xf32>
+//
+func.func @add(%A : tensor<16x8xf32>, %B: tensor<16x8xf32>, %C : tensor<16x8xf32>) -> tensor<16x8xf32> {
+ %add = linalg.add ins(%A, %B : tensor<16x8xf32>, tensor<16x8xf32>) outs(%C : tensor<16x8xf32>) -> tensor<16x8xf32>
+ return %add : tensor<16x8xf32>
+}
+
+// ----
+
+// CHECK: @sub(%[[A:.+]]: memref<16x8xf32>, %[[B:.+]]: memref<16x8xf32>, %[[C:.+]]: memref<16x8xf32>) {
+// CHECK: linalg.elementwise
+// CHECK-SAME: kind=#linalg.elementwise_kind<sub>
+// CHECK-SAME: ins(%[[A]], %[[B]] : memref<16x8xf32>, memref<16x8xf32>)
+// CHECK-SAME: outs(%[[C]] : memref<16x8xf32>)
+//
+func.func @sub(%A : memref<16x8xf32>, %B: memref<16x8xf32>, %C : memref<16x8xf32>) {
+ linalg.sub ins(%A, %B : memref<16x8xf32>, memref<16x8xf32>) outs(%C : memref<16x8xf32>)
+ return
+}
>From a5d1622da8eef7bb6d6ab7791899d4f1c69de8d2 Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Sat, 19 Jul 2025 06:51:54 -0400
Subject: [PATCH 2/5] Add linalg-morph pass
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 50 +++++++++--
.../Dialect/Linalg/Transforms/CMakeLists.txt | 1 +
.../Dialect/Linalg/Transforms/MorphOps.cpp | 82 +++++++++++++++++++
.../Linalg/Transforms/NamedToElementwise.cpp | 21 -----
.../elementwise/named_to_elementwise.mlir | 2 +-
.../test/Dialect/Linalg/linalg-morph-ops.mlir | 25 ++++++
6 files changed, 154 insertions(+), 27 deletions(-)
create mode 100644 mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
create mode 100644 mlir/test/Dialect/Linalg/linalg-morph-ops.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index f2c1b99b138bc..25b635726fbf5 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -89,6 +89,51 @@ def LinalgInlineScalarOperandsPass : Pass<"linalg-inline-scalar-operands"> {
];
}
+def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
+ let summary = "Convert named op to category ops or generic and vice-versa";
+
+ let description = [{
+ Convert a linalg op from one representation to another equivalent.
+
+ For example, a linalg named op `linalg.add` can also be written as an
+ category op `linalg.elementwise`, and can also be re-written as
+ a `linalg.generic`, giving the morphism:
+
+ named-op <--> category_op (elementwise, contraction, ..) <--> generic
+
+ Generic is a bigger set than named and category ops and so not all generics
+ can be converted to single category-op or named-op. Similarly, category
+ ops are bigger set than named ops.
+
+ Note:
+ Legacy converters (will be deprecated):
+ `--linalg-generalize-named-ops` is the path `named-op --> generic-op`
+ `--linalg-specialize-generic-ops` is the path `named-op <-- generic-op`
+ }];
+ let dependentDialects = ["linalg::LinalgDialect"];
+
+ let options = [
+ // named-op <--> category <--> generic
+ Option<"namedToCategory", "named-to-category", "bool", /*default=*/"false",
+ "convert named ops to category op e.g. `linalg.elementwise`">,
+
+ Option<"categoryToGeneric", "category-to-generic", "bool", /*default=*/"false",
+ "convert category ops e.g. `linalg.elementwise` to `linalg.generic`">,
+
+ Option<"namedToGeneric", "named-to-generic", "bool", /*default=*/"false",
+ "convert named ops e.g. `linalg.add` to `linalg.generic`">,
+
+ Option<"genericToCategory", "generic-to-category", "bool", /*default=*/"false",
+ "convert generic ops to category op e.g. `linalg.contraction`">,
+
+ Option<"categoryToNamed", "category-to-named", "bool", /*default=*/"false",
+ "convert category ops to equivalent named ops">,
+
+ Option<"genericToNamed", "generic-to-named", "bool", /*default=*/"false",
+ "convert linalg.generic to equivalent named ops">
+ ];
+}
+
def LinalgGeneralizeNamedOpsPass : Pass<"linalg-generalize-named-ops"> {
let summary = "Convert named ops into generic ops";
let dependentDialects = ["linalg::LinalgDialect"];
@@ -99,11 +144,6 @@ def LinalgSpecializeGenericOpsPass : Pass<"linalg-specialize-generic-ops"> {
let dependentDialects = ["linalg::LinalgDialect"];
}
-def LinalgNamedToElementwisePass : Pass<"linalg-named-to-elementwise"> {
- let summary = "Convert linalg named ops to elementwise where possible";
- let dependentDialects = ["linalg::LinalgDialect"];
-}
-
def LinalgFoldIntoElementwisePass : Pass<"linalg-fold-into-elementwise"> {
let summary = "Fold transform, broadcast and other ops into elementwise";
let dependentDialects = ["linalg::LinalgDialect"];
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
index cf375792ae308..6ec2e9fd0be7d 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
@@ -23,6 +23,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
InlineScalarOperands.cpp
Interchange.cpp
Loops.cpp
+ MorphOps.cpp
TransposeMatmul.cpp
ShardingInterfaceImpl.cpp
NamedOpConversions.cpp
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
new file mode 100644
index 0000000000000..c0e800be3917a
--- /dev/null
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -0,0 +1,82 @@
+//===- MorphOps.cpp - conversion between named,category and generic ops ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements conversions between:
+// named <--> category (elementwise, contraction, ..) <--> generic ops.
+//
+// For example, a named op such `linalg.add` can also be re-written as an
+// equivalent category op `linalg.elementwise` and also as a `linalg.generic`.
+//
+// Generic is a bigger set than named ops and so not all generics can be
+// converted to single category-op or named-op. Similarly, category-ops
+// are bigger in representational possiblities than named ops e.g.
+// `linalg.add` has no affine maps attached, but `linalg.elementwise` does.
+//
+// Note:
+// Legacy converters (will be deprecated):
+// `--linalg-generalize-named-ops` is the path `named-op --> generic-op`
+// `--linalg-specialize-generic-ops` is the path `named-op <-- generic-op`
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_LINALGMORPHOPSPASS
+#include "mlir/Dialect/Linalg/Passes.h.inc"
+} // namespace mlir
+
+#define DEBUG_TYPE "linalg-morphism"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+namespace {
+struct LinalgMorphOpsPass
+ : public impl::LinalgMorphOpsPassBase<LinalgMorphOpsPass> {
+
+ using impl::LinalgMorphOpsPassBase<
+ LinalgMorphOpsPass>::LinalgMorphOpsPassBase;
+
+ void runOnOperation() override;
+};
+
+void LinalgMorphOpsPass::runOnOperation() {
+
+ RewritePatternSet patterns(&getContext());
+
+ // Lowering paths (named -> category -> generic)
+ if (namedToCategory) {
+ // TODO: named -> contraction-op
+ populateLinalgNamedToElementwisePatterns(patterns);
+ }
+ if (namedToGeneric || categoryToGeneric) {
+ populateLinalgNamedOpsGeneralizationPatterns(patterns);
+ }
+
+ // Lifting paths (named <- category <- generic)
+ if (genericToCategory) {
+ // TODO.
+ }
+ if (categoryToNamed) {
+ // TODO: if there is a case for this.
+ }
+ if (genericToNamed) {
+ populateLinalgGenericOpsSpecializationPatterns(patterns);
+ }
+
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ signalPassFailure();
+}
+} // namespace
diff --git a/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
index 1303b7cb5f6f9..81430b0432269 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
@@ -20,11 +20,6 @@
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
-namespace mlir {
-#define GEN_PASS_DEF_LINALGNAMEDTOELEMENTWISEPASS
-#include "mlir/Dialect/Linalg/Passes.h.inc"
-} // namespace mlir
-
using namespace mlir;
using namespace mlir::linalg;
@@ -76,22 +71,6 @@ struct NamedToElementwisePattern : public OpRewritePattern<NamedOpTy> {
return success();
}
};
-
-struct LinalgNamedToElementwisePass
- : public impl::LinalgNamedToElementwisePassBase<
- LinalgNamedToElementwisePass> {
- using impl::LinalgNamedToElementwisePassBase<
- LinalgNamedToElementwisePass>::LinalgNamedToElementwisePassBase;
-
- void runOnOperation() override {
- Operation *op = getOperation();
- RewritePatternSet patterns(op->getContext());
- populateLinalgNamedToElementwisePatterns(patterns);
-
- if (failed(applyPatternsGreedily(op, std::move(patterns))))
- return signalPassFailure();
- }
-};
} // namespace
void mlir::linalg::populateLinalgNamedToElementwisePatterns(
diff --git a/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir b/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
index 3dc8275117336..0920cbbd6e15f 100644
--- a/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -linalg-named-to-elementwise -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -linalg-morph-ops=named-to-category -split-input-file | FileCheck %s
// CHECK: @exp(%[[A:.+]]: tensor<16x8xf32>, %[[B:.+]]: tensor<16x8xf32>) -> tensor<16x8xf32> {
// CHECK: {{.*}} = linalg.elementwise
diff --git a/mlir/test/Dialect/Linalg/linalg-morph-ops.mlir b/mlir/test/Dialect/Linalg/linalg-morph-ops.mlir
new file mode 100644
index 0000000000000..d15d29b4fd532
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/linalg-morph-ops.mlir
@@ -0,0 +1,25 @@
+// Forward path `named -> category -> generic`
+// RUN: mlir-opt %s -linalg-morph-ops=named-to-category | FileCheck %s --check-prefix=NAMED_TO_CATEGORY
+// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | FileCheck %s --check-prefix=NAMED_TO_GENERIC
+// RUN: mlir-opt %s -linalg-morph-ops=named-to-category | \
+// RUN: mlir-opt %s -linalg-morph-ops=category-to-generic | FileCheck %s --check-prefix=CATEGORY_TO_GENERIC
+//
+// Backward path `named <- category <- generic`
+// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | mlir-opt %s -linalg-morph-ops=generic-to-named | \
+// RUN: FileCheck %s --check-prefix=GENERIC_TO_NAMED
+
+func.func @exp(%A : tensor<16x8xf32>, %B : tensor<16x8xf32>) -> tensor<16x8xf32> {
+ %exp = linalg.exp ins(%A : tensor<16x8xf32>) outs(%B : tensor<16x8xf32>) -> tensor<16x8xf32>
+ return %exp : tensor<16x8xf32>
+}
+// NAMED_TO_CATEGORY: linalg.elementwise
+// NAMED_TO_CATEGORY-NOT: linalg.exp
+
+// NAMED_TO_GENERIC: linalg.generic
+// NAMED_TO_GENERIC-NOT: linalg.exp
+
+// CATEGORY_TO_GENERIC: linalg.generic
+// CATEGORY_TO_GENERIC-NOT: linalg.elementwise
+
+// GENERIC_TO_NAMED: linalg.exp
+// GENERIC_TO_NAMED-NOT: linalg.generic
>From 4f763e510a7205ff206516cab7283e8f3c69ec80 Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Sat, 19 Jul 2025 09:31:59 -0400
Subject: [PATCH 3/5] address comment
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 9 +--------
.../lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp | 2 +-
2 files changed, 2 insertions(+), 9 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 25b635726fbf5..707c5439cd3d3 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -94,7 +94,6 @@ def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
let description = [{
Convert a linalg op from one representation to another equivalent.
-
For example, a linalg named op `linalg.add` can also be written as an
category op `linalg.elementwise`, and can also be re-written as
a `linalg.generic`, giving the morphism:
@@ -116,22 +115,16 @@ def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
// named-op <--> category <--> generic
Option<"namedToCategory", "named-to-category", "bool", /*default=*/"false",
"convert named ops to category op e.g. `linalg.elementwise`">,
-
Option<"categoryToGeneric", "category-to-generic", "bool", /*default=*/"false",
"convert category ops e.g. `linalg.elementwise` to `linalg.generic`">,
-
Option<"namedToGeneric", "named-to-generic", "bool", /*default=*/"false",
"convert named ops e.g. `linalg.add` to `linalg.generic`">,
-
Option<"genericToCategory", "generic-to-category", "bool", /*default=*/"false",
"convert generic ops to category op e.g. `linalg.contraction`">,
-
Option<"categoryToNamed", "category-to-named", "bool", /*default=*/"false",
"convert category ops to equivalent named ops">,
-
Option<"genericToNamed", "generic-to-named", "bool", /*default=*/"false",
- "convert linalg.generic to equivalent named ops">
- ];
+ "convert linalg.generic to equivalent named ops"> ];
}
def LinalgGeneralizeNamedOpsPass : Pass<"linalg-generalize-named-ops"> {
diff --git a/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
index 81430b0432269..26eedc7ee31c7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
@@ -49,7 +49,7 @@ ElementwiseKind getKind(Operation *op) {
.Case([](TanhOp) { return ElementwiseKind::tanh; })
.Case([](ErfOp) { return ElementwiseKind::erf; })
.Default([&](Operation *op) {
- assert(false && "unexpected op");
+ llvm_unreachable("unhandled case in named to elementwise");
return ElementwiseKind::sub;
});
}
>From de9c5ce8e2ce944ca67775d8b1b4c2cbc5682915 Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Wed, 30 Jul 2025 05:29:09 -0400
Subject: [PATCH 4/5] address review comments Mostly renaming and restructuring
of code and tests.
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 17 ++++++-------
.../Dialect/Linalg/Transforms/MorphOps.cpp | 24 ++-----------------
.../Linalg/Transforms/NamedToElementwise.cpp | 1 +
...entwise.mlir => named-to-elementwise.mlir} | 18 ++++++++++++++
...ps.mlir => linalg-morph-category-ops.mlir} | 12 +---------
.../Linalg/linalg-morph-multi-step.mlir | 14 +++++++++++
6 files changed, 45 insertions(+), 41 deletions(-)
rename mlir/test/Dialect/Linalg/elementwise/{named_to_elementwise.mlir => named-to-elementwise.mlir} (67%)
rename mlir/test/Dialect/Linalg/{linalg-morph-ops.mlir => linalg-morph-category-ops.mlir} (60%)
create mode 100644 mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index 707c5439cd3d3..f23662930accc 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -100,12 +100,12 @@ def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
named-op <--> category_op (elementwise, contraction, ..) <--> generic
- Generic is a bigger set than named and category ops and so not all generics
- can be converted to single category-op or named-op. Similarly, category
- ops are bigger set than named ops.
+ Note that the set of `linalg.generic` subsumes named and category ops
+ and therefore not all `linalg.genric` can be converted to named or
+ category op. Similarly, catgory ops subsume named ops.
Note:
- Legacy converters (will be deprecated):
+ Legacy converters:
`--linalg-generalize-named-ops` is the path `named-op --> generic-op`
`--linalg-specialize-generic-ops` is the path `named-op <-- generic-op`
}];
@@ -113,16 +113,17 @@ def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
let options = [
// named-op <--> category <--> generic
+
+ // Lowering options
Option<"namedToCategory", "named-to-category", "bool", /*default=*/"false",
"convert named ops to category op e.g. `linalg.elementwise`">,
Option<"categoryToGeneric", "category-to-generic", "bool", /*default=*/"false",
"convert category ops e.g. `linalg.elementwise` to `linalg.generic`">,
Option<"namedToGeneric", "named-to-generic", "bool", /*default=*/"false",
"convert named ops e.g. `linalg.add` to `linalg.generic`">,
- Option<"genericToCategory", "generic-to-category", "bool", /*default=*/"false",
- "convert generic ops to category op e.g. `linalg.contraction`">,
- Option<"categoryToNamed", "category-to-named", "bool", /*default=*/"false",
- "convert category ops to equivalent named ops">,
+
+ // Lifting options
+ // TODOs: `generic-to-category`, `category-to-named`
Option<"genericToNamed", "generic-to-named", "bool", /*default=*/"false",
"convert linalg.generic to equivalent named ops"> ];
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
index c0e800be3917a..f261ccb1415fe 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -6,21 +6,8 @@
//
//===----------------------------------------------------------------------===//
//
-// This file implements conversions between:
-// named <--> category (elementwise, contraction, ..) <--> generic ops.
-//
-// For example, a named op such `linalg.add` can also be re-written as an
-// equivalent category op `linalg.elementwise` and also as a `linalg.generic`.
-//
-// Generic is a bigger set than named ops and so not all generics can be
-// converted to single category-op or named-op. Similarly, category-ops
-// are bigger in representational possiblities than named ops e.g.
-// `linalg.add` has no affine maps attached, but `linalg.elementwise` does.
-//
-// Note:
-// Legacy converters (will be deprecated):
-// `--linalg-generalize-named-ops` is the path `named-op --> generic-op`
-// `--linalg-specialize-generic-ops` is the path `named-op <-- generic-op`
+// This file implements conversions between linalg ops:
+// named <--> category (elementwise, contraction, ..) <--> generic.
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Complex/IR/Complex.h"
@@ -58,7 +45,6 @@ void LinalgMorphOpsPass::runOnOperation() {
// Lowering paths (named -> category -> generic)
if (namedToCategory) {
- // TODO: named -> contraction-op
populateLinalgNamedToElementwisePatterns(patterns);
}
if (namedToGeneric || categoryToGeneric) {
@@ -66,12 +52,6 @@ void LinalgMorphOpsPass::runOnOperation() {
}
// Lifting paths (named <- category <- generic)
- if (genericToCategory) {
- // TODO.
- }
- if (categoryToNamed) {
- // TODO: if there is a case for this.
- }
if (genericToNamed) {
populateLinalgGenericOpsSpecializationPatterns(patterns);
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
index 26eedc7ee31c7..00a076b6e9746 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/NamedToElementwise.cpp
@@ -75,6 +75,7 @@ struct NamedToElementwisePattern : public OpRewritePattern<NamedOpTy> {
void mlir::linalg::populateLinalgNamedToElementwisePatterns(
RewritePatternSet &patterns) {
+ patterns.add<NamedToElementwisePattern<SelectOp>>(patterns.getContext());
patterns.add<NamedToElementwisePattern<AddOp>>(patterns.getContext());
patterns.add<NamedToElementwisePattern<SubOp>>(patterns.getContext());
patterns.add<NamedToElementwisePattern<MulOp>>(patterns.getContext());
diff --git a/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir b/mlir/test/Dialect/Linalg/elementwise/named-to-elementwise.mlir
similarity index 67%
rename from mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
rename to mlir/test/Dialect/Linalg/elementwise/named-to-elementwise.mlir
index 0920cbbd6e15f..fa86160aa1f95 100644
--- a/mlir/test/Dialect/Linalg/elementwise/named_to_elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/elementwise/named-to-elementwise.mlir
@@ -36,3 +36,21 @@ func.func @sub(%A : memref<16x8xf32>, %B: memref<16x8xf32>, %C : memref<16x8xf32
linalg.sub ins(%A, %B : memref<16x8xf32>, memref<16x8xf32>) outs(%C : memref<16x8xf32>)
return
}
+
+// ----
+
+// CHECK: @ternary_select(%[[A:.+]]: tensor<4x8x16xi1>, %[[B:.+]]: tensor<4x8x16xf32>, %[[C:.+]]: tensor<4x8x16xf32>)
+// CHECK: %[[E:.+]] = tensor.empty() : tensor<4x8x16xf32>
+// CHECK: {{.*}} = linalg.elementwise
+// CHECK-SAME: kind=#linalg.elementwise_kind<select>
+// CHECK-SAME: ins(%[[A]], %[[B]], %[[C]] : tensor<4x8x16xi1>, tensor<4x8x16xf32>, tensor<4x8x16xf32>)
+// CHECK-SAME: outs(%[[E]] : tensor<4x8x16xf32>) -> tensor<4x8x16xf32>
+//
+func.func @ternary_select(%A: tensor<4x8x16xi1>, %B: tensor<4x8x16xf32>, %C: tensor<4x8x16xf32>)
+ -> tensor<4x8x16xf32> {
+ %empty = tensor.empty() : tensor<4x8x16xf32>
+ %select = linalg.select
+ ins(%A, %B, %C : tensor<4x8x16xi1>, tensor<4x8x16xf32>, tensor<4x8x16xf32>)
+ outs(%empty: tensor<4x8x16xf32>) -> tensor<4x8x16xf32>
+ return %select : tensor<4x8x16xf32>
+}
diff --git a/mlir/test/Dialect/Linalg/linalg-morph-ops.mlir b/mlir/test/Dialect/Linalg/linalg-morph-category-ops.mlir
similarity index 60%
rename from mlir/test/Dialect/Linalg/linalg-morph-ops.mlir
rename to mlir/test/Dialect/Linalg/linalg-morph-category-ops.mlir
index d15d29b4fd532..00602c4a36010 100644
--- a/mlir/test/Dialect/Linalg/linalg-morph-ops.mlir
+++ b/mlir/test/Dialect/Linalg/linalg-morph-category-ops.mlir
@@ -1,12 +1,8 @@
// Forward path `named -> category -> generic`
// RUN: mlir-opt %s -linalg-morph-ops=named-to-category | FileCheck %s --check-prefix=NAMED_TO_CATEGORY
-// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | FileCheck %s --check-prefix=NAMED_TO_GENERIC
+
// RUN: mlir-opt %s -linalg-morph-ops=named-to-category | \
// RUN: mlir-opt %s -linalg-morph-ops=category-to-generic | FileCheck %s --check-prefix=CATEGORY_TO_GENERIC
-//
-// Backward path `named <- category <- generic`
-// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | mlir-opt %s -linalg-morph-ops=generic-to-named | \
-// RUN: FileCheck %s --check-prefix=GENERIC_TO_NAMED
func.func @exp(%A : tensor<16x8xf32>, %B : tensor<16x8xf32>) -> tensor<16x8xf32> {
%exp = linalg.exp ins(%A : tensor<16x8xf32>) outs(%B : tensor<16x8xf32>) -> tensor<16x8xf32>
@@ -15,11 +11,5 @@ func.func @exp(%A : tensor<16x8xf32>, %B : tensor<16x8xf32>) -> tensor<16x8xf32
// NAMED_TO_CATEGORY: linalg.elementwise
// NAMED_TO_CATEGORY-NOT: linalg.exp
-// NAMED_TO_GENERIC: linalg.generic
-// NAMED_TO_GENERIC-NOT: linalg.exp
-
// CATEGORY_TO_GENERIC: linalg.generic
// CATEGORY_TO_GENERIC-NOT: linalg.elementwise
-
-// GENERIC_TO_NAMED: linalg.exp
-// GENERIC_TO_NAMED-NOT: linalg.generic
diff --git a/mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir b/mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir
new file mode 100644
index 0000000000000..3c40d1f2092c8
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | FileCheck %s --check-prefix=NAMED_TO_GENERIC
+// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | mlir-opt %s -linalg-morph-ops=generic-to-named | \
+// RUN: FileCheck %s --check-prefix=GENERIC_TO_NAMED
+
+func.func @exp(%A : tensor<16x8xf32>, %B : tensor<16x8xf32>) -> tensor<16x8xf32> {
+ %exp = linalg.exp ins(%A : tensor<16x8xf32>) outs(%B : tensor<16x8xf32>) -> tensor<16x8xf32>
+ return %exp : tensor<16x8xf32>
+}
+
+// NAMED_TO_GENERIC: linalg.generic
+// NAMED_TO_GENERIC-NOT: linalg.exp
+
+// GENERIC_TO_NAMED: linalg.exp
+// GENERIC_TO_NAMED-NOT: linalg.generic
>From aa5b3e17fbd1c86da09324f31fd6395e7dafbc5d Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Thu, 7 Aug 2025 06:19:06 -0400
Subject: [PATCH 5/5] review comment address
---
.../Linalg/elementwise/named-to-elementwise.mlir | 14 +++++++-------
.../Dialect/Linalg/linalg-morph-multi-step.mlir | 6 +++---
2 files changed, 10 insertions(+), 10 deletions(-)
diff --git a/mlir/test/Dialect/Linalg/elementwise/named-to-elementwise.mlir b/mlir/test/Dialect/Linalg/elementwise/named-to-elementwise.mlir
index fa86160aa1f95..2332b287ace8d 100644
--- a/mlir/test/Dialect/Linalg/elementwise/named-to-elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/elementwise/named-to-elementwise.mlir
@@ -26,15 +26,15 @@ func.func @add(%A : tensor<16x8xf32>, %B: tensor<16x8xf32>, %C : tensor<16x8xf32
// ----
-// CHECK: @sub(%[[A:.+]]: memref<16x8xf32>, %[[B:.+]]: memref<16x8xf32>, %[[C:.+]]: memref<16x8xf32>) {
-// CHECK: linalg.elementwise
+// CHECK: @sub(%[[A:.+]]: tensor<16x8xf32>, %[[B:.+]]: tensor<16x8xf32>, %[[C:.+]]: tensor<16x8xf32>) -> tensor<16x8xf32> {
+// CHECK: {{.*}} = linalg.elementwise
// CHECK-SAME: kind=#linalg.elementwise_kind<sub>
-// CHECK-SAME: ins(%[[A]], %[[B]] : memref<16x8xf32>, memref<16x8xf32>)
-// CHECK-SAME: outs(%[[C]] : memref<16x8xf32>)
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<16x8xf32>, tensor<16x8xf32>)
+// CHECK-SAME: outs(%[[C]] : tensor<16x8xf32>)
//
-func.func @sub(%A : memref<16x8xf32>, %B: memref<16x8xf32>, %C : memref<16x8xf32>) {
- linalg.sub ins(%A, %B : memref<16x8xf32>, memref<16x8xf32>) outs(%C : memref<16x8xf32>)
- return
+func.func @sub(%A : tensor<16x8xf32>, %B: tensor<16x8xf32>, %C : tensor<16x8xf32>) -> tensor<16x8xf32> {
+ %sub = linalg.sub ins(%A, %B : tensor<16x8xf32>, tensor<16x8xf32>) outs(%C : tensor<16x8xf32>) -> tensor<16x8xf32>
+ return %sub : tensor<16x8xf32>
}
// ----
diff --git a/mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir b/mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir
index 3c40d1f2092c8..ab50a44a37067 100644
--- a/mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir
+++ b/mlir/test/Dialect/Linalg/linalg-morph-multi-step.mlir
@@ -1,6 +1,6 @@
// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | FileCheck %s --check-prefix=NAMED_TO_GENERIC
// RUN: mlir-opt %s -linalg-morph-ops=named-to-generic | mlir-opt %s -linalg-morph-ops=generic-to-named | \
-// RUN: FileCheck %s --check-prefix=GENERIC_TO_NAMED
+// RUN: FileCheck %s --check-prefix=ROUND_TRIP
func.func @exp(%A : tensor<16x8xf32>, %B : tensor<16x8xf32>) -> tensor<16x8xf32> {
%exp = linalg.exp ins(%A : tensor<16x8xf32>) outs(%B : tensor<16x8xf32>) -> tensor<16x8xf32>
@@ -10,5 +10,5 @@ func.func @exp(%A : tensor<16x8xf32>, %B : tensor<16x8xf32>) -> tensor<16x8xf32
// NAMED_TO_GENERIC: linalg.generic
// NAMED_TO_GENERIC-NOT: linalg.exp
-// GENERIC_TO_NAMED: linalg.exp
-// GENERIC_TO_NAMED-NOT: linalg.generic
+// ROUND_TRIP: linalg.exp
+// ROUND_TRIP-NOT: linalg.generic
More information about the Mlir-commits
mailing list