[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/include/mlir/Dialect/Affine/Passes.h

 /// operations out of affine loops.
 std::unique_ptr<OperationPass<func::FuncOp>>
 createAffineLoopInvariantCodeMotionPass();
 
 /// Creates a pass to convert all parallel affine.for's into 1-d affine.parallel
 /// ops.
 std::unique_ptr<OperationPass<func::FuncOp>> createAffineParallelizePass();
 
-/// Apply normalization transformations to affine loop-like ops.
-std::unique_ptr<OperationPass<func::FuncOp>> createAffineLoopNormalizePass();
+/// Apply normalization transformations to affine loop-like ops. If
+/// `promoteSingleIter` is true, single iteration loops are promoted (i.e., the
+/// loop is replaced by its loop body).
+std::unique_ptr<OperationPass<func::FuncOp>>
+createAffineLoopNormalizePass(bool promoteSingleIter = false);
 
 /// Performs packing (or explicit copying) of accessed memref regions into
 /// buffers in the specified faster memory space through either pointwise copies
 /// or DMA operations.
 std::unique_ptr<OperationPass<func::FuncOp>> createAffineDataCopyGenerationPass(
     unsigned slowMemorySpace, unsigned fastMemorySpace,
     unsigned tagMemorySpace = 0, int minDmaTransferSize = 1024,
     uint64_t fastMemCapacityBytes = std::numeric_limits<uint64_t>::max());

mlir/include/mlir/Dialect/Affine/Passes.td

Option<"parallelReductions", "parallel-reductions", "bool",
            /*default=*/"false",
            "Whether to parallelize reduction loops. Defaults to false.">
   ];
 }
 
 def AffineLoopNormalize : Pass<"affine-loop-normalize", "func::FuncOp"> {
   let summary = "Apply normalization transformations to affine loop-like ops";
   let constructor = "mlir::createAffineLoopNormalizePass()";
+  let options = [
+    Option<"promoteSingleIter", "promote-single-iter", "bool",
+           /*default=*/"true", "Promote single iteration loops">,
+  ];
 }
 
 def LoopCoalescing : Pass<"affine-loop-coalescing", "func::FuncOp"> {
   let summary = "Coalesce nested loops with independent bounds into a single "
                 "loop";
   let constructor = "mlir::createLoopCoalescingPass()";
   let dependentDialects = ["arith::ArithDialect"];
 }

mlir/include/mlir/Dialect/Affine/Utils.h

 vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
                         const VectorizationStrategy &strategy);
 
 /// Normalize a affine.parallel op so that lower bounds are 0 and steps are 1.
 /// As currently implemented, this transformation cannot fail and will return
 /// early if the op is already in a normalized form.
 void normalizeAffineParallel(AffineParallelOp op);
 
