[Mlir-commits] [mlir] Add category-to-named specialization to linalg-morph-ops (PR #190116)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 16 04:57:19 PDT 2026
https://github.com/LekkalaSravya3 updated https://github.com/llvm/llvm-project/pull/190116
>From 7e13f1dc0f2ffed4b3dc98960e71fc1300715148 Mon Sep 17 00:00:00 2001
From: LekkalaSravya3 <lekkala.sravya at multicorewareinc.com>
Date: Thu, 2 Apr 2026 05:49:15 +0000
Subject: [PATCH 1/3] Add linalg-morph-ops category-to-named specialization
path
Signed-off-by: LekkalaSravya3 <lekkala.sravya at multicorewareinc.com>
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 5 +-
.../Dialect/Linalg/Transforms/Transforms.h | 4 +
.../Dialect/Linalg/Transforms/CMakeLists.txt | 1 +
.../Linalg/Transforms/CategoryToNamed.cpp | 160 +++++++++++++
.../Dialect/Linalg/Transforms/MorphOps.cpp | 5 +-
.../linalg-morph-category-to-named.mlir | 211 ++++++++++++++++++
6 files changed, 381 insertions(+), 5 deletions(-)
create mode 100644 mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
create mode 100644 mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index b873f260e7d92..db7b342373b2e 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -67,13 +67,14 @@ def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
"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`">,
-
+
// Specialization path is not guaranteed.
+ Option<"categoryToNamed", "category-to-named", "bool", /*default=*/"false",
+ "convert category ops to equivalent named ops where possible">,
Option<"genericToNamed", "generic-to-named", "bool", /*default=*/"false",
"convert linalg.generic to equivalent named ops">,
Option<"genericToCategory", "generic-to-category", "bool", /*default=*/"false",
"convert linalg.generic to equivalent category ops"> ];
- // TODOs: `category-to-named`
}
def LinalgGeneralizeNamedOpsPass : Pass<"linalg-generalize-named-ops">,
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 486ef75b76859..e86ba7fa59936 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1926,6 +1926,10 @@ void populateLinalgNamedToElementwisePatterns(RewritePatternSet &patterns);
/// `linalg.transform` into elementwise op map.
void populateLinalgFoldIntoElementwisePatterns(RewritePatternSet &patterns);
+/// Populates `patterns` that convert linalg category ops to equivalent
+/// named ops where possible.
+void populateLinalgCategoryToNamedPatterns(RewritePatternSet &patterns);
+
/// Linalg decompose convolutions patterns
/// Populates patterns to decompose high-D convolution ops into low-D ones.
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
index a2149478e4c2d..38ae99cf25a96 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
AllInterfaces.cpp
BubbleUpExtractSlice.cpp
BufferizableOpInterfaceImpl.cpp
+ CategoryToNamed.cpp
ConstantFold.cpp
ConvertToDestinationStyle.cpp
ConvertConv2DToImg2Col.cpp
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp b/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
new file mode 100644
index 0000000000000..f5c3f7374be59
--- /dev/null
+++ b/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
@@ -0,0 +1,160 @@
+//===- CategoryToNamed.cpp - convert linalg category ops into named 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 rewriting those linalg category ops that can be
+// represented by named ops, e.g. `linalg.elementwise<exp>` to `linalg.exp` or
+// `linalg.contract` to `linalg.matmul`.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Verifier.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+#define DEBUG_TYPE "linalg-category-to-named"
+
+namespace {
+
+template <typename NamedOpTy>
+static FailureOr<LinalgOp> replaceElementwiseOp(ElementwiseOp op,
+ PatternRewriter &rewriter) {
+ SmallVector<NamedAttribute> attrs;
+ attrs.push_back(rewriter.getNamedAttr("indexing_maps", op.getIndexingMaps()));
+
+ auto namedOp = NamedOpTy::create(rewriter, op.getLoc(), op.getDpsInputs(),
+ op.getDpsInits(), attrs);
+
+ {
+ ScopedDiagnosticHandler handler(op.getContext(), [](Diagnostic &) {});
+ if (failed(verify(namedOp.getOperation()))) {
+ rewriter.eraseOp(namedOp);
+ return rewriter.notifyMatchFailure(
+ op, "elementwise op does not satisfy named op constraints");
+ }
+ }
+
+ rewriter.replaceOp(op, namedOp->getResults());
+ return cast<LinalgOp>(namedOp.getOperation());
+}
+
+static FailureOr<LinalgOp> specializeElementwiseOp(ElementwiseOp op,
+ PatternRewriter &rewriter) {
+ switch (op.getKind()) {
+ case ElementwiseKind::select:
+ return replaceElementwiseOp<SelectOp>(op, rewriter);
+ case ElementwiseKind::add:
+ return replaceElementwiseOp<AddOp>(op, rewriter);
+ case ElementwiseKind::sub:
+ return replaceElementwiseOp<SubOp>(op, rewriter);
+ case ElementwiseKind::mul:
+ return replaceElementwiseOp<MulOp>(op, rewriter);
+ case ElementwiseKind::div:
+ return replaceElementwiseOp<DivOp>(op, rewriter);
+ case ElementwiseKind::div_unsigned:
+ return replaceElementwiseOp<DivUnsignedOp>(op, rewriter);
+ case ElementwiseKind::max_signed:
+ return replaceElementwiseOp<MaxOp>(op, rewriter);
+ case ElementwiseKind::min_signed:
+ return replaceElementwiseOp<MinOp>(op, rewriter);
+ case ElementwiseKind::max_unsigned:
+ case ElementwiseKind::min_unsigned:
+ break;
+ case ElementwiseKind::powf:
+ return replaceElementwiseOp<PowFOp>(op, rewriter);
+ case ElementwiseKind::exp:
+ return replaceElementwiseOp<ExpOp>(op, rewriter);
+ case ElementwiseKind::log:
+ return replaceElementwiseOp<LogOp>(op, rewriter);
+ case ElementwiseKind::abs:
+ return replaceElementwiseOp<AbsOp>(op, rewriter);
+ case ElementwiseKind::ceil:
+ return replaceElementwiseOp<CeilOp>(op, rewriter);
+ case ElementwiseKind::floor:
+ return replaceElementwiseOp<FloorOp>(op, rewriter);
+ case ElementwiseKind::negf:
+ return replaceElementwiseOp<NegFOp>(op, rewriter);
+ case ElementwiseKind::reciprocal:
+ return replaceElementwiseOp<ReciprocalOp>(op, rewriter);
+ case ElementwiseKind::round:
+ return replaceElementwiseOp<RoundOp>(op, rewriter);
+ case ElementwiseKind::sqrt:
+ return replaceElementwiseOp<SqrtOp>(op, rewriter);
+ case ElementwiseKind::rsqrt:
+ return replaceElementwiseOp<RsqrtOp>(op, rewriter);
+ case ElementwiseKind::square:
+ return replaceElementwiseOp<SquareOp>(op, rewriter);
+ case ElementwiseKind::tanh:
+ return replaceElementwiseOp<TanhOp>(op, rewriter);
+ case ElementwiseKind::erf:
+ return replaceElementwiseOp<ErfOp>(op, rewriter);
+ }
+
+ return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
+ diag << "unsupported elementwise kind for named specialization: "
+ << stringifyElementwiseKind(op.getKind());
+ });
+}
+
+struct ElementwiseToNamedPattern : public OpRewritePattern<ElementwiseOp> {
+ using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ElementwiseOp op,
+ PatternRewriter &rewriter) const override {
+ return succeeded(specializeElementwiseOp(op, rewriter)) ? success()
+ : failure();
+ }
+};
+
+struct ContractToNamedPattern : public OpRewritePattern<ContractOp> {
+ using OpRewritePattern<ContractOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ContractOp op,
+ PatternRewriter &rewriter) const override {
+ // Route through a cloned generic op so we can reuse the existing
+ // contraction-to-named specialization without mutating the original op
+ // on unsuccessful matches.
+ auto *clonedOp = rewriter.clone(*op.getOperation());
+ auto clonedLinalgOp = cast<LinalgOp>(clonedOp);
+
+ FailureOr<GenericOp> genericOp =
+ generalizeNamedOp(rewriter, clonedLinalgOp);
+ if (failed(genericOp)) {
+ rewriter.eraseOp(clonedOp);
+ return failure();
+ }
+
+ GenericOpSpecializationOptions options;
+ FailureOr<LinalgOp> namedOp =
+ specializeGenericOp(rewriter, *genericOp, options);
+ if (failed(namedOp)) {
+ rewriter.eraseOp(*genericOp);
+ return failure();
+ }
+
+ if (op->getNumResults() == 0) {
+ rewriter.eraseOp(op);
+ return success();
+ }
+
+ rewriter.replaceOp(op, (*namedOp)->getResults());
+ return success();
+ }
+};
+
+} // namespace
+
+void mlir::linalg::populateLinalgCategoryToNamedPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<ElementwiseToNamedPattern, ContractToNamedPattern>(
+ patterns.getContext());
+}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
index fee293647deda..2c7c320eff866 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -15,8 +15,6 @@
#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 {
@@ -55,7 +53,8 @@ void LinalgMorphOpsPass::runOnOperation() {
opts.emitCategoryOps = genericToCategory;
populateLinalgGenericOpsSpecializationPatterns(patterns, opts);
}
-
+ if (categoryToNamed)
+ populateLinalgCategoryToNamedPatterns(patterns);
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
signalPassFailure();
}
diff --git a/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir b/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
new file mode 100644
index 0000000000000..51a5e016bcde4
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
@@ -0,0 +1,211 @@
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=category-to-named | \
+// RUN: FileCheck %s
+
+func.func @elementwise_unary(%arg0: tensor<?x?xf32>,
+ %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.elementwise kind = #linalg.elementwise_kind<exp>
+ ins(%arg0 : tensor<?x?xf32>)
+ outs(%arg1 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: @elementwise_unary
+// CHECK-SAME: %[[IN:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>)
+// CHECK-NOT: linalg.elementwise
+// CHECK: linalg.exp
+// CHECK-SAME: ins(%[[IN]] : tensor<?x?xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+
+// -----
+
+func.func @elementwise_binary(%arg0: tensor<?x?xf32>,
+ %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.elementwise
+ kind = #linalg.elementwise_kind<powf>
+ ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: @elementwise_binary
+// CHECK-SAME: %[[LHS:.+]]: tensor<?x?xf32>, %[[RHS:.+]]: tensor<?x?xf32>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK-NOT: linalg.elementwise
+// CHECK: linalg.powf
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+
+// -----
+
+#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @contract_to_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+ %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.contract indexing_maps = [#map_a, #map_b, #map_c]
+ ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: @contract_to_matmul
+// CHECK-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK-NOT: linalg.contract
+// CHECK: linalg.matmul
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+
+// -----
+
+#cast_map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
+#cast_map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
+#cast_map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @contract_to_matmul_unsigned_cast(%arg0: tensor<16x8xi16>,
+ %arg1: tensor<8x32xi64>, %arg2: tensor<16x32xi32>) -> tensor<16x32xi32> {
+ %0 = linalg.contract indexing_maps = [#cast_map_a, #cast_map_b, #cast_map_c]
+ {cast = #linalg.type_fn<cast_unsigned>}
+ ins(%arg0, %arg1 : tensor<16x8xi16>, tensor<8x32xi64>)
+ outs(%arg2 : tensor<16x32xi32>) -> tensor<16x32xi32>
+ return %0 : tensor<16x32xi32>
+}
+
+// CHECK-LABEL: @contract_to_matmul_unsigned_cast
+// CHECK-SAME: %[[A:.+]]: tensor<16x8xi16>, %[[B:.+]]: tensor<8x32xi64>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<16x32xi32>) -> tensor<16x32xi32>
+// CHECK-NOT: linalg.contract
+// CHECK: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<16x8xi16>, tensor<8x32xi64>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<16x32xi32>) -> tensor<16x32xi32>
+
+// -----
+
+#map_ta = affine_map<(d0, d1, d2) -> (d2, d0)>
+#map_tb_base = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map_tc_base = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @contract_to_matmul_transpose_a(%arg0: tensor<8x16xf32>,
+ %arg1: tensor<8x32xf32>, %arg2: tensor<16x32xf32>) -> tensor<16x32xf32> {
+ %0 = linalg.contract indexing_maps = [#map_ta, #map_tb_base, #map_tc_base]
+ ins(%arg0, %arg1 : tensor<8x16xf32>, tensor<8x32xf32>)
+ outs(%arg2 : tensor<16x32xf32>) -> tensor<16x32xf32>
+ return %0 : tensor<16x32xf32>
+}
+
+// CHECK-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
+// CHECK-DAG: #[[$MAP_TB_BASE:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[$MAP_TC_BASE:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: @contract_to_matmul_transpose_a
+// CHECK-SAME: %[[A:.+]]: tensor<8x16xf32>, %[[B:.+]]: tensor<8x32xf32>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<16x32xf32>) -> tensor<16x32xf32>
+// CHECK-NOT: linalg.contract
+// CHECK: linalg.matmul
+// CHECK-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_TB_BASE]], #[[$MAP_TC_BASE]]]
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<8x16xf32>, tensor<8x32xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<16x32xf32>) -> tensor<16x32xf32>
+
+// -----
+
+#batch_map_a = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#batch_map_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#batch_map_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+func.func @contract_to_batch_matmul(%arg0: tensor<2x16x8xf32>,
+ %arg1: tensor<2x8x16xf32>, %arg2: tensor<2x16x16xf32>)
+ -> tensor<2x16x16xf32> {
+ %0 = linalg.contract indexing_maps = [#batch_map_a, #batch_map_b, #batch_map_c]
+ ins(%arg0, %arg1 : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
+ outs(%arg2 : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
+ return %0 : tensor<2x16x16xf32>
+}
+
+// CHECK-LABEL: @contract_to_batch_matmul
+// CHECK-SAME: %[[A:.+]]: tensor<2x16x8xf32>, %[[B:.+]]: tensor<2x8x16xf32>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
+// CHECK-NOT: linalg.contract
+// CHECK: linalg.batch_matmul
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
+
+// -----
+
+#batch_map_ta = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
+#batch_map_tb = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#batch_map_tc = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+func.func @contract_to_batch_matmul_transpose_a(%arg0: tensor<2x8x16xf32>,
+ %arg1: tensor<2x8x32xf32>, %arg2: tensor<2x16x32xf32>)
+ -> tensor<2x16x32xf32> {
+ %0 = linalg.contract indexing_maps = [#batch_map_ta, #batch_map_tb, #batch_map_tc]
+ ins(%arg0, %arg1 : tensor<2x8x16xf32>, tensor<2x8x32xf32>)
+ outs(%arg2 : tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
+ return %0 : tensor<2x16x32xf32>
+}
+
+// CHECK-DAG: #[[$BATCH_MAP_TA:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
+// CHECK-DAG: #[[$BATCH_MAP_TB:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[$BATCH_MAP_TC:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+// CHECK-LABEL: @contract_to_batch_matmul_transpose_a
+// CHECK-SAME: %[[A:.+]]: tensor<2x8x16xf32>, %[[B:.+]]: tensor<2x8x32xf32>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
+// CHECK-NOT: linalg.contract
+// CHECK: linalg.batch_matmul
+// CHECK-SAME: indexing_maps = [#[[$BATCH_MAP_TA]], #[[$BATCH_MAP_TB]], #[[$BATCH_MAP_TC]]]
+// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<2x8x16xf32>, tensor<2x8x32xf32>)
+// CHECK-SAME: outs(%[[OUT]] : tensor<2x16x32xf32>) -> tensor<2x16x32xf32>
+
+// -----
+
+#non_identity_batch_a = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
+#non_identity_batch_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#non_identity_batch_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+// this stays as `linalg.contract` because the batch dimension is
+// not in identity position across the operand/result maps. Named
+// `linalg.batch_matmul` does not model such non-identity batch permutations.
+func.func @contract_non_identity_batch(%arg0: tensor<4x2x8xf32>,
+ %arg1: tensor<2x8x16xf32>, %arg2: tensor<2x4x16xf32>)
+ -> tensor<2x4x16xf32> {
+ %0 = linalg.contract indexing_maps = [#non_identity_batch_a, #non_identity_batch_b, #non_identity_batch_c]
+ ins(%arg0, %arg1 : tensor<4x2x8xf32>, tensor<2x8x16xf32>)
+ outs(%arg2 : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
+ return %0 : tensor<2x4x16xf32>
+}
+
+// CHECK-DAG: #[[$NON_ID_BATCH_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
+// CHECK-DAG: #[[$NON_ID_BATCH_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[$NON_ID_BATCH_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+
+// CHECK-LABEL: @contract_non_identity_batch
+// CHECK-SAME: %[[A:.+]]: tensor<4x2x8xf32>, %[[B:.+]]: tensor<2x8x16xf32>,
+// CHECK-SAME: %[[OUT:.+]]: tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#[[$NON_ID_BATCH_A]], #[[$NON_ID_BATCH_B]], #[[$NON_ID_BATCH_C]]]
+// CHECK-NOT: linalg.batch_matmul
+
+// -----
+
+#map_d = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map_e = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)>
+#map_f = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
+// this stays as `linalg.contract` because it has two reduction
+// dimensions. Named matmul-like ops require exactly one M dim, one N dim, and
+// one K dim.
+func.func @contract_multi_reduction(%arg0: tensor<10x20x30xf32>,
+ %arg1: tensor<30x20x40xf32>,
+ %arg2: tensor<10x40xf32>) -> tensor<10x40xf32> {
+ %0 = linalg.contract indexing_maps = [#map_d, #map_e, #map_f]
+ ins(%arg0, %arg1 : tensor<10x20x30xf32>, tensor<30x20x40xf32>)
+ outs(%arg2 : tensor<10x40xf32>) -> tensor<10x40xf32>
+ return %0 : tensor<10x40xf32>
+}
+
+// CHECK-LABEL: @contract_multi_reduction
+// CHECK-NOT: linalg.matmul
+// CHECK: linalg.contract
+// CHECK-SAME: indexing_maps = [#{{.+}}, #{{.+}}, #{{.+}}]
>From e1b6604463be647f4f343270f1c81ac55cd42835 Mon Sep 17 00:00:00 2001
From: LekkalaSravya3 <lekkala.sravya at multicorewareinc.com>
Date: Mon, 6 Apr 2026 10:56:26 +0000
Subject: [PATCH 2/3] Handle contract op conversion and legality checks
Signed-off-by: LekkalaSravya3 <lekkala.sravya at multicorewareinc.com>
---
.../Linalg/Transforms/CategoryToNamed.cpp | 186 +++++++++++++-----
.../Dialect/Linalg/Transforms/MorphOps.cpp | 1 +
.../linalg-morph-category-to-named.mlir | 24 +++
3 files changed, 164 insertions(+), 47 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp b/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
index f5c3f7374be59..b9f7c515c9e54 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
@@ -6,7 +6,7 @@
//
//===----------------------------------------------------------------------===//
//
-// This file implements rewriting those linalg category ops that can be
+// This file implements rewriting the subset of linalg category ops that can be
// represented by named ops, e.g. `linalg.elementwise<exp>` to `linalg.exp` or
// `linalg.contract` to `linalg.matmul`.
//
@@ -25,30 +25,71 @@ using namespace mlir::linalg;
namespace {
+enum class IndexMatchResult { Match = 0, Transposed, Mismatch };
+
+static IndexMatchResult matchOperandMap(AffineMap map, unsigned rowDimIdx,
+ unsigned expectedPosOfRowDim,
+ unsigned expectedPosOfColDim) {
+ auto exprOfRowDim = map.getResults()[rowDimIdx];
+ auto exprOfColDim = map.getResults()[rowDimIdx + 1];
+
+ if (exprOfRowDim.getKind() != AffineExprKind::DimId ||
+ exprOfColDim.getKind() != AffineExprKind::DimId)
+ return IndexMatchResult::Mismatch;
+
+ auto posRowDim = cast<AffineDimExpr>(exprOfRowDim).getPosition();
+ auto posColDim = cast<AffineDimExpr>(exprOfColDim).getPosition();
+
+ if (expectedPosOfRowDim == posRowDim && expectedPosOfColDim == posColDim)
+ return IndexMatchResult::Match;
+ if (expectedPosOfRowDim == posColDim && expectedPosOfColDim == posRowDim)
+ return IndexMatchResult::Transposed;
+ return IndexMatchResult::Mismatch;
+}
+
template <typename NamedOpTy>
-static FailureOr<LinalgOp> replaceElementwiseOp(ElementwiseOp op,
- PatternRewriter &rewriter) {
- SmallVector<NamedAttribute> attrs;
- attrs.push_back(rewriter.getNamedAttr("indexing_maps", op.getIndexingMaps()));
+static LogicalResult replaceElementwiseOp(ElementwiseOp op,
+ PatternRewriter &rewriter) {
+ unsigned numMaps = op.getNumDpsInputs() + op.getNumDpsInits();
+ SmallVector<AffineMap> defaultMaps = ElementwiseOp::getDefaultIndexingMaps(
+ numMaps, op.getResultRank(), op.getContext());
+ if (!llvm::equal(op.getIndexingMapsArray(), defaultMaps))
+ return rewriter.notifyMatchFailure(
+ op, "named elementwise ops require default indexing maps");
+
+ auto namedOp =
+ NamedOpTy::create(rewriter, op.getLoc(), op.getDpsInputs(),
+ op.getDpsInits(), ArrayRef<NamedAttribute>{});
+ rewriter.replaceOp(op, namedOp->getResults());
+ return success();
+}
- auto namedOp = NamedOpTy::create(rewriter, op.getLoc(), op.getDpsInputs(),
- op.getDpsInits(), attrs);
+template <typename NamedOpTy>
+static LogicalResult replaceContractOp(ContractOp op, PatternRewriter &rewriter,
+ ArrayRef<AffineMap> indexingMaps) {
+ SmallVector<NamedAttribute> attrs;
+ if (op.getCast() == TypeFn::cast_unsigned) {
+ attrs.push_back(rewriter.getNamedAttr(
+ "cast", TypeFnAttr::get(rewriter.getContext(), op.getCast())));
+ }
- {
- ScopedDiagnosticHandler handler(op.getContext(), [](Diagnostic &) {});
- if (failed(verify(namedOp.getOperation()))) {
- rewriter.eraseOp(namedOp);
- return rewriter.notifyMatchFailure(
- op, "elementwise op does not satisfy named op constraints");
- }
+ SmallVector<Attribute> indexingMapsAttrVal =
+ llvm::map_to_vector(indexingMaps, [](AffineMap map) -> Attribute {
+ return AffineMapAttr::get(map);
+ });
+ auto indexingMapsAttr = rewriter.getArrayAttr(indexingMapsAttrVal);
+ if (!NamedOpTy::isDefaultIndexingMaps(indexingMapsAttr)) {
+ attrs.push_back(rewriter.getNamedAttr("indexing_maps", indexingMapsAttr));
}
+ auto namedOp = NamedOpTy::create(rewriter, op.getLoc(), op.getDpsInputs(),
+ op.getDpsInits(), attrs);
rewriter.replaceOp(op, namedOp->getResults());
- return cast<LinalgOp>(namedOp.getOperation());
+ return success();
}
-static FailureOr<LinalgOp> specializeElementwiseOp(ElementwiseOp op,
- PatternRewriter &rewriter) {
+static LogicalResult specializeElementwiseOp(ElementwiseOp op,
+ PatternRewriter &rewriter) {
switch (op.getKind()) {
case ElementwiseKind::select:
return replaceElementwiseOp<SelectOp>(op, rewriter);
@@ -68,6 +109,8 @@ static FailureOr<LinalgOp> specializeElementwiseOp(ElementwiseOp op,
return replaceElementwiseOp<MinOp>(op, rewriter);
case ElementwiseKind::max_unsigned:
case ElementwiseKind::min_unsigned:
+ // There are no named unsigned max/min ops yet, so these category ops
+ // cannot currently be represented as named ops.
break;
case ElementwiseKind::powf:
return replaceElementwiseOp<PowFOp>(op, rewriter);
@@ -105,13 +148,89 @@ static FailureOr<LinalgOp> specializeElementwiseOp(ElementwiseOp op,
});
}
+static LogicalResult specializeContractOp(ContractOp op,
+ PatternRewriter &rewriter) {
+ if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+ return failure();
+
+ // Named matmul-like ops only admit permutation-style indexing maps.
+ auto indexingMaps = op.getIndexingMapsArray();
+ if (llvm::any_of(indexingMaps,
+ [](AffineMap map) { return !map.isProjectedPermutation(); }))
+ return failure();
+
+ // Restrict the contraction to the matmul family shape: one M, one N, one K,
+ // plus an optional prefix of batch dimensions.
+ auto res = inferContractionDims(op);
+ if (failed(res))
+ return failure();
+ auto dims = *res;
+ if (dims.m.size() != 1 || dims.n.size() != 1 || dims.k.size() != 1)
+ return failure();
+
+ if (llvm::any_of(indexingMaps, [&dims](AffineMap map) {
+ return map.getResults().size() !=
+ dims.batch.size() + 2 /* any two of {m, n, k} */;
+ }))
+ return failure();
+
+ auto numBatchDims = dims.batch.size();
+ if (indexingMaps[0].getNumDims() != numBatchDims + 3)
+ return failure();
+
+ // Batch dimensions must stay in canonical order for named batch_matmul.
+ if (numBatchDims && llvm::any_of(indexingMaps, [numBatchDims](AffineMap map) {
+ for (unsigned i = 0; i < numBatchDims; ++i) {
+ auto expr = map.getResults()[i];
+ if (expr.getKind() != AffineExprKind::DimId ||
+ cast<AffineDimExpr>(expr).getPosition() != i)
+ return true;
+ }
+ return false;
+ }))
+ return failure();
+
+ auto a = matchOperandMap(indexingMaps[0], numBatchDims, dims.m[0], dims.k[0]);
+ auto b = matchOperandMap(indexingMaps[1], numBatchDims, dims.k[0], dims.n[0]);
+ auto c = matchOperandMap(indexingMaps[2], numBatchDims, dims.m[0], dims.n[0]);
+ if (llvm::is_contained({a, b, c}, IndexMatchResult::Mismatch))
+ return failure();
+
+ auto *ctx = op.getContext();
+ unsigned numLoopDims = numBatchDims + 3;
+ unsigned mIdx = numBatchDims;
+ unsigned nIdx = mIdx + 1;
+ unsigned kIdx = mIdx + 2;
+
+ // Rebuild indexing maps in the named op's canonical loop order while
+ // preserving legal operand transpositions.
+ auto makeMap = [&](IndexMatchResult match, unsigned rowIdx, unsigned colIdx) {
+ SmallVector<unsigned> tensorDims;
+ for (unsigned i = 0; i < numBatchDims; ++i)
+ tensorDims.push_back(i);
+ if (match == IndexMatchResult::Transposed)
+ llvm::append_values(tensorDims, colIdx, rowIdx);
+ else
+ llvm::append_values(tensorDims, rowIdx, colIdx);
+ return AffineMap::getMultiDimMapWithTargets(numLoopDims, tensorDims, ctx);
+ };
+
+ auto mapA = makeMap(a, mIdx, kIdx);
+ auto mapB = makeMap(b, kIdx, nIdx);
+ auto mapC = makeMap(c, mIdx, nIdx);
+ SmallVector<AffineMap, 3> namedOpMaps = {mapA, mapB, mapC};
+
+ if (numBatchDims)
+ return replaceContractOp<BatchMatmulOp>(op, rewriter, namedOpMaps);
+ return replaceContractOp<MatmulOp>(op, rewriter, namedOpMaps);
+}
+
struct ElementwiseToNamedPattern : public OpRewritePattern<ElementwiseOp> {
using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
LogicalResult matchAndRewrite(ElementwiseOp op,
PatternRewriter &rewriter) const override {
- return succeeded(specializeElementwiseOp(op, rewriter)) ? success()
- : failure();
+ return specializeElementwiseOp(op, rewriter);
}
};
@@ -120,34 +239,7 @@ struct ContractToNamedPattern : public OpRewritePattern<ContractOp> {
LogicalResult matchAndRewrite(ContractOp op,
PatternRewriter &rewriter) const override {
- // Route through a cloned generic op so we can reuse the existing
- // contraction-to-named specialization without mutating the original op
- // on unsuccessful matches.
- auto *clonedOp = rewriter.clone(*op.getOperation());
- auto clonedLinalgOp = cast<LinalgOp>(clonedOp);
-
- FailureOr<GenericOp> genericOp =
- generalizeNamedOp(rewriter, clonedLinalgOp);
- if (failed(genericOp)) {
- rewriter.eraseOp(clonedOp);
- return failure();
- }
-
- GenericOpSpecializationOptions options;
- FailureOr<LinalgOp> namedOp =
- specializeGenericOp(rewriter, *genericOp, options);
- if (failed(namedOp)) {
- rewriter.eraseOp(*genericOp);
- return failure();
- }
-
- if (op->getNumResults() == 0) {
- rewriter.eraseOp(op);
- return success();
- }
-
- rewriter.replaceOp(op, (*namedOp)->getResults());
- return success();
+ return specializeContractOp(op, rewriter);
}
};
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
index 2c7c320eff866..d34f89738dd5a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -53,6 +53,7 @@ void LinalgMorphOpsPass::runOnOperation() {
opts.emitCategoryOps = genericToCategory;
populateLinalgGenericOpsSpecializationPatterns(patterns, opts);
}
+
if (categoryToNamed)
populateLinalgCategoryToNamedPatterns(patterns);
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
diff --git a/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir b/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
index 51a5e016bcde4..0c5b26c9f4523 100644
--- a/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
+++ b/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
@@ -37,6 +37,30 @@ func.func @elementwise_binary(%arg0: tensor<?x?xf32>,
// -----
+#elt_map_in = affine_map<(d0, d1) -> (d1, d0)>
+#elt_map_out = affine_map<(d0, d1) -> (d0, d1)>
+
+// This stays as `linalg.elementwise` because named elementwise ops do not
+// support user-defined indexing maps.
+func.func @elementwise_non_default_maps(%arg0: tensor<?x?xf32>,
+ %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.elementwise
+ kind = #linalg.elementwise_kind<exp>
+ indexing_maps = [#elt_map_in, #elt_map_out]
+ ins(%arg0 : tensor<?x?xf32>)
+ outs(%arg1 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CHECK-DAG: #[[$ELT_MAP_IN:.+]] = affine_map<(d0, d1) -> (d1, d0)>
+// CHECK-DAG: #[[$ELT_MAP_OUT:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: @elementwise_non_default_maps
+// CHECK: linalg.elementwise
+// CHECK-SAME: indexing_maps = [#[[$ELT_MAP_IN]], #[[$ELT_MAP_OUT]]]
+// CHECK-NOT: linalg.exp
+
+// -----
+
#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
>From 7bf4a2258f9374b3e734adfef62798af6cf8e47e Mon Sep 17 00:00:00 2001
From: LekkalaSravya3 <lekkala.sravya at multicorewareinc.com>
Date: Tue, 16 Jun 2026 11:55:59 +0000
Subject: [PATCH 3/3] extract the shared matmul-like index-map into shared
utility and address the review feedback
Signed-off-by: LekkalaSravya3 <lekkala.sravya at multicorewareinc.com>
---
mlir/include/mlir/Dialect/Linalg/Passes.td | 2 +-
.../Dialect/Linalg/Transforms/Transforms.h | 6 +
.../Linalg/Transforms/CategoryToNamed.cpp | 109 ++-----------
.../Dialect/Linalg/Transforms/MorphOps.cpp | 1 +
.../Dialect/Linalg/Transforms/Specialize.cpp | 145 +-----------------
.../Dialect/Linalg/Transforms/Transforms.cpp | 103 +++++++++++++
.../linalg-morph-category-to-named.mlir | 12 +-
...rip-morphism-linalg-category-to-named.mlir | 46 ++++++
8 files changed, 179 insertions(+), 245 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-to-named.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td
index db7b342373b2e..d4ef5059f3469 100644
--- a/mlir/include/mlir/Dialect/Linalg/Passes.td
+++ b/mlir/include/mlir/Dialect/Linalg/Passes.td
@@ -70,7 +70,7 @@ def LinalgMorphOpsPass : Pass<"linalg-morph-ops"> {
// Specialization path is not guaranteed.
Option<"categoryToNamed", "category-to-named", "bool", /*default=*/"false",
- "convert category ops to equivalent named ops where possible">,
+ "convert category ops to equivalent named ops">,
Option<"genericToNamed", "generic-to-named", "bool", /*default=*/"false",
"convert linalg.generic to equivalent named ops">,
Option<"genericToCategory", "generic-to-category", "bool", /*default=*/"false",
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index e86ba7fa59936..026560af155f1 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -45,6 +45,12 @@ enum class WinogradConv2DFmr : uint32_t;
/// Return vector::CombiningKind for the given op.
std::optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
+/// Return the canonical matmul-like indexing maps for `linalgOp` if it has the
+/// operand count, projected-permutation maps, and contraction dimensions needed
+/// to represent it as `linalg.matmul` or `linalg.batch_matmul`.
+FailureOr<SmallVector<AffineMap, 3>>
+inferMatmulLikeIndexingMaps(LinalgOp linalgOp);
+
//===----------------------------------------------------------------------===//
// Bufferization-related transforms.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp b/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
index b9f7c515c9e54..d2babf94be8c2 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/CategoryToNamed.cpp
@@ -23,30 +23,6 @@ using namespace mlir::linalg;
#define DEBUG_TYPE "linalg-category-to-named"
-namespace {
-
-enum class IndexMatchResult { Match = 0, Transposed, Mismatch };
-
-static IndexMatchResult matchOperandMap(AffineMap map, unsigned rowDimIdx,
- unsigned expectedPosOfRowDim,
- unsigned expectedPosOfColDim) {
- auto exprOfRowDim = map.getResults()[rowDimIdx];
- auto exprOfColDim = map.getResults()[rowDimIdx + 1];
-
- if (exprOfRowDim.getKind() != AffineExprKind::DimId ||
- exprOfColDim.getKind() != AffineExprKind::DimId)
- return IndexMatchResult::Mismatch;
-
- auto posRowDim = cast<AffineDimExpr>(exprOfRowDim).getPosition();
- auto posColDim = cast<AffineDimExpr>(exprOfColDim).getPosition();
-
- if (expectedPosOfRowDim == posRowDim && expectedPosOfColDim == posColDim)
- return IndexMatchResult::Match;
- if (expectedPosOfRowDim == posColDim && expectedPosOfColDim == posRowDim)
- return IndexMatchResult::Transposed;
- return IndexMatchResult::Mismatch;
-}
-
template <typename NamedOpTy>
static LogicalResult replaceElementwiseOp(ElementwiseOp op,
PatternRewriter &rewriter) {
@@ -73,11 +49,7 @@ static LogicalResult replaceContractOp(ContractOp op, PatternRewriter &rewriter,
"cast", TypeFnAttr::get(rewriter.getContext(), op.getCast())));
}
- SmallVector<Attribute> indexingMapsAttrVal =
- llvm::map_to_vector(indexingMaps, [](AffineMap map) -> Attribute {
- return AffineMapAttr::get(map);
- });
- auto indexingMapsAttr = rewriter.getArrayAttr(indexingMapsAttrVal);
+ auto indexingMapsAttr = rewriter.getAffineMapArrayAttr(indexingMaps);
if (!NamedOpTy::isDefaultIndexingMaps(indexingMapsAttr)) {
attrs.push_back(rewriter.getNamedAttr("indexing_maps", indexingMapsAttr));
}
@@ -150,81 +122,18 @@ static LogicalResult specializeElementwiseOp(ElementwiseOp op,
static LogicalResult specializeContractOp(ContractOp op,
PatternRewriter &rewriter) {
- if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
- return failure();
-
- // Named matmul-like ops only admit permutation-style indexing maps.
- auto indexingMaps = op.getIndexingMapsArray();
- if (llvm::any_of(indexingMaps,
- [](AffineMap map) { return !map.isProjectedPermutation(); }))
+ FailureOr<SmallVector<AffineMap, 3>> namedOpMaps =
+ inferMatmulLikeIndexingMaps(op);
+ if (failed(namedOpMaps))
return failure();
- // Restrict the contraction to the matmul family shape: one M, one N, one K,
- // plus an optional prefix of batch dimensions.
- auto res = inferContractionDims(op);
- if (failed(res))
- return failure();
- auto dims = *res;
- if (dims.m.size() != 1 || dims.n.size() != 1 || dims.k.size() != 1)
- return failure();
-
- if (llvm::any_of(indexingMaps, [&dims](AffineMap map) {
- return map.getResults().size() !=
- dims.batch.size() + 2 /* any two of {m, n, k} */;
- }))
- return failure();
-
- auto numBatchDims = dims.batch.size();
- if (indexingMaps[0].getNumDims() != numBatchDims + 3)
- return failure();
-
- // Batch dimensions must stay in canonical order for named batch_matmul.
- if (numBatchDims && llvm::any_of(indexingMaps, [numBatchDims](AffineMap map) {
- for (unsigned i = 0; i < numBatchDims; ++i) {
- auto expr = map.getResults()[i];
- if (expr.getKind() != AffineExprKind::DimId ||
- cast<AffineDimExpr>(expr).getPosition() != i)
- return true;
- }
- return false;
- }))
- return failure();
-
- auto a = matchOperandMap(indexingMaps[0], numBatchDims, dims.m[0], dims.k[0]);
- auto b = matchOperandMap(indexingMaps[1], numBatchDims, dims.k[0], dims.n[0]);
- auto c = matchOperandMap(indexingMaps[2], numBatchDims, dims.m[0], dims.n[0]);
- if (llvm::is_contained({a, b, c}, IndexMatchResult::Mismatch))
- return failure();
-
- auto *ctx = op.getContext();
- unsigned numLoopDims = numBatchDims + 3;
- unsigned mIdx = numBatchDims;
- unsigned nIdx = mIdx + 1;
- unsigned kIdx = mIdx + 2;
-
- // Rebuild indexing maps in the named op's canonical loop order while
- // preserving legal operand transpositions.
- auto makeMap = [&](IndexMatchResult match, unsigned rowIdx, unsigned colIdx) {
- SmallVector<unsigned> tensorDims;
- for (unsigned i = 0; i < numBatchDims; ++i)
- tensorDims.push_back(i);
- if (match == IndexMatchResult::Transposed)
- llvm::append_values(tensorDims, colIdx, rowIdx);
- else
- llvm::append_values(tensorDims, rowIdx, colIdx);
- return AffineMap::getMultiDimMapWithTargets(numLoopDims, tensorDims, ctx);
- };
-
- auto mapA = makeMap(a, mIdx, kIdx);
- auto mapB = makeMap(b, kIdx, nIdx);
- auto mapC = makeMap(c, mIdx, nIdx);
- SmallVector<AffineMap, 3> namedOpMaps = {mapA, mapB, mapC};
-
- if (numBatchDims)
- return replaceContractOp<BatchMatmulOp>(op, rewriter, namedOpMaps);
- return replaceContractOp<MatmulOp>(op, rewriter, namedOpMaps);
+ if ((*namedOpMaps)[0].getNumDims() > 3)
+ return replaceContractOp<BatchMatmulOp>(op, rewriter, *namedOpMaps);
+ return replaceContractOp<MatmulOp>(op, rewriter, *namedOpMaps);
}
+namespace {
+
struct ElementwiseToNamedPattern : public OpRewritePattern<ElementwiseOp> {
using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
index d34f89738dd5a..e1319214a108b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MorphOps.cpp
@@ -56,6 +56,7 @@ void LinalgMorphOpsPass::runOnOperation() {
if (categoryToNamed)
populateLinalgCategoryToNamedPatterns(patterns);
+
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
signalPassFailure();
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index 60b18fb2e8d93..7698be8e46fae 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -92,45 +92,6 @@ static bool areBinOpsSwapped(GenericOp genericOp) {
// It is not possible to represent above as named op.
// e.g. linalg.batch_matmul(%A, %B : tensor<20x20x20xf32>, ...) is
// not the same as linalg.generic above.
-namespace {
-enum class IndexMatchResult {
- Match = 0, // identity map.
- Transposed, // transposed map.
- Mismatch // none of the above.
-};
-
-// Checks whether the input Affine `map` contains two consecutive dims that
-// can be interpreted as accessing a 2D matrix. It is assumed that the row
-// column dimension are adjacent axis (in this order) and start at
-// `rowDimIdx` in the input map.
-//
-// e.g. consider A matrix in `C[M,N] = A[M,K] * B[K,N]`. We will check
-// whether the map of A is identity (match), transposed, or something
-// completely different (mis-match). Similar for B and C.
-static IndexMatchResult matchOperandMap(AffineMap map, unsigned rowDimIdx,
- unsigned expectedPosOfRowDim,
- unsigned expectedPosOfColDim) {
- // Get the matrix multiply indices. They are past the batch indices.
- auto exprOfRowDim = map.getResults()[rowDimIdx];
- auto exprOfColDim = map.getResults()[rowDimIdx + 1];
-
- // They should be pure dimension ids.
- if (exprOfRowDim.getKind() != AffineExprKind::DimId ||
- exprOfColDim.getKind() != AffineExprKind::DimId)
- return IndexMatchResult::Mismatch;
-
- auto posRowDim = cast<AffineDimExpr>(exprOfRowDim).getPosition();
- auto posColDim = cast<AffineDimExpr>(exprOfColDim).getPosition();
-
- if (expectedPosOfRowDim == posRowDim && expectedPosOfColDim == posColDim)
- return IndexMatchResult::Match;
-
- if (expectedPosOfRowDim == posColDim && expectedPosOfColDim == posRowDim)
- return IndexMatchResult::Transposed;
-
- return IndexMatchResult::Mismatch;
-}
-
// Replaces genericOp with `NamedOpTy` op, supplied as a template arg.
// All the variants expressed as pseudo regular expression:
// `linalg.{batch_}?matmul` have same number of ins/out, so it's easy to
@@ -152,12 +113,8 @@ static LinalgOp replaceWithMatmulVariant(RewriterBase &rewriter, GenericOp op,
// Set the original generic's maps to preserve operand indexing semantics like
// transposition.
- SmallVector<Attribute, 3> indexingMapsAttrVal =
- llvm::map_to_vector(indexingMaps, [](AffineMap map) -> Attribute {
- return AffineMapAttr::get(map);
- });
auto indexingMapsAttr = rewriter.getNamedAttr(
- "indexing_maps", rewriter.getArrayAttr(indexingMapsAttrVal));
+ "indexing_maps", rewriter.getAffineMapArrayAttr(indexingMaps));
attributes.push_back(indexingMapsAttr);
LinalgOp namedOp = rewriter.replaceOpWithNewOp<NamedOpTy>(
@@ -252,104 +209,18 @@ static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
return replaceWithMatmulVariant<ContractOp>(
rewriter, genericOp, castTy, genericOp.getIndexingMapsArray());
- // Further checks for named variants.
- //
- // Linalg generic contraction can be across multiple axis e.g.
- // ```
- // linalg.generic
- // {indexing_maps = [affine_map<(m, n, k1, k2) -> (m, k1, k2)>,
- // affine_map<(m, n, k1, k2) -> (k2, k1, n)>,
- // affine_map<(m, n, k1, k2) -> (m, n)>],
- // 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>
- // ```
- // In above contraction, there are two reduction dimensions {k1, k2}
- // and although a valid linalg contraction, it is not a named-op
- // matrix multiply kind. Therefore, reject multi-dim reduction.
- auto res = inferContractionDims(genericOp);
- if (!succeeded(res))
- return failure();
- auto dims = *res;
- if (dims.m.size() != 1 || dims.n.size() != 1 || dims.k.size() != 1)
+ FailureOr<SmallVector<AffineMap, 3>> namedOpMaps =
+ inferMatmulLikeIndexingMaps(genericOp);
+ if (failed(namedOpMaps))
return failure();
- // Check rank of operands
- auto indexingMaps = genericOp.getIndexingMapsArray();
- if (llvm::any_of(indexingMaps, [&dims](AffineMap m) {
- return m.getResults().size() !=
- dims.batch.size() + 2 /* any two of {m,n,k} */;
- }))
- return failure();
-
- auto numOfBatchDims = dims.batch.size();
- if (indexingMaps[0].getNumDims() != numOfBatchDims + 3)
- return failure();
-
- if (numOfBatchDims) {
- // Each operand in a linalg generic contraction could express different
- // permutations for its batch dimension. But for named op it must be
- // identity since separate maps are not specified.
- if (llvm::any_of(indexingMaps, [numOfBatchDims](AffineMap m) {
- for (unsigned i = 0; i < numOfBatchDims; ++i) {
- auto expr = m.getResults()[i];
- if (expr.getKind() != AffineExprKind::DimId ||
- cast<AffineDimExpr>(expr).getPosition() != i)
- return true;
- }
- return false;
- }))
- return failure();
- }
-
- auto a =
- matchOperandMap(indexingMaps[0], numOfBatchDims, dims.m[0], dims.k[0]);
- auto b =
- matchOperandMap(indexingMaps[1], numOfBatchDims, dims.k[0], dims.n[0]);
- auto c =
- matchOperandMap(indexingMaps[2], numOfBatchDims, dims.m[0], dims.n[0]);
-
- if (llvm::is_contained({a, b, c}, IndexMatchResult::Mismatch))
- return failure();
-
- // Build indexing maps for the named op in its canonical dimension ordering
- auto *ctx = genericOp.getContext();
- unsigned numLoopDims = numOfBatchDims + 3;
- unsigned mIdx = numOfBatchDims;
- unsigned nIdx = mIdx + 1;
- unsigned kIdx = mIdx + 2;
-
- // TODO: add support for indexing_maps with broadcasts.
- auto makeMap = [&](IndexMatchResult match, unsigned rowIdx, unsigned colIdx) {
- SmallVector<unsigned> tensorDims;
- for (unsigned i = 0; i < numOfBatchDims; ++i)
- tensorDims.push_back(i);
- if (match == IndexMatchResult::Transposed)
- llvm::append_values(tensorDims, colIdx, rowIdx);
- else
- llvm::append_values(tensorDims, rowIdx, colIdx);
- return AffineMap::getMultiDimMapWithTargets(numLoopDims, tensorDims, ctx);
- };
-
- auto mapA = makeMap(a, mIdx, kIdx);
- auto mapB = makeMap(b, kIdx, nIdx);
- auto mapC = makeMap(c, mIdx, nIdx);
-
- SmallVector<AffineMap> namedOpMaps = {mapA, mapB, mapC};
-
// Codegen the different matmul variants.
- if (numOfBatchDims) {
+ if ((*namedOpMaps)[0].getNumDims() > 3) {
return replaceWithMatmulVariant<BatchMatmulOp>(rewriter, genericOp, castTy,
- namedOpMaps);
+ *namedOpMaps);
}
return replaceWithMatmulVariant<MatmulOp>(rewriter, genericOp, castTy,
- namedOpMaps);
+ *namedOpMaps);
}
/// Utility to specialize a `genericOp` with a convolution op of type `ConvOpTy`
@@ -447,8 +318,6 @@ static FailureOr<LinalgOp> specializeLinalgConvolutions(RewriterBase &rewriter,
return failure();
}
-} // namespace
-
//===----------------------------------------------------------------------===//
// Categorize linalg generic to named op where possible.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 260e36fb47f04..2923935cf226e 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -43,6 +43,109 @@
using namespace mlir;
using namespace mlir::linalg;
+namespace {
+enum class IndexMatchResult {
+ Match = 0, // identity map.
+ Transposed, // transposed map.
+ Mismatch // none of the above.
+};
+} // namespace
+
+// Checks whether the input Affine `map` contains two consecutive dims that
+// can be interpreted as accessing a 2D matrix. It is assumed that the row and
+// column dimensions are adjacent axes (in this order) and start at `rowDimIdx`
+// in the input map.
+static IndexMatchResult matchOperandMap(AffineMap map, unsigned rowDimIdx,
+ unsigned expectedPosOfRowDim,
+ unsigned expectedPosOfColDim) {
+ auto exprOfRowDim = map.getResults()[rowDimIdx];
+ auto exprOfColDim = map.getResults()[rowDimIdx + 1];
+
+ if (exprOfRowDim.getKind() != AffineExprKind::DimId ||
+ exprOfColDim.getKind() != AffineExprKind::DimId)
+ return IndexMatchResult::Mismatch;
+
+ auto posRowDim = cast<AffineDimExpr>(exprOfRowDim).getPosition();
+ auto posColDim = cast<AffineDimExpr>(exprOfColDim).getPosition();
+
+ if (expectedPosOfRowDim == posRowDim && expectedPosOfColDim == posColDim)
+ return IndexMatchResult::Match;
+
+ if (expectedPosOfRowDim == posColDim && expectedPosOfColDim == posRowDim)
+ return IndexMatchResult::Transposed;
+
+ return IndexMatchResult::Mismatch;
+}
+
+FailureOr<SmallVector<AffineMap, 3>>
+mlir::linalg::inferMatmulLikeIndexingMaps(LinalgOp linalgOp) {
+ if (linalgOp.getNumDpsInputs() != 2 || linalgOp.getNumDpsInits() != 1)
+ return failure();
+
+ auto indexingMaps = linalgOp.getIndexingMapsArray();
+ if (llvm::any_of(indexingMaps,
+ [](AffineMap map) { return !map.isProjectedPermutation(); }))
+ return failure();
+
+ auto res = inferContractionDims(linalgOp);
+ if (failed(res))
+ return failure();
+ auto dims = *res;
+ if (dims.m.size() != 1 || dims.n.size() != 1 || dims.k.size() != 1)
+ return failure();
+
+ if (llvm::any_of(indexingMaps, [&dims](AffineMap map) {
+ return map.getResults().size() !=
+ dims.batch.size() + 2 /* any two of {m, n, k} */;
+ }))
+ return failure();
+
+ unsigned numBatchDims = dims.batch.size();
+ if (indexingMaps[0].getNumDims() != numBatchDims + 3)
+ return failure();
+
+ // Batch dimensions must stay in canonical order for named batch_matmul.
+ if (numBatchDims && llvm::any_of(indexingMaps, [numBatchDims](AffineMap map) {
+ for (unsigned i = 0; i < numBatchDims; ++i) {
+ auto expr = map.getResults()[i];
+ if (expr.getKind() != AffineExprKind::DimId ||
+ cast<AffineDimExpr>(expr).getPosition() != i)
+ return true;
+ }
+ return false;
+ }))
+ return failure();
+
+ auto a = matchOperandMap(indexingMaps[0], numBatchDims, dims.m[0], dims.k[0]);
+ auto b = matchOperandMap(indexingMaps[1], numBatchDims, dims.k[0], dims.n[0]);
+ auto c = matchOperandMap(indexingMaps[2], numBatchDims, dims.m[0], dims.n[0]);
+ if (llvm::is_contained({a, b, c}, IndexMatchResult::Mismatch))
+ return failure();
+
+ auto *ctx = linalgOp.getContext();
+ unsigned numLoopDims = numBatchDims + 3;
+ unsigned mIdx = numBatchDims;
+ unsigned nIdx = mIdx + 1;
+ unsigned kIdx = mIdx + 2;
+
+ // Rebuild indexing maps in the named op's canonical loop order while
+ // preserving legal operand transpositions.
+ auto makeMap = [&](IndexMatchResult match, unsigned rowIdx, unsigned colIdx) {
+ SmallVector<unsigned> tensorDims;
+ for (unsigned i = 0; i < numBatchDims; ++i)
+ tensorDims.push_back(i);
+ if (match == IndexMatchResult::Transposed)
+ llvm::append_values(tensorDims, colIdx, rowIdx);
+ else
+ llvm::append_values(tensorDims, rowIdx, colIdx);
+ return AffineMap::getMultiDimMapWithTargets(numLoopDims, tensorDims, ctx);
+ };
+
+ return SmallVector<AffineMap, 3>{makeMap(a, mIdx, kIdx),
+ makeMap(b, kIdx, nIdx),
+ makeMap(c, mIdx, nIdx)};
+}
+
//===----------------------------------------------------------------------===//
// Transformations exposed as functional-style API calls.
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir b/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
index 0c5b26c9f4523..8e0a894cfca37 100644
--- a/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
+++ b/mlir/test/Dialect/Linalg/linalg-morph-category-to-named.mlir
@@ -42,7 +42,7 @@ func.func @elementwise_binary(%arg0: tensor<?x?xf32>,
// This stays as `linalg.elementwise` because named elementwise ops do not
// support user-defined indexing maps.
-func.func @elementwise_non_default_maps(%arg0: tensor<?x?xf32>,
+func.func @negative_elementwise_non_default_maps(%arg0: tensor<?x?xf32>,
%arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.elementwise
kind = #linalg.elementwise_kind<exp>
@@ -54,7 +54,7 @@ func.func @elementwise_non_default_maps(%arg0: tensor<?x?xf32>,
// CHECK-DAG: #[[$ELT_MAP_IN:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK-DAG: #[[$ELT_MAP_OUT:.+]] = affine_map<(d0, d1) -> (d0, d1)>
-// CHECK-LABEL: @elementwise_non_default_maps
+// CHECK-LABEL: @negative_elementwise_non_default_maps
// CHECK: linalg.elementwise
// CHECK-SAME: indexing_maps = [#[[$ELT_MAP_IN]], #[[$ELT_MAP_OUT]]]
// CHECK-NOT: linalg.exp
@@ -191,7 +191,7 @@ func.func @contract_to_batch_matmul_transpose_a(%arg0: tensor<2x8x16xf32>,
// this stays as `linalg.contract` because the batch dimension is
// not in identity position across the operand/result maps. Named
// `linalg.batch_matmul` does not model such non-identity batch permutations.
-func.func @contract_non_identity_batch(%arg0: tensor<4x2x8xf32>,
+func.func @negative_contract_non_identity_batch(%arg0: tensor<4x2x8xf32>,
%arg1: tensor<2x8x16xf32>, %arg2: tensor<2x4x16xf32>)
-> tensor<2x4x16xf32> {
%0 = linalg.contract indexing_maps = [#non_identity_batch_a, #non_identity_batch_b, #non_identity_batch_c]
@@ -204,7 +204,7 @@ func.func @contract_non_identity_batch(%arg0: tensor<4x2x8xf32>,
// CHECK-DAG: #[[$NON_ID_BATCH_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[$NON_ID_BATCH_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
-// CHECK-LABEL: @contract_non_identity_batch
+// CHECK-LABEL: @negative_contract_non_identity_batch
// CHECK-SAME: %[[A:.+]]: tensor<4x2x8xf32>, %[[B:.+]]: tensor<2x8x16xf32>,
// CHECK-SAME: %[[OUT:.+]]: tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
// CHECK: linalg.contract
@@ -220,7 +220,7 @@ func.func @contract_non_identity_batch(%arg0: tensor<4x2x8xf32>,
// this stays as `linalg.contract` because it has two reduction
// dimensions. Named matmul-like ops require exactly one M dim, one N dim, and
// one K dim.
-func.func @contract_multi_reduction(%arg0: tensor<10x20x30xf32>,
+func.func @negative_contract_multi_reduction(%arg0: tensor<10x20x30xf32>,
%arg1: tensor<30x20x40xf32>,
%arg2: tensor<10x40xf32>) -> tensor<10x40xf32> {
%0 = linalg.contract indexing_maps = [#map_d, #map_e, #map_f]
@@ -229,7 +229,7 @@ func.func @contract_multi_reduction(%arg0: tensor<10x20x30xf32>,
return %0 : tensor<10x40xf32>
}
-// CHECK-LABEL: @contract_multi_reduction
+// CHECK-LABEL: @negative_contract_multi_reduction
// CHECK-NOT: linalg.matmul
// CHECK: linalg.contract
// CHECK-SAME: indexing_maps = [#{{.+}}, #{{.+}}, #{{.+}}]
diff --git a/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-to-named.mlir b/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-to-named.mlir
new file mode 100644
index 0000000000000..50e9e178186e3
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-to-named.mlir
@@ -0,0 +1,46 @@
+// The following test examples of linalg named/category ops lowered to the other
+// form and then lifted back up.
+
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=category-to-named \
+// RUN: | mlir-opt -split-input-file -linalg-morph-ops=named-to-category \
+// RUN: | FileCheck %s --check-prefix=CATEGORY
+
+// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=named-to-category \
+// RUN: | mlir-opt -split-input-file -linalg-morph-ops=category-to-named \
+// RUN: | FileCheck %s --check-prefix=NAMED
+
+func.func @elementwise_exp(%arg0: tensor<?x?xf32>,
+ %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.elementwise kind = #linalg.elementwise_kind<exp>
+ ins(%arg0 : tensor<?x?xf32>)
+ outs(%arg1 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CATEGORY-LABEL: @elementwise_exp
+// CATEGORY-NOT: linalg.exp
+// CATEGORY: linalg.elementwise
+// CATEGORY-SAME: kind=#linalg.elementwise_kind<exp>
+
+// NAMED-LABEL: @elementwise_exp
+// NAMED-NOT: linalg.elementwise
+// NAMED: linalg.exp
+
+// -----
+
+func.func @named_add(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+ %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %0 = linalg.add
+ ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// CATEGORY-LABEL: @named_add
+// CATEGORY-NOT: linalg.add
+// CATEGORY: linalg.elementwise
+// CATEGORY-SAME: kind=#linalg.elementwise_kind<add>
+
+// NAMED-LABEL: @named_add
+// NAMED-NOT: linalg.elementwise
+// NAMED: linalg.add
More information about the Mlir-commits
mailing list