[Mlir-commits] [mlir] [mlir][scf] Propagate if conditions to captured select uses (PR #235433)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 25 05:21:24 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: meinie (meinie0826)
<details>
<summary>Changes</summary>
## Summary
Simplify branch-local uses of `arith.select` when an enclosing `scf.if` uses the same scalar condition. A use in the then region takes the select's true operand, and a use in the else region takes its false operand. The rule also covers uses nested in those regions and preserves uses outside the if.
For example:
```python
%x = arith.select %p, %a, %b : i32
scf.if %p {
func.call @<!-- -->consume(%x) : (i32) -> ()
} else {
func.call @<!-- -->consume(%x) : (i32) -> ()
}
```
The calls can use `%a` and `%b`, respectively. When no other uses remain, ordinary dead-code elimination removes the select. If the select also has outside uses, it remains for those uses.
## Implementation
Extend SCF condition propagation to inspect operands used in each branch and replace select results whose condition is the if's condition. The condition must be the same SSA value; a select on an unrelated scalar condition or an elementwise vector condition is not refined.
No computation, call, load, store, allocation, or region is moved. Memref uses select the same reference on the executed path; the rewrite does not copy data or coalesce allocations. This is a canonicalization rule, with no target-specific or memory-planning option.
## Tests
Regression tests cover then/else and then-only uses, repeated operands, outside uses, nested if/loop regions, memref stores, scalar selects with vector results, negated if conditions, and unrelated scalar/vector conditions.
Built upstream `cced08fcc92db61044c0e32ef35444fa48246345` with assertions enabled using system Clang 14.
- The regression test fails on the unmodified upstream build and passes with this change (12 cases).
- SCF, Arith, and the canonicalize/CSE/SCCP transform tests: 80 passed, 1 expected failure, no unexpected failures.
- Additional canonicalization and structured SCCP tests: 13 passed.
- `clang-format` on the modified pattern and `git diff --check` pass.
These are local regression-test results, not upstream CI or measured SRAM/runtime improvements.
---
Full diff: https://github.com/llvm/llvm-project/pull/235433.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+23)
- (added) mlir/test/Dialect/SCF/condition-propagation-select.mlir (+275)
``````````diff
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index ae548cf1d2e10b..3ab2cce1627df9 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2330,6 +2330,7 @@ struct ConditionPropagation : public OpRewritePattern<IfOp> {
return failure();
bool changed = false;
+ bool hasCapturedSelect = false;
mlir::Type i1Ty = rewriter.getI1Type();
// These variables serve to prevent creating duplicate constants
@@ -2365,10 +2366,32 @@ struct ConditionPropagation : public OpRewritePattern<IfOp> {
break;
}
case Parent::None:
+ if (isa<arith::SelectOp>(use.getOwner()) && use.getOperandNumber() == 0)
+ hasCapturedSelect = true;
break;
}
}
+ // A select captured by a branch can be defined outside the if, where its
+ // condition is not constant. Replace its branch-local uses with the value
+ // selected on that path, without changing uses outside the branch.
+ auto propagateSelectUses = [&](Region ®ion, bool condition) {
+ visitUsedValuesDefinedAbove(region, [&](OpOperand *use) {
+ auto selectOp = use->get().getDefiningOp<arith::SelectOp>();
+ if (!selectOp || selectOp.getCondition() != op.getCondition())
+ return;
+ Value selected =
+ condition ? selectOp.getTrueValue() : selectOp.getFalseValue();
+ rewriter.modifyOpInPlace(use->getOwner(),
+ [&]() { use->set(selected); });
+ changed = true;
+ });
+ };
+ if (hasCapturedSelect) {
+ propagateSelectUses(op.getThenRegion(), true);
+ propagateSelectUses(op.getElseRegion(), false);
+ }
+
return success(changed);
}
};
diff --git a/mlir/test/Dialect/SCF/condition-propagation-select.mlir b/mlir/test/Dialect/SCF/condition-propagation-select.mlir
new file mode 100644
index 00000000000000..122062d26d43f9
--- /dev/null
+++ b/mlir/test/Dialect/SCF/condition-propagation-select.mlir
@@ -0,0 +1,275 @@
+// RUN: mlir-opt %s -canonicalize='test-convergence' -split-input-file | FileCheck %s
+
+// Branch-local uses can be simplified even when the select must remain for an
+// outside use. Repeated operands of the same operation are all replaced.
+// CHECK-LABEL: func.func @captured_select(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : i32
+// CHECK: scf.if %[[P]] {
+// CHECK: func.call @consume_pair(%[[A]], %[[A]])
+// CHECK: } else {
+// CHECK: func.call @consume_pair(%[[B]], %[[B]])
+// CHECK: return %[[X]] : i32
+func.func @captured_select(%p: i1, %a: i32, %b: i32) -> i32 {
+ %x = arith.select %p, %a, %b : i32
+ scf.if %p {
+ func.call @consume_pair(%x, %x) : (i32, i32) -> ()
+ } else {
+ func.call @consume_pair(%x, %x) : (i32, i32) -> ()
+ }
+ return %x : i32
+}
+func.func private @consume_pair(i32, i32)
+
+// -----
+
+// CHECK-LABEL: func.func @then_only(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32
+// CHECK-NOT: arith.select
+// CHECK: scf.if %[[P]] {
+// CHECK: func.call @consume(%[[A]])
+// CHECK-NOT: arith.select
+// CHECK: return
+func.func @then_only(%p: i1, %a: i32, %b: i32) {
+ %x = arith.select %p, %a, %b : i32
+ scf.if %p {
+ func.call @consume(%x) : (i32) -> ()
+ }
+ return
+}
+func.func private @consume(i32)
+
+// -----
+
+// Selecting a memref does not copy its contents. Stores stay in their original
+// branches and use the selected destination; outside uses keep the select.
+// CHECK-LABEL: func.func @memref_destination(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: memref<16xi32>, %[[B:.*]]: memref<16xi32>
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : memref<16xi32>
+// CHECK: scf.if %[[P]] {
+// CHECK: memref.store %{{.*}}, %[[A]][%{{.*}}]
+// CHECK: } else {
+// CHECK: memref.store %{{.*}}, %[[B]][%{{.*}}]
+// CHECK: return %[[X]] : memref<16xi32>
+func.func @memref_destination(%p: i1, %a: memref<16xi32>, %b: memref<16xi32>) -> memref<16xi32> {
+ %c0 = arith.constant 0 : index
+ %v = arith.constant 42 : i32
+ %x = arith.select %p, %a, %b : memref<16xi32>
+ scf.if %p {
+ memref.store %v, %x[%c0] : memref<16xi32>
+ } else {
+ memref.store %v, %x[%c0] : memref<16xi32>
+ }
+ return %x : memref<16xi32>
+}
+
+// -----
+
+// The outer branch's condition also applies to uses in nested regions.
+// CHECK-LABEL: func.func @nested_regions(
+// CHECK-SAME: %[[P:.*]]: i1, %[[Q:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32
+// CHECK-NOT: arith.select
+// CHECK: scf.if %[[P]] {
+// CHECK: scf.if %[[Q]] {
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: func.call @consume(%[[B]])
+func.func @nested_regions(%p: i1, %q: i1, %a: i32, %b: i32) {
+ %x = arith.select %p, %a, %b : i32
+ scf.if %p {
+ scf.if %q {
+ func.call @consume(%x) : (i32) -> ()
+ } else {
+ func.call @consume(%x) : (i32) -> ()
+ }
+ } else {
+ func.call @consume(%x) : (i32) -> ()
+ }
+ return
+}
+func.func private @consume(i32)
+
+// -----
+
+// CHECK-LABEL: func.func @loop_region(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32, %[[N:.*]]: index
+// CHECK-NOT: arith.select
+// CHECK: scf.if %[[P]] {
+// CHECK: scf.for
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: func.call @consume(%[[B]])
+func.func @loop_region(%p: i1, %a: i32, %b: i32, %n: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %x = arith.select %p, %a, %b : i32
+ scf.if %p {
+ scf.for %i = %c0 to %n step %c1 {
+ func.call @consume(%x) : (i32) -> ()
+ }
+ } else {
+ func.call @consume(%x) : (i32) -> ()
+ }
+ return
+}
+func.func private @consume(i32)
+
+// -----
+
+// A branch on an unrelated condition cannot refine the selected value.
+// CHECK-LABEL: func.func @unrelated_condition(
+// CHECK-SAME: %[[P:.*]]: i1, %[[Q:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : i32
+// CHECK: scf.if %[[Q]] {
+// CHECK: func.call @consume(%[[X]])
+// CHECK: } else {
+// CHECK: func.call @consume(%[[X]])
+func.func @unrelated_condition(%p: i1, %q: i1, %a: i32, %b: i32) {
+ %x = arith.select %p, %a, %b : i32
+ scf.if %q {
+ func.call @consume(%x) : (i32) -> ()
+ } else {
+ func.call @consume(%x) : (i32) -> ()
+ }
+ return
+}
+func.func private @consume(i32)
+
+// -----
+
+// Merely using the branch condition as a select arm is not sufficient.
+// CHECK-LABEL: func.func @condition_as_select_arm(
+// CHECK-SAME: %[[P:.*]]: i1, %[[Q:.*]]: i1, %[[B:.*]]: i1
+// CHECK: %[[X:.*]] = arith.select %[[Q]], %[[P]], %[[B]] : i1
+// CHECK: scf.if %[[P]] {
+// CHECK: func.call @consume_i1(%[[X]])
+// CHECK: } else {
+// CHECK: func.call @consume_i1(%[[X]])
+func.func @condition_as_select_arm(%p: i1, %q: i1, %b: i1) {
+ %x = arith.select %q, %p, %b : i1
+ scf.if %p {
+ func.call @consume_i1(%x) : (i1) -> ()
+ } else {
+ func.call @consume_i1(%x) : (i1) -> ()
+ }
+ return
+}
+func.func private @consume_i1(i1)
+
+// -----
+
+// Canonicalization of a negated if condition must preserve which arm is used.
+// CHECK-LABEL: func.func @negated_condition(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32
+// CHECK-NOT: arith.select
+// CHECK: scf.if %[[P]] {
+// CHECK: func.call @consume_pair(%[[A]], %[[B]])
+// CHECK: } else {
+// CHECK: func.call @consume_pair(%[[B]], %[[A]])
+func.func @negated_condition(%p: i1, %a: i32, %b: i32) {
+ %true = arith.constant true
+ %not_p = arith.xori %p, %true : i1
+ %x = arith.select %p, %a, %b : i32
+ scf.if %not_p {
+ func.call @consume_pair(%x, %a) : (i32, i32) -> ()
+ } else {
+ func.call @consume_pair(%x, %b) : (i32, i32) -> ()
+ }
+ return
+}
+func.func private @consume_pair(i32, i32)
+
+// -----
+
+// A scalar branch condition does not refine an elementwise vector condition.
+// CHECK-LABEL: func.func @vector_condition(
+// CHECK-SAME: %[[P:.*]]: i1, %[[MASK:.*]]: vector<4xi1>, %[[A:.*]]: vector<4xi32>, %[[B:.*]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[MASK]], %[[A]], %[[B]]
+// CHECK: scf.if %[[P]] {
+// CHECK: func.call @consume_vector(%[[X]])
+// CHECK: } else {
+// CHECK: func.call @consume_vector(%[[X]])
+func.func @vector_condition(%p: i1, %mask: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>) {
+ %x = arith.select %mask, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %p {
+ func.call @consume_vector(%x) : (vector<4xi32>) -> ()
+ } else {
+ func.call @consume_vector(%x) : (vector<4xi32>) -> ()
+ }
+ return
+}
+func.func private @consume_vector(vector<4xi32>)
+
+// -----
+
+// A scalar select can select an entire vector without changing its elements.
+// CHECK-LABEL: func.func @scalar_condition_vector_result(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: vector<4xi32>, %[[B:.*]]: vector<4xi32>
+// CHECK-NOT: arith.select
+// CHECK: scf.if %[[P]] {
+// CHECK: func.call @consume_vector(%[[A]])
+// CHECK: } else {
+// CHECK: func.call @consume_vector(%[[B]])
+func.func @scalar_condition_vector_result(%p: i1, %a: vector<4xi32>, %b: vector<4xi32>) {
+ %x = arith.select %p, %a, %b : vector<4xi32>
+ scf.if %p {
+ func.call @consume_vector(%x) : (vector<4xi32>) -> ()
+ } else {
+ func.call @consume_vector(%x) : (vector<4xi32>) -> ()
+ }
+ return
+}
+func.func private @consume_vector(vector<4xi32>)
+
+// -----
+
+// An isolated region receives the selected value as an operand. Its block
+// argument remains local to that region; it must not capture either select arm.
+// CHECK-LABEL: func.func @isolated_region_operand(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: index, %[[B:.*]]: index
+// CHECK-NOT: arith.select
+// CHECK: scf.if %[[P]] {
+// CHECK: test.isolated_region %[[A]] {
+// CHECK: "test.consume"(%[[A]]) : (index) -> ()
+// CHECK: } else {
+// CHECK: test.isolated_region %[[B]] {
+// CHECK: "test.consume"(%[[B]]) : (index) -> ()
+func.func @isolated_region_operand(%p: i1, %a: index, %b: index) {
+ %x = arith.select %p, %a, %b : index
+ scf.if %p {
+ test.isolated_region %x {
+ "test.consume"(%x) : (index) -> ()
+ }
+ } else {
+ test.isolated_region %x {
+ "test.consume"(%x) : (index) -> ()
+ }
+ }
+ return
+}
+
+// -----
+
+// Correlated selects may form a chain. Both branch-local operands are refined,
+// while the result returned outside the if still needs a select.
+// CHECK-LABEL: func.func @select_chain(
+// CHECK-SAME: %[[P:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32, %[[C:.*]]: i32
+// CHECK: %[[RESULT:.*]] = arith.select %[[P]], %[[C]], %[[B]] : i32
+// CHECK: scf.if %[[P]] {
+// CHECK: func.call @consume_pair(%[[A]], %[[C]])
+// CHECK: } else {
+// CHECK: func.call @consume_pair(%[[B]], %[[B]])
+// CHECK: return %[[RESULT]] : i32
+func.func @select_chain(%p: i1, %a: i32, %b: i32, %c: i32) -> i32 {
+ %x = arith.select %p, %a, %b : i32
+ %y = arith.select %p, %c, %x : i32
+ scf.if %p {
+ func.call @consume_pair(%x, %y) : (i32, i32) -> ()
+ } else {
+ func.call @consume_pair(%x, %y) : (i32, i32) -> ()
+ }
+ return %y : i32
+}
+func.func private @consume_pair(i32, i32)
``````````
</details>
https://github.com/llvm/llvm-project/pull/235433
More information about the Mlir-commits
mailing list