[Mlir-commits] [mlir] [mlir][linalg] Don't assert on non-canonical operand order when specializing elementwise generics (PR #235340)
Stanislav Bardyuk
llvmlistbot at llvm.org
Sat Oct 24 11:48:45 PDT 2026
https://github.com/kodlan created https://github.com/llvm/llvm-project/pull/235340
specializeLinalgElementwise calls areBinOpsSwapped on every binary or ternary linalg.generic that passes the elementwise matchers. The helper only looks at the first value operand of the body op, and if that one is not block argument 0 it asserts that the body must be the exact swapped form. Since the matchers accept body ops with more operands than inputs (that is what makes reciprocal and square work for unary generics), a two-input generic with a body like arith.select %m, %c, %x, where %c is captured from outside, reaches the assert and mlir-opt aborts.
The same helper has a second gap for ternaries: it never checks that the mask is the first operand of the select. With all-i1 inputs, three operand permutations hit the same assert, and one (select %in_1, %in_0, %in) is quietly specialized to linalg.elementwise <select> with the inputs in the wrong order.
This changes the helper to check the whole mapping and return std::nullopt when the operands are neither in canonical nor in swapped order (or when a ternary does not use the mask as operand 0). The caller then fails the match and leaves the generic alone instead of asserting. Inputs that specialized correctly before still do, since all of them have the two value operands as the two block arguments in one of the two orders.
Hoisting the captured scalar of a binary select into a ternary linalg.elementwise <select> with a broadcast map would be a possible follow up, the same way unary generics already hoist a scalar operand.
Fixes #228854
>From 421a7ea060ebd842aafb7f7ed1723e4057b79b6d Mon Sep 17 00:00:00 2001
From: Stanislav Bardyuk <sbardyuk at google.com>
Date: Sat, 24 Oct 2026 18:11:37 +0000
Subject: [PATCH] [mlir][linalg] Don't assert on non-canonical operand order
when specializing elementwise generics
`specializeLinalgElementwise` calls `areBinOpsSwapped` on every binary or ternary `linalg.generic` that passes the elementwise matchers. The helper only looks at the first value operand of the body op, and if that one is not block argument 0 it asserts that the body must be the exact swapped form. Since the matchers accept body ops with more operands than inputs (that is what makes reciprocal and square work for unary generics), a two-input generic with a body like `arith.select %m, %c, %x`, where `%c` is captured from outside, reaches the assert and mlir-opt aborts.
The same helper has a second gap for ternaries: it never checks that the mask is the first operand of the select. With all-i1 inputs, three operand permutations hit the same assert, and one (`select %in_1, %in_0, %in`) is quietly specialized to `linalg.elementwise <select>` with the inputs in the wrong order.
This changes the helper to check the whole mapping and return `std::nullopt` when the operands are neither in canonical nor in swapped order (or when a ternary does not use the mask as operand 0). The caller then fails the match and leaves the generic alone instead of asserting. Inputs that specialized correctly before still do, since all of them have the two value operands as the two block arguments in one of the two orders.
Fixes #228854
---
.../Dialect/Linalg/Transforms/Specialize.cpp | 37 +++++---
.../Linalg/specialize-generic-ops.mlir | 92 +++++++++++++++++++
2 files changed, 115 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index 8fce91acf4b71f..c5b3c3e9588478 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -47,20 +47,23 @@ using namespace mlir::linalg;
// %0 = arith.subf %b, %a : f32
// linalg.yield %0: f32
// Former is linalg.sub(a,b), latter is linalg.sub(b,a).
-static bool areBinOpsSwapped(GenericOp genericOp, bool isTernary) {
+// Returns std::nullopt if the operands are not block arguments in either order.
+static std::optional<bool> areBinOpsSwapped(GenericOp genericOp,
+ bool isTernary) {
Block *body = genericOp.getBody();
Operation *op = &body->front();
- bool swapped = false;
- if (op->getOpOperand(0 + isTernary).get() !=
- body->getArgument(0 + isTernary)) {
- swapped = true;
- assert(op->getOpOperand(0 + isTernary).get() ==
- body->getArgument(1 + isTernary) &&
- op->getOpOperand(1 + isTernary).get() ==
- body->getArgument(0 + isTernary) &&
- "binary op uses just one block arg");
- }
- return swapped;
+ // The mask of a ternary op is not subject to swapping.
+ if (isTernary && op->getOperand(0) != body->getArgument(0))
+ return std::nullopt;
+ Value lhs = op->getOperand(0 + isTernary);
+ Value rhs = op->getOperand(1 + isTernary);
+ BlockArgument lhsArg = body->getArgument(0 + isTernary);
+ BlockArgument rhsArg = body->getArgument(1 + isTernary);
+ if (lhs == lhsArg && rhs == rhsArg)
+ return false;
+ if (lhs == rhsArg && rhs == lhsArg)
+ return true;
+ return std::nullopt;
}
// Given an elementwise single unary linalg generic op whose body operation is a
@@ -181,8 +184,14 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
Operation *op = &genericOp.getBody()->front();
// Detect variations from canonical forms.
- bool hasSwappedOperands =
- (isBinary || isTernary) && areBinOpsSwapped(genericOp, isTernary);
+ bool hasSwappedOperands = false;
+ if (isBinary || isTernary) {
+ std::optional<bool> swapped = areBinOpsSwapped(genericOp, isTernary);
+ if (!swapped)
+ return rewriter.notifyMatchFailure(
+ genericOp, "operands are not block arguments in either order");
+ hasSwappedOperands = *swapped;
+ }
int scalarOprIdx = -1;
bool hasScalarOperand = isUnary && op->getNumOperands() == 2 &&
findIndexOfScalarOperand(genericOp, scalarOprIdx);
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index da51685fc584f4..40e0253b54e2ca 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -797,6 +797,98 @@ func.func @negative_ternary_op_select(%C: tensor<?x?xi1>, %T: tensor<?x?xf32>,
// -----
+// One select operand is a scalar captured from outside the generic
+#map = affine_map<(d0) -> (d0)>
+func.func @negative_binary_op_select_captured_scalar(%A: tensor<?xi32>, %C: tensor<?xi1>,
+ %Out: tensor<?xi32>, %s: i32)
+ -> tensor<?xi32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel"]}
+ ins(%A, %C : tensor<?xi32>, tensor<?xi1>)
+ outs(%Out : tensor<?xi32>) {
+ ^bb0(%in: i32, %in_0: i1, %out: i32):
+ %v = arith.select %in_0, %s, %in : i32
+ linalg.yield %v : i32
+ } -> tensor<?xi32>
+ return %0 : tensor<?xi32>
+}
+
+// ALL-LABEL: negative_binary_op_select_captured_scalar
+// ALL: linalg.generic
+// ALL-NOT: linalg.elementwise <select>
+
+// -----
+
+// The select mask is a scalar captured from outside the generic
+#map = affine_map<(d0) -> (d0)>
+func.func @negative_binary_op_select_captured_mask(%T: tensor<?xi32>, %F: tensor<?xi32>,
+ %Out: tensor<?xi32>, %m: i1)
+ -> tensor<?xi32> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map],
+ iterator_types = ["parallel"]}
+ ins(%T, %F : tensor<?xi32>, tensor<?xi32>)
+ outs(%Out : tensor<?xi32>) {
+ ^bb0(%in: i32, %in_0: i32, %out: i32):
+ %v = arith.select %m, %in, %in_0 : i32
+ linalg.yield %v : i32
+ } -> tensor<?xi32>
+ return %0 : tensor<?xi32>
+}
+
+// ALL-LABEL: negative_binary_op_select_captured_mask
+// ALL: linalg.generic
+// ALL-NOT: linalg.elementwise <select>
+
+// -----
+
+// The mask is not the first select operand
+#map = affine_map<(d0) -> (d0)>
+func.func @negative_ternary_op_select_mask_not_first(%C: tensor<?xi1>, %T: tensor<?xi1>,
+ %F: tensor<?xi1>, %Out: tensor<?xi1>)
+ -> tensor<?xi1> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map, #map],
+ iterator_types = ["parallel"]}
+ ins(%C, %T, %F : tensor<?xi1>, tensor<?xi1>, tensor<?xi1>)
+ outs(%Out : tensor<?xi1>) {
+ ^bb0(%in: i1, %in_0: i1, %in_1: i1, %out: i1):
+ %v = arith.select %in_0, %in, %in_1 : i1
+ linalg.yield %v : i1
+ } -> tensor<?xi1>
+ return %0 : tensor<?xi1>
+}
+
+// ALL-LABEL: negative_ternary_op_select_mask_not_first
+// ALL: linalg.generic
+// ALL-NOT: linalg.elementwise <select>
+
+// -----
+
+// The select operands are a permutation that is neither canonical nor swapped
+#map = affine_map<(d0) -> (d0)>
+func.func @negative_ternary_op_select_permuted(%C: tensor<?xi1>, %T: tensor<?xi1>,
+ %F: tensor<?xi1>, %Out: tensor<?xi1>)
+ -> tensor<?xi1> {
+ %0 = linalg.generic
+ {indexing_maps = [#map, #map, #map, #map],
+ iterator_types = ["parallel"]}
+ ins(%C, %T, %F : tensor<?xi1>, tensor<?xi1>, tensor<?xi1>)
+ outs(%Out : tensor<?xi1>) {
+ ^bb0(%in: i1, %in_0: i1, %in_1: i1, %out: i1):
+ %v = arith.select %in_1, %in_0, %in : i1
+ linalg.yield %v : i1
+ } -> tensor<?xi1>
+ return %0 : tensor<?xi1>
+}
+
+// ALL-LABEL: negative_ternary_op_select_permuted
+// ALL: linalg.generic
+// ALL-NOT: linalg.elementwise <select>
+
+// -----
+
///----------------------------------------------------------------------------------------
/// Tests for linalg.matmul
///----------------------------------------------------------------------------------------
More information about the Mlir-commits
mailing list