From a005f0c80a8dfc579eec4a3bcb7c84cf1acb07d2 Mon Sep 17 00:00:00 2001 From: Javed Absar 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(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 +struct NamedToElementwisePattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(NamedOpTy op, + PatternRewriter &rewriter) const override { + SmallVector 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(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>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(patterns.getContext()); + patterns.add>(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 +// 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 +// 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 +// 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 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 { + + 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 { 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 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 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 { void mlir::linalg::populateLinalgNamedToElementwisePatterns( RewritePatternSet &patterns) { + patterns.add>(patterns.getContext()); patterns.add>(patterns.getContext()); patterns.add>(patterns.getContext()); patterns.add>(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