-/// Normalize an affine.for op. If the affine.for op has only a single iteration
-/// only then it is simply promoted, else it is normalized in the traditional
-/// way, by converting the lower bound to zero and loop step to one. The upper
-/// bound is set to the trip count of the loop. Original loops must have a
-/// lower bound with only a single result. There is no such restriction on upper
-/// bounds. Returns success if the loop has been normalized (or is already in
-/// the normal form).
-LogicalResult normalizeAffineFor(AffineForOp op);
+/// Normalize an affine.for op. An affine.for op is normalized by converting the
+/// lower bound to zero and loop step to one. The upper bound is set to the trip
+/// count of the loop. Original loops must have a lower bound with only a single
+/// result. There is no such restriction on upper bounds. Returns success if the
+/// loop has been normalized (or is already in the normal form). If
+/// `promoteSingleIter` is true, the loop is simply promoted if it has a single
+/// iteration.
+LogicalResult normalizeAffineFor(AffineForOp op,
+                                 bool promoteSingleIter = false);
 
 /// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
 /// been replaced by either:
 ///  - `min` if `positivePath` is true when we reach an occurrence of `dim`
 ///  - `max` if `positivePath` is true when we reach an occurrence of `dim`
 /// `positivePath` is negated each time we hit a multiplicative or divisive
 /// binary op with a constant negative coefficient.
 AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,

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

 
 namespace {
 
 /// Normalize affine.parallel ops so that lower bounds are 0 and steps are 1.
 /// As currently implemented, this pass cannot fail, but it might skip over ops
 /// that are already in a normalized form.
 struct AffineLoopNormalizePass
     : public impl::AffineLoopNormalizeBase<AffineLoopNormalizePass> {
+  explicit AffineLoopNormalizePass(bool promoteSingleIter) {
+    this->promoteSingleIter = promoteSingleIter;
+  }
 
   void runOnOperation() override {
-    getOperation().walk([](Operation *op) {
+    getOperation().walk([&](Operation *op) {
       if (auto affineParallel = dyn_cast<AffineParallelOp>(op))
         normalizeAffineParallel(affineParallel);
       else if (auto affineFor = dyn_cast<AffineForOp>(op))
-        (void)normalizeAffineFor(affineFor);
+        (void)normalizeAffineFor(affineFor, promoteSingleIter);
     });
   }
 };
 
 } // namespace
 
 std::unique_ptr<OperationPass<func::FuncOp>>
-mlir::createAffineLoopNormalizePass() {
-  return std::make_unique<AffineLoopNormalizePass>();
+mlir::createAffineLoopNormalizePass(bool promoteSingleIter) {
+  return std::make_unique<AffineLoopNormalizePass>(promoteSingleIter);
 }

mlir/lib/Dialect/Affine/Utils/Utils.cpp

void mlir::normalizeAffineParallel(AffineParallelOp op) {
   op.setSteps(newSteps);
   auto newLowerMap = AffineMap::get(
       /*dimCount=*/0, /*symbolCount=*/0, lbExprs, op.getContext());
   op.setLowerBounds({}, newLowerMap);
   auto newUpperMap = AffineMap::get(ranges.getNumDims(), ranges.getNumSymbols(),
                                     ubExprs, op.getContext());
   op.setUpperBounds(ranges.getOperands(), newUpperMap);
 }
-
-/// Normalizes affine.for ops. If the affine.for op has only a single iteration
-/// only then it is simply promoted, else it is normalized in the traditional
-/// way, by converting the lower bound to zero and loop step to one. The upper
-/// bound is set to the trip count of the loop. For now, original loops must
-/// have lower bound with a single result only. There is no such restriction on
-/// upper bounds.
-LogicalResult mlir::normalizeAffineFor(AffineForOp op) {
-  if (succeeded(promoteIfSingleIteration(op)))
+LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
+  if (promoteSingleIter && succeeded(promoteIfSingleIteration(op)))
     return success();
 
   // Check if the forop is already normalized.
   if (op.hasConstantLowerBound() && (op.getConstantLowerBound() == 0) &&
       (op.getStep() == 1))
     return success();
 
   // Check if the lower bound has a single result only. Loops with a max lower

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

 // RUN: mlir-opt %s -affine-loop-normalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -affine-loop-normalize='promote-single-iter=1' -split-input-file | FileCheck %s --check-prefix=PROMOTE-SINGLE-ITER
 
 // Normalize steps to 1 and lower bounds to 0.
 
 // CHECK-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 * 3)>
 // CHECK-DAG: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (d0 * 2 + 1)>
 // CHECK-DAG: [[$MAP2:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + d1)>
 
 // CHECK-LABEL: func @normalize_parallel()
func.func @relative_bounds(%arg: index) {
   return
 }
 
 // -----
 
 // Check that single iteration loop is removed and its body is promoted to the
 // parent block.
 
-// CHECK-LABEL: func @single_iteration_loop
-func.func @single_iteration_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
+// CHECK-LABEL: func @promote_single_iter_loop
+// PROMOTE-SINGLE-ITER-LABEL: func @promote_single_iter_loop
+func.func @promote_single_iter_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
   affine.for %i = 0 to 1 {
     %1 = affine.load %in[%i] : memref<1xf32>
     affine.store %1, %out[%i] : memref<1xf32>
   }
   return
 }
 
-// CHECK-NOT:  affine.for
-// CHECK:      affine.load
-// CHECK-NEXT: affine.store
-// CHECK-NEXT: return
+// PROMOTE-SINGLE-ITER-NEXT: arith.constant
+// PROMOTE-SINGLE-ITER-NEXT: affine.load
+// PROMOTE-SINGLE-ITER-NEXT: affine.store
+// PROMOTE-SINGLE-ITER-NEXT: return
 
 // -----
 
 // CHECK-DAG: [[$IV0:#map[0-9]*]] = affine_map<(d0) -> (d0 * 2 + 2)>
 // CHECK-DAG: [[$IV1:#map[0-9]*]] = affine_map<(d0) -> (d0 * 3)>
 
 // CHECK-LABEL: func @simple_loop_nest()
 // CHECK-NEXT:   affine.for %[[I:.*]] = 0 to 15 {