[MLIR] Make promote single iteration optional on affine-loop-normalize

[bondhugula]

Summary

Make promote single iteration optional on affine-loop-normalize:
introduce a command-line flag and an argument on the utility. Disable it
by default since such a promotion isn't normally expected with loop
normalization: it could drop certain structure or information on the
loops that a user wanted to preserve.

[dcaballe]

We initially added this to --affine-loop-normalize. Probably it's not the best place but removing loops as part of canonicalization or a common simplification pass that is run all over the pipeline can be problematic as it drops information about the input dimensions. This may prevent or make some transformations harder. I'd suggest that we use the loop specific normalization pass or create a new pass to simplify loops and other control-flow structures. What do you think?

[bondhugula]

In D137605#3914067, @dcaballe wrote:

We initially added this to --affine-loop-normalize. Probably it's not the best place but removing loops as part of canonicalization or a common simplification pass that is run all over the pipeline can be problematic as it drops information about the input dimensions. This may prevent or make some transformations harder. I'd suggest that we use the loop specific normalization pass or create a new pass to simplify loops and other control-flow structures. What do you think?

Another alternative/solution is to perform this promotion under a pass option? Can be enabled or disabled and perhaps disabled by default?

[bondhugula]

In D137605#3914174, @bondhugula wrote:

In D137605#3914067, @dcaballe wrote:

We initially added this to --affine-loop-normalize. Probably it's not the best place but removing loops as part of canonicalization or a common simplification pass that is run all over the pipeline can be problematic as it drops information about the input dimensions. This may prevent or make some transformations harder. I'd suggest that we use the loop specific normalization pass or create a new pass to simplify loops and other control-flow structures. What do you think?

Another alternative/solution is to perform this promotion under a pass option? Can be enabled or disabled and perhaps disabled by default?

Thanks. I just noticed that it's already part of affine-loop-normalize - had forgotten about it. Perhaps we should either remove or disable it by default from there as well for the reasons you mention.

556 LogicalResult mlir::normalizeAffineFor(AffineForOp op) {
557   if (succeeded(promoteIfSingleIteration(op)))
558     return success();

[dcaballe]

Yeah, sorry, my comment sounded as if we had removed it from affine-loop-normalize and that's not the case :). A flag or a new pass sound good to me.

[bondhugula]

Revise to add flag on affine-loop-normalize pass.

[bondhugula]

Update commit summary.

[bondhugula]

In D137605#3914244, @dcaballe wrote:

Yeah, sorry, my comment sounded as if we had removed it from affine-loop-normalize and that's not the case :). A flag or a new pass sound good to me.

Done - added a flag.

[dcaballe]

Thanks!

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp

 //===- SimplifyAffineStructures.cpp ---------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to simplify affine structures in operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Passes.h"
 
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Dialect/Affine/LoopUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_SIMPLIFYAFFINESTRUCTURES
 #include "mlir/Dialect/Affine/Passes.h.inc"
 } // namespace mlir

 } // namespace
 
 std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::createSimplifyAffineStructuresPass() {
   return std::make_unique<SimplifyAffineStructures>();
 }
 
 void SimplifyAffineStructures::runOnOperation() {
-  auto func = getOperation();
+  func::FuncOp func = getOperation();
+
+  // Promote any single iteration affine loops. This relies on an
+  // analysis utility, and is thus not part of canonicalization.
+  func.walk([&](AffineForOp forOp) { (void)promoteIfSingleIteration(forOp); });
+
   simplifiedAttributes.clear();
   RewritePatternSet patterns(func.getContext());
   AffineApplyOp::getCanonicalizationPatterns(patterns, func.getContext());
   AffineForOp::getCanonicalizationPatterns(patterns, func.getContext());
   AffineIfOp::getCanonicalizationPatterns(patterns, func.getContext());
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
 
   // The simplification of affine attributes will likely simplify the op. Try to

mlir/test/Dialect/Affine/simplify-structures.mlir

affine.for %m = 0 to 10 {
     affine.if affine_set<(d0) : (d0 mod 2 - 3 == 0)> (%m) {
       "foo"() : () -> ()
     }
   }
 
   return
 }
 
+// CHECK-LABEL: func.func @promote_single_iter
+func.func @promote_single_iter(%N: index) {
+  // CHECK-NEXT: test.foo
+  // CHECK-NEXT: return
+  affine.for %i = %N to affine_map<(d0) -> (d0 + 1)>(%N) {
+    "test.foo"() : () -> ()
+  }
+  return
+}
+
 // -----
 
 // An external function that we will use in bodies to avoid DCE.
 func.func private @external() -> ()
 
 // CHECK-DAG: #[[$SET:.*]] = affine_set<()[s0] : (s0 >= 0, -s0 + 50 >= 0)
 
 // CHECK-LABEL: func @simplify_set