[Mlir-commits] [mlir] 788390c - Make scf.for and affine.for conditionally speculatable
Sanjoy Das
llvmlistbot at llvm.org
Sun Oct 30 16:24:45 PDT 2022
Author: Sanjoy Das
Date: 2022-10-30T16:08:42-07:00
New Revision: 788390c1309cbac772f318deff2c09921dadac88
URL: https://github.com/llvm/llvm-project/commit/788390c1309cbac772f318deff2c09921dadac88
DIFF: https://github.com/llvm/llvm-project/commit/788390c1309cbac772f318deff2c09921dadac88.diff
LOG: Make scf.for and affine.for conditionally speculatable
for (I = Start; I < End; I += 1) always terminates so mark
{scf|affine}.for as RecursivelySpeculatable when step is known to be
1.
Reviewed By: chelini
Differential Revision: https://reviews.llvm.org/D136376
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/lib/Dialect/SCF/IR/SCF.cpp
mlir/test/Transforms/loop-invariant-code-motion.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index bb505c52168d6..148e2a4174ac2 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -105,7 +105,7 @@ def AffineApplyOp : Affine_Op<"apply", [Pure]> {
}
def AffineForOp : Affine_Op<"for",
- [AutomaticAllocationScope, ImplicitAffineTerminator, RecursivelySpeculatable,
+ [AutomaticAllocationScope, ImplicitAffineTerminator, ConditionallySpeculatable,
RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
"getSingleUpperBound"]>,
@@ -340,6 +340,9 @@ def AffineForOp : Affine_Op<"for",
/// Returns true if both the lower and upper bound have the same operand
/// lists (same operands in the same order).
bool matchingBoundOperandList();
+
+ /// Interface method for ConditionallySpeculatable.
+ Speculation::Speculatability getSpeculatability();
}];
let hasCanonicalizer = 1;
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 38dd0acc6a69e..d576ca28efff7 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -119,6 +119,7 @@ def ForOp : SCF_Op<"for",
[AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
"getSingleUpperBound"]>,
+ ConditionallySpeculatable,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
SingleBlockImplicitTerminator<"scf::YieldOp">,
RecursiveMemoryEffects]> {
@@ -330,6 +331,12 @@ def ForOp : SCF_Op<"for",
/// induction variable. LoopOp only has one region, so 0 is the only valid
/// value for `index`.
OperandRange getSuccessorEntryOperands(Optional<unsigned> index);
+
+ /// Returns the step as an `APInt` if it is constant.
+ Optional<APInt> getConstantStep();
+
+ /// Interface method for ConditionallySpeculatable.
+ Speculation::Speculatability getSpeculatability();
}];
let hasCanonicalizer = 1;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 01cc8b31f7c21..4a581727e190c 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2273,6 +2273,16 @@ Optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()));
}
+Speculation::Speculatability AffineForOp::getSpeculatability() {
+ // `affine.for (I = Start; I < End; I += 1)` terminates for all values of
+ // Start and End.
+ //
+ // For Step != 1, the loop may not terminate. We can add more smarts here if
+ // needed.
+ return getStep() == 1 ? Speculation::RecursivelySpeculatable
+ : Speculation::NotSpeculatable;
+}
+
/// Returns true if the provided value is the induction variable of a
/// AffineForOp.
bool mlir::isForInductionVar(Value val) {
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index cd1d3829cab9a..684b911f55d45 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -762,13 +762,13 @@ struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
return success();
}
- IntegerAttr step;
- if (!matchPattern(op.getStep(), m_Constant(&step)))
+ llvm::Optional<llvm::APInt> maybeStepValue = op.getConstantStep();
+ if (!maybeStepValue)
return failure();
// If the loop is known to have 1 iteration, inline its body and remove the
// loop.
- llvm::APInt stepValue = step.getValue();
+ llvm::APInt stepValue = *maybeStepValue;
if (stepValue.sge(*
diff )) {
SmallVector<Value, 4> blockArgs;
blockArgs.reserve(op.getNumIterOperands() + 1);
@@ -1065,6 +1065,25 @@ void ForOp::getCanonicalizationPatterns(RewritePatternSet &results,
LastTensorLoadCanonicalization, ForOpTensorCastFolder>(context);
}
+Optional<APInt> ForOp::getConstantStep() {
+ IntegerAttr step;
+ if (matchPattern(getStep(), m_Constant(&step)))
+ return step.getValue();
+ return {};
+}
+
+Speculation::Speculatability ForOp::getSpeculatability() {
+ // `scf.for (I = Start; I < End; I += 1)` terminates for all values of Start
+ // and End.
+ if (auto constantStep = getConstantStep())
+ if (*constantStep == 1)
+ return Speculation::RecursivelySpeculatable;
+
+ // For Step != 1, the loop may not terminate. We can add more smarts here if
+ // needed.
+ return Speculation::NotSpeculatable;
+}
+
//===----------------------------------------------------------------------===//
// ForeachThreadOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index b8d3450862a39..9f5d9a55053ed 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -103,7 +103,7 @@ func.func @invariant_affine_if() {
%m = memref.alloc() : memref<10xf32>
%cf8 = arith.constant 8.0 : f32
affine.for %arg0 = 0 to 10 {
- affine.for %arg1 = 0 to 10 {
+ affine.for %arg1 = 0 to 20 {
affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
%cf9 = arith.addf %cf8, %cf8 : f32
}
@@ -112,7 +112,7 @@ func.func @invariant_affine_if() {
// CHECK: memref.alloc() : memref<10xf32>
// CHECK-NEXT: %[[CST:.*]] = arith.constant 8.000000e+00 : f32
- // CHECK-NEXT: affine.for %[[ARG:.*]] = 0 to 10 {
+ // CHECK-NEXT: affine.for %[[ARG:.*]] = 0 to 20 {
// CHECK-NEXT: }
// CHECK-NEXT: affine.for %[[ARG:.*]] = 0 to 10 {
// CHECK-NEXT: affine.if #set(%[[ARG]], %[[ARG]]) {
@@ -124,6 +124,95 @@ func.func @invariant_affine_if() {
// -----
+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 {
+ }
+ }
+
+ // CHECK: @hoist_affine_for_with_unknown_trip_count(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) {
+ // CHECK-NEXT: affine.for %[[ARG2:.*]] = %[[ARG0]] to %[[ARG1]] {
+ // CHECK-NEXT: }
+ // CHECK-NEXT: affine.for %[[ARG3:.*]] = 0 to 10 {
+ // CHECK-NEXT: }
+
+ return
+}
+
+// -----
+
+func.func @hoist_affine_for_with_unknown_trip_count_non_unit_step(%lb: index, %ub: index) {
+ affine.for %arg0 = 0 to 10 {
+ affine.for %arg1 = %lb to %ub step 2 {
+ }
+ }
+
+ // CHECK: @hoist_affine_for_with_unknown_trip_count_non_unit_step(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) {
+ // CHECK-NEXT: affine.for %[[ARG2:.*]] = 0 to 10 {
+ // CHECK-NEXT: affine.for %[[ARG3:.*]] = %[[ARG0]] to %[[ARG1]] step 2 {
+ // CHECK-NEXT: }
+ // CHECK-NEXT: }
+
+ return
+}
+
+// -----
+
+func.func @hoist_scf_for_with_unknown_trip_count_unit_step(%lb: index, %ub: index) {
+ %c1 = arith.constant 1 : index
+ scf.for %arg0 = %lb to %ub step %c1 {
+ scf.for %arg1 = %lb to %ub step %c1 {
+ }
+ }
+
+ // CHECK: @hoist_scf_for_with_unknown_trip_count_unit_step(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) {
+ // CHECK: scf.for %[[ARG2:.*]] = %[[ARG0]] to %[[ARG1]]
+ // CHECK-NEXT: }
+ // CHECK-NEXT: scf.for %[[ARG3:.*]] = %[[ARG0]] to %[[ARG1]]
+ // CHECK-NEXT: }
+
+ return
+}
+
+// -----
+
+func.func @hoist_scf_for_with_unknown_trip_count_non_unit_constant_step(%lb: index, %ub: index) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ scf.for %arg0 = %lb to %ub step %c1 {
+ scf.for %arg1 = %lb to %ub step %c2 {
+ }
+ }
+
+ // CHECK: @hoist_scf_for_with_unknown_trip_count_non_unit_constant_step(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) {
+ // CHECK: scf.for %[[ARG2:.*]] = %[[ARG0]] to %[[ARG1]]
+ // CHECK-NEXT: scf.for %[[ARG3:.*]] = %[[ARG0]] to %[[ARG1]]
+ // CHECK-NEXT: }
+ // CHECK-NEXT: }
+
+ return
+}
+
+// -----
+
+func.func @hoist_scf_for_with_unknown_trip_count_unknown_step(%lb: index, %ub: index, %step: index) {
+ %c1 = arith.constant 1 : index
+ scf.for %arg0 = %lb to %ub step %c1 {
+ scf.for %arg1 = %lb to %ub step %step {
+ }
+ }
+
+ // CHECK: @hoist_scf_for_with_unknown_trip_count_unknown_step(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[STEP:.*]]: index) {
+ // CHECK: scf.for %[[ARG2:.*]] = %[[ARG0]] to %[[ARG1]]
+ // CHECK-NEXT: scf.for %[[ARG3:.*]] = %[[ARG0]] to %[[ARG1]] step %[[STEP]]
+ // CHECK-NEXT: }
+ // CHECK-NEXT: }
+
+ return
+}
+
+// -----
+
func.func @invariant_affine_if2() {
%m = memref.alloc() : memref<10xf32>
%cf8 = arith.constant 8.0 : f32
More information about the Mlir-commits
mailing list