[mlir] canonicalize away zero-iteration SCF for loops

[ftynse]

Summary

An SCF 'for' loop does not iterate if its lower bound is equal to its upper
bound. Remove loops where both bounds are the same SSA value as such bounds are
guaranteed to be equal. Similarly, remove 'parallel' loops where at least one
pair of respective lower/upper bounds is specified by the same SSA value.

[gysit]

mlir/test/Dialect/SCF/canonicalize.mlir

35

I would probably remove this test? It originally tested that the "CollapseSingleIterationLoops" pattern does not accidentally remove no/zero iteration loops. I think your new pattern/test now covers this scenario.

[ftynse]

Address review

mlir/lib/Dialect/SCF/SCF.cpp

 
 /// Rewriting pattern that erases loops that are known not to iterate and
 /// replaces single-iteration loops with their bodies.
 struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
   using OpRewritePattern<ForOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForOp op,
                                 PatternRewriter &rewriter) const override {
+    // If the upper bound is the same as the lower bound, the loop does not
+    // iterate, just remove it.
+    if (op.lowerBound() == op.upperBound()) {
+      rewriter.replaceOp(op, op.getIterOperands());
+      return success();
+    }
+
     auto lb = op.lowerBound().getDefiningOp<ConstantOp>();
     auto ub = op.upperBound().getDefiningOp<ConstantOp>();
     if (!lb || !ub)
       return failure();
 
     // If the loop is known to have 0 iterations, remove it.
     llvm::APInt lbValue = lb.getValue().cast<IntegerAttr>().getValue();
     llvm::APInt ubValue = ub.getValue().cast<IntegerAttr>().getValue();
LogicalResult matchAndRewrite(ParallelOp op,
     // Clone the loop body and remap the block arguments of the collapsed loops
     // (inlining does not support a cancellable block argument mapping).
     rewriter.cloneRegionBefore(op.region(), newOp.region(),
                                newOp.region().begin(), mapping);
     rewriter.replaceOp(op, newOp.getResults());
     return success();
   }
 };
+
+/// Removes parallel loops in which at least one lower/upper bound pair consists
+/// of the same values - such loops have an empty iteration domain.
+struct RemoveEmptyParallelLoops : public OpRewritePattern<ParallelOp> {
+  using OpRewritePattern<ParallelOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ParallelOp op,
+                                PatternRewriter &rewriter) const override {
+    for (auto dim : llvm::zip(op.lowerBound(), op.upperBound())) {
+      if (std::get<0>(dim) == std::get<1>(dim)) {
+        rewriter.replaceOp(op, op.initVals());
+        return success();
+      }
+    }
+    return failure();
+  }
+};
+
 } // namespace
 
 void ParallelOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                              MLIRContext *context) {
-  results.insert<CollapseSingleIterationLoops>(context);
+  results.insert<CollapseSingleIterationLoops, RemoveEmptyParallelLoops>(
+      context);
 }
 
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
 void ReduceOp::build(
     OpBuilder &builder, OperationState &result, Value operand,

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK:           scf.parallel ([[V0:%.*]]) = ([[C3]]) to ([[C6]]) step ([[C2]]) {
 // CHECK:             store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V0]], [[C7]]] : memref<?x?x?xi32>
 // CHECK:             scf.yield
 // CHECK:           }
 // CHECK:           return
 
 // -----
 
-func @no_iteration(%A: memref<?x?xi32>) {
+func @no_iteration(%A: memref<?x?xi32>, %ub: index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
-  scf.parallel (%i0, %i1) = (%c0, %c0) to (%c1, %c0) step (%c1, %c1) {
+  scf.parallel (%i0, %i1) = (%c0, %c0) to (%c1, %ub) step (%c1, %c1) {
     %c42 = constant 42 : i32
     store %c42, %A[%i0, %i1] : memref<?x?xi32>
     scf.yield
   }
   return
 }
 
 // CHECK-LABEL:   func @no_iteration(
-// CHECK-SAME:                        [[ARG0:%.*]]: memref<?x?xi32>) {
+// CHECK-SAME:                 [[ARG0:%.*]]: memref<?x?xi32>, [[UB:%.*]]: index) {
 // CHECK:           [[C0:%.*]] = constant 0 : index
 // CHECK:           [[C1:%.*]] = constant 1 : index
 // CHECK:           [[C42:%.*]] = constant 42 : i32
-// CHECK:           scf.parallel ([[V1:%.*]]) = ([[C0]]) to ([[C0]]) step ([[C1]]) {
+// CHECK:           scf.parallel ([[V1:%.*]]) = ([[C0]]) to ([[UB]]) step ([[C1]]) {
 // CHECK:             store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V1]]] : memref<?x?xi32>
 // CHECK:             scf.yield
 // CHECK:           }
 // CHECK:           return
 
 // -----
 
 func @one_unused(%cond: i1) -> (index) {
%0 = scf.for %i = %c42 to %c1 step %c1 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[INIT]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// CHECK-LABEL: @remove_zero_iteration_loop_vals
+func @remove_zero_iteration_loop_vals(%arg0: index) {
+  %c2 = constant 2 : index
+  // CHECK: %[[INIT:.*]] = "test.init"
+  %init = "test.init"() : () -> i32
+  // CHECK-NOT: scf.for
+  // CHECK-NOT: test.op
+  %0 = scf.for %i = %arg0 to %arg0 step %c2 iter_args(%arg = %init) -> (i32) {
+    %1 = "test.op"(%i, %arg) : (index, i32) -> i32
+    scf.yield %1 : i32
+  }
+  // CHECK: "test.consume"(%[[INIT]])
+  "test.consume"(%0) : (i32) -> ()
+  return
+}
+
 // CHECK-LABEL: @replace_single_iteration_loop
 func @replace_single_iteration_loop() {
   // CHECK: %[[LB:.*]] = constant 42
   %c42 = constant 42 : index
   %c43 = constant 43 : index
   %c1 = constant 1 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
func @replace_single_iteration_loop_non_unit_step() {
   %0 = scf.for %i = %c42 to %c47 step %c5 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
+
+// CHECK-LABEL: @remove_empty_parallel_loop
+func @remove_empty_parallel_loop(%lb: index, %ub: index, %s: index) {
+  // CHECK: %[[INIT:.*]] = "test.init"
+  %init = "test.init"() : () -> f32
+  // CHECK-NOT: scf.parallel
+  // CHECK-NOT: test.produce
+  // CHECK-NOT: test.transform
+  %0 = scf.parallel (%i, %j, %k) = (%lb, %ub, %lb) to (%ub, %ub, %ub) step (%s, %s, %s) init(%init) -> f32 {
+    %1 = "test.produce"() : () -> f32
+    scf.reduce(%1) : f32 {
+    ^bb0(%lhs: f32, %rhs: f32):
+      %2 = "test.transform"(%lhs, %rhs) : (f32, f32) -> f32
+      scf.reduce.return %2 : f32
+    }
+    scf.yield
+  }
+  // CHECK: "test.consume"(%[[INIT]])
+  "test.consume"(%0) : (f32) -> ()
+  return
+}