[Mlir-commits] [mlir] b3ce6dc - [mlir][licm] Make scf.if recursively speculatable (#122031)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jan 8 09:54:22 PST 2025
Author: Arda Unal
Date: 2025-01-08T09:54:18-08:00
New Revision: b3ce6dc7232c566c21b84ac5d5795341a355ff79
URL: https://github.com/llvm/llvm-project/commit/b3ce6dc7232c566c21b84ac5d5795341a355ff79
DIFF: https://github.com/llvm/llvm-project/commit/b3ce6dc7232c566c21b84ac5d5795341a355ff79.diff
LOG: [mlir][licm] Make scf.if recursively speculatable (#122031)
This change:
- makes **scf.if** recursively speculatable like **affine.if** is.
- also introduces related LICM tests for both **scf.if** and
**affine.if**
Added:
Modified:
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
mlir/test/Transforms/loop-invariant-code-motion.mlir
Removed:
################################################################################
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
More information about the Mlir-commits
mailing list