[Mlir-commits] [mlir] [mlir][licm] Make scf.if recursively speculatable (PR #122031)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jan 7 16:42:25 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-scf
Author: Arda Unal (ardaunal)
<details>
<summary>Changes</summary>
This change:
- makes **scf.if** recursively speculatable like **affine.if** is.
- also introduces related tests for both **scf.if** and **affine.if**
---
Full diff: https://github.com/llvm/llvm-project/pull/122031.diff
2 Files Affected:
- (modified) mlir/include/mlir/Dialect/SCF/IR/SCFOps.td (+2-2)
- (modified) mlir/test/Transforms/loop-invariant-code-motion.mlir (+121)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 23c597a1ca5108..6f408b3c924de8 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -302,7 +302,7 @@ def ForallOp : SCF_Op<"forall", [
AttrSizedOperandSegments,
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getInitsMutable", "getRegionIterArgs", "getLoopInductionVars",
+ ["getInitsMutable", "getRegionIterArgs", "getLoopInductionVars",
"getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps",
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
@@ -671,7 +671,7 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
"getNumRegionInvocations", "getRegionInvocationBounds",
"getEntrySuccessorRegions"]>,
InferTypeOpAdaptor, SingleBlockImplicitTerminator<"scf::YieldOp">,
- RecursiveMemoryEffects, NoRegionArguments]> {
+ RecursiveMemoryEffects, RecursivelySpeculatable, NoRegionArguments]> {
let summary = "if-then-else operation";
let description = [{
The `scf.if` operation represents an if-then-else construct for
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index e4c423ce7052bf..5133c14414c978 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -124,6 +124,64 @@ func.func @invariant_affine_if() {
// -----
+func.func @hoist_invariant_affine_if_success(%lb: index, %ub: index, %step: index) -> i32 {
+ %cst_0 = arith.constant 0 : i32
+ %cst_42 = arith.constant 42 : i32
+ %sum_result = affine.for %i = %lb to %ub iter_args(%acc = %cst_0) -> i32 {
+ %conditional_add = affine.if affine_set<() : ()> () -> (i32) {
+ %add = arith.addi %cst_42, %cst_42 : i32
+ affine.yield %add : i32
+ } else {
+ %poison = ub.poison : i32
+ affine.yield %poison : i32
+ }
+ %sum = arith.addi %acc, %conditional_add : i32
+ affine.yield %sum : i32
+ }
+
+ // CHECK-LABEL: hoist_invariant_affine_if_success
+ // CHECK-NEXT: arith.constant 0 : i32
+ // CHECK-NEXT: %[[CST:.*]] = arith.constant 42 : i32
+ // CHECK-NEXT: %[[IF:.*]] = affine.if
+ // CHECK-NEXT: arith.addi %[[CST]], %[[CST]] : i32
+ // CHECK: affine.for
+ // CHECK-NOT: affine.if
+ // CHECK-NEXT: arith.addi %{{.*}}, %[[IF]]
+
+ return %sum_result : i32
+}
+
+// -----
+
+func.func @hoist_variant_affine_if_failure(%lb: index, %ub: index, %step: index) -> i32 {
+ %cst_0 = arith.constant 0 : i32
+ %cst_42 = arith.constant 42 : i32
+ %ind_7 = arith.constant 7 : index
+ %sum_result = affine.for %i = %lb to %ub iter_args(%acc = %cst_0) -> i32 {
+ %conditional_add = affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%i, %ind_7) -> (i32) {
+ %add = arith.addi %cst_42, %cst_42 : i32
+ affine.yield %add : i32
+ } else {
+ %poison = ub.poison : i32
+ affine.yield %poison : i32
+ }
+ %sum = arith.addi %acc, %conditional_add : i32
+ affine.yield %sum : i32
+ }
+
+ // CHECK-LABEL: hoist_variant_affine_if_failure
+ // CHECK-NEXT: arith.constant 0 : i32
+ // CHECK-NEXT: %[[CST:.*]] = arith.constant 42 : i32
+ // CHECK-NEXT: arith.constant 7 : index
+ // CHECK-NEXT: affine.for
+ // CHECK-NEXT: %[[IF:.*]] = affine.if
+ // CHECK: arith.addi %{{.*}}, %[[IF]]
+
+ return %sum_result : i32
+}
+
+// -----
+
func.func @hoist_affine_for_with_unknown_trip_count(%lb: index, %ub: index) {
affine.for %arg0 = 0 to 10 {
affine.for %arg1 = %lb to %ub {
@@ -383,6 +441,69 @@ func.func @parallel_loop_with_invariant() {
// -----
+func.func @hoist_invariant_scf_if_success(%lb: index, %ub: index, %step: index) -> i32 {
+ %cst_0 = arith.constant 0 : i32
+ %cst_42 = arith.constant 42 : i32
+ %true = arith.constant true
+ %sum_result = scf.for %i = %lb to %ub step %step iter_args(%acc = %cst_0) -> i32 {
+ %conditional_add = scf.if %true -> (i32) {
+ %add = arith.addi %cst_42, %cst_42 : i32
+ scf.yield %add : i32
+ } else {
+ %poison = ub.poison : i32
+ scf.yield %poison : i32
+ }
+ %sum = arith.addi %acc, %conditional_add : i32
+ scf.yield %sum : i32
+ }
+
+ // CHECK-LABEL: hoist_invariant_scf_if_success
+ // CHECK-NEXT: arith.constant 0 : i32
+ // CHECK-NEXT: %[[CST:.*]] = arith.constant 42 : i32
+ // CHECK-NEXT: %[[TRUE:.*]] = arith.constant true
+ // CHECK-NEXT: %[[IF:.*]] = scf.if %[[TRUE]]
+ // CHECK-NEXT: arith.addi %[[CST]], %[[CST]] : i32
+ // CHECK: scf.for
+ // CHECK-NOT: scf.if
+ // CHECK-NEXT: arith.addi %{{.*}}, %[[IF]]
+
+ return %sum_result : i32
+}
+
+// -----
+
+func.func @hoist_variant_scf_if_failure(%lb: index, %ub: index, %step: index) -> i32 {
+ %cst_0 = arith.constant 0 : i32
+ %cst_42 = arith.constant 42 : i32
+ %ind_7 = arith.constant 7 : index
+ %sum_result = scf.for %i = %lb to %ub step %step iter_args(%acc = %cst_0) -> i32 {
+ %cond = arith.cmpi ult, %i, %ind_7 : index
+ %conditional_add = scf.if %cond -> (i32) {
+ %add = arith.addi %cst_42, %cst_42 : i32
+ scf.yield %add : i32
+ } else {
+ %poison = ub.poison : i32
+ scf.yield %poison : i32
+ }
+ %sum = arith.addi %acc, %conditional_add : i32
+ scf.yield %sum : i32
+ }
+
+ // CHECK-LABEL: hoist_variant_scf_if_failure
+ // CHECK-NEXT: arith.constant 0 : i32
+ // CHECK-NEXT: %[[CST_42:.*]] = arith.constant 42 : i32
+ // CHECK-NEXT: %[[CST_7:.*]] = arith.constant 7 : index
+ // CHECK-NEXT: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}}
+ // CHECK-NEXT: %[[CMP:.*]] = arith.cmpi ult, %[[IV]], %[[CST_7]]
+ // CHECK-NEXT: %[[IF:.*]] = scf.if %[[CMP]]
+ // CHECK-NEXT: arith.addi %[[CST_42]], %[[CST_42]] : i32
+ // CHECK: arith.addi %{{.*}}, %[[IF]]
+
+ return %sum_result : i32
+}
+
+// -----
+
func.func private @make_val() -> (index)
// CHECK-LABEL: func @nested_uses_inside
``````````
</details>
https://github.com/llvm/llvm-project/pull/122031
More information about the Mlir-commits
mailing list