Make scf.for and affine.for conditionally speculatable

[sanjoy]

Summary

for (I = Start; I < End; I += 1) always terminates so mark
{scf|affine}.for as RecursivelySpeculatable when step is known to be
1.

[zero9178]

mlir/lib/Dialect/SCF/IR/SCF.cpp

1069–1070

Using matchPattern with m_Constant and friends makes the code indepdent of any specific dialects ConstantLike op

[sanjoy]

• Use m_Constant
• Change getConstantStep() to return APInt (that's more robust)
• Refactor one obvious place in the code to use ForOp::getConstantStep()

[bondhugula]

mlir/test/Transforms/loop-invariant-code-motion.mlir

106

Any reason this has become 20?

[sanjoy]

mlir/test/Transforms/loop-invariant-code-motion.mlir

106

Just so that I could visually distinguish between the two loops.

[rengolin]

mlir/lib/Dialect/SCF/IR/SCF.cpp

1076

Don't you also have to make sure neither the induction variable nor the range change within the loop?

This is an scf.for, the most basic loop, anything can happen inside...

[sanjoy]

• Address review

mlir/lib/Dialect/SCF/IR/SCF.cpp

1076

Don't you also have to make sure neither the induction variable nor the range change within the loop?

I used a C for loop just for illustration, scf.for has more restrictive semantics and cannot change the IV or the bounds in the loop body. I'll change it to say scf.for and affine.for to make this clear.

[sanjoy]

Friendly ping!

[rengolin]

@chelini

[chelini]

mlir/test/Transforms/loop-invariant-code-motion.mlir

200

nit: %c2 is not used here. Please remove it.

[chelini]

Thanks to me looks good. But please have a look at the build status, as it is currently failing.

[sanjoy]

• Address review

[sanjoy]

In D136376#3892425, @chelini wrote:

Thanks to me looks good. But please have a look at the build status, as it is currently failing.

They're passing after the rebase.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

def AffineApplyOp : Affine_Op<"apply", [Pure]> {
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def AffineForOp : Affine_Op<"for",
-    [AutomaticAllocationScope, ImplicitAffineTerminator, RecursivelySpeculatable,
+    [AutomaticAllocationScope, ImplicitAffineTerminator, ConditionallySpeculatable,
      RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
      ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
       "getSingleUpperBound"]>,
      DeclareOpInterfaceMethods<RegionBranchOpInterface,
      ["getSuccessorEntryOperands"]>]> {
   let summary = "for operation";
   let description = [{
     Syntax:
let extraClassDeclaration = [{
     /// Sets the lower bound to the given constant value.
     void setConstantLowerBound(int64_t value);
     /// Sets the upper bound to the given constant value.
     void setConstantUpperBound(int64_t value);
 
     /// 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;
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
   let hasRegionVerifier = 1;
 }
 

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

 //===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 def ForOp : SCF_Op<"for",
       [AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
        ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
         "getSingleUpperBound"]>,
+       ConditionallySpeculatable,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">,
        RecursiveMemoryEffects]> {
   let summary = "for operation";
   let description = [{
     The "scf.for" operation represents a loop taking 3 SSA value as operands
     that represent the lower bound, upper bound and step respectively. The
     operation defines an SSA value for its induction variable. It has one
OpOperand &getOpOperandForResult(OpResult opResult) {
         getNumControlOperands() + opResult.getResultNumber());
     }
 
     /// Return operands used when entering the region at 'index'. These operands
     /// correspond to the loop iterator operands, i.e., those exclusing the
     /// 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;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
 }
 

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 
 Optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
   if (!hasConstantUpperBound())
     return llvm::None;
   OpBuilder b(getContext());
   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) {
   return getForInductionVarOwner(val) != AffineForOp();
 }
 
 /// Returns the loop parent of an induction variable. If the provided value is
 /// not an induction variable, then return nullptr.

mlir/lib/Dialect/SCF/IR/SCF.cpp

if (!diff)
       return failure();
 
     // If the loop is known to have 0 iterations, remove it.
     if (*diff <= 0) {
       rewriter.replaceOp(op, op.getIterOperands());
       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);
       blockArgs.push_back(op.getLowerBound());
       llvm::append_range(blockArgs, op.getIterOperands());
       replaceOpWithRegion(rewriter, op, op.getLoopBody(), blockArgs);
       return success();
     }

 } // namespace
 
 void ForOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                         MLIRContext *context) {
   results.add<ForOpIterArgsFolder, SimplifyTrivialLoops,
               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
 //===----------------------------------------------------------------------===//
 
 LogicalResult ForeachThreadOp::verify() {
   // Call terminator's verify to produce most informative error messages.
   if (failed(getTerminator().verify()))
     return failure();

mlir/test/Transforms/loop-invariant-code-motion.mlir

 }
 
 // -----
 
 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
       }
     }
   }
 
   // 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]]) {
   // CHECK-NEXT: arith.addf %[[CST]], %[[CST]] : f32
   // CHECK-NEXT: }
 
   return
 }
 
 // -----
 
+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
+  %c2 = arith.constant 2 : 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
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
           %cf9 = arith.addf %cf8, %cf8 : f32
           affine.store %cf9, %m[%arg1] : memref<10xf32>