[Mlir-commits] [mlir] [mlir][linalg] Don't assert on non-canonical operand order when specializing elementwise generics (PR #235340)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 24 11:49:19 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Stanislav Bardyuk (kodlan)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/235340.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp (+23-14)
- (modified) mlir/test/Dialect/Linalg/specialize-generic-ops.mlir (+92)
``````````diff
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
///----------------------------------------------------------------------------------------
``````````
</details>
https://github.com/llvm/llvm-project/pull/235340
More information about the Mlir-commits
mailing list