[mlir][Affine] Remove single iteration affine.for ops in AffineLoopNormalize

[dcaballe]

Summary

This patch renames AffineParallelNormalize to AffineLoopNormalize to make it
more generic and be able to hold more loop normalization transformations in
the future for affine.for and affine.parallel ops. Eventually, it could also be
extended to support scf.for and scf.parallel. As a starting point for affine.for,
the patch also adds support for removing single iteration affine.for ops to the
the pass.

[flaub]

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

23

This needs to be defined in the mlir namespace (this was a bug I introduced originally).

97

This also needs to be defined in the mlir namespace:

void mlir::normalizeAffineFor(...

[ftynse]

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

97

This should only be defined in the mlir namespace if it was declared in it, which does not seem to be the case. For functions scoped in a translation unit, MLIR uses static instead of anonymous namespaces.

[dcaballe]

Thanks for the review!

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

23

Ok, I'll fix it before committing it.

97

Ok, I'll make it static before committing. We can make it public once it implements all the functionality.

[flaub]

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

85

I'd prefer if this was part of the public interface so that I can use it in downstream projects. Shouldn't this then be in the mlir namespace?

[dcaballe]

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

85

I think if you want to use this as a utility you may probably want to invoke the promoteIfSingleIteration directly since it's the only thing this method is doing for now. If we make normalizeAffineFor public when it's actually not normalizing the loop bounds it's going to be a bit confusing for external users, I think.

mlir/include/mlir/Dialect/Affine/Passes.h

 /// operations out of affine loops.
 std::unique_ptr<OperationPass<FuncOp>>
 createAffineLoopInvariantCodeMotionPass();
 
 /// Creates a pass to convert all parallel affine.for's into 1-d affine.parallel
 /// ops.
 std::unique_ptr<OperationPass<FuncOp>> createAffineParallelizePass();
 
-/// Normalize affine.parallel ops so that lower bounds are 0 and steps are 1.
-std::unique_ptr<OperationPass<FuncOp>> createAffineParallelNormalizePass();
+/// Apply normalization transformations to affine loop-like ops.
+std::unique_ptr<OperationPass<FuncOp>> createAffineLoopNormalizePass();
 
 /// 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<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

def AffineVectorize : FunctionPass<"affine-super-vectorize"> {
   ];
 }
 
 def AffineParallelize : FunctionPass<"affine-parallelize"> {
   let summary = "Convert affine.for ops into 1-D affine.parallel";
   let constructor = "mlir::createAffineParallelizePass()";
 }
 
-def AffineParallelNormalize : FunctionPass<"affine-parallel-normalize"> {
-  let summary = "Normalize affine.parallel ops so that lower bounds are 0 and "
-                "steps are 1";
-  let constructor = "mlir::createAffineParallelNormalizePass()";
+def AffineLoopNormalize : FunctionPass<"affine-loop-normalize"> {
+  let summary = "Apply normalization transformations to affine loop-like ops";
+  let constructor = "mlir::createAffineLoopNormalizePass()";
 }
 
 def SimplifyAffineStructures : FunctionPass<"simplify-affine-structures"> {
   let summary = "Simplify affine expressions in maps/sets and normalize "
                 "memrefs";
   let constructor = "mlir::createSimplifyAffineStructuresPass()";
 }
 
 #endif // MLIR_DIALECT_AFFINE_PASSES

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

-//===- AffineParallelNormalize.cpp - AffineParallelNormalize Pass ---------===//
+//===- AffineLoopNormalize.cpp - AffineLoopNormalize Pass -----------------===//
 //
 // 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 normalizer for affine parallel loops.
+// This file implements a normalizer for affine loop-like ops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/LoopUtils.h"
 
 using namespace mlir;
 
-void normalizeAffineParallel(AffineParallelOp op) {
+void mlir::normalizeAffineParallel(AffineParallelOp op) {
   AffineMap lbMap = op.lowerBoundsMap();
   SmallVector<int64_t, 8> steps = op.getSteps();
   // No need to do any work if the parallel op is already normalized.
   bool isAlreadyNormalized =
       llvm::all_of(llvm::zip(steps, lbMap.getResults()), [](auto tuple) {
         int64_t step = std::get<0>(tuple);
         auto lbExpr =
             std::get<1>(tuple).template dyn_cast<AffineConstantExpr>();
void mlir::normalizeAffineParallel(AffineParallelOp op) {
   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);
 }
 
+/// Normalization transformations for affine.for ops. For now, it only removes
+/// single iteration loops. We may want to consider separating redundant loop
+/// elimitation from loop bound normalization, if needed in the future.
+static void normalizeAffineFor(AffineForOp op) {
+  if (succeeded(promoteIfSingleIteration(op)))
+    return;
+
+  // TODO: Normalize loop bounds.
+}
+
 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 AffineParallelNormalizePass
-    : public AffineParallelNormalizeBase<AffineParallelNormalizePass> {
+struct AffineLoopNormalizePass
+    : public AffineLoopNormalizeBase<AffineLoopNormalizePass> {
 
-  void runOnFunction() override { getFunction().walk(normalizeAffineParallel); }
+  void runOnFunction() override {
+    getFunction().walk([](Operation *op) {
+      if (auto affineParallel = dyn_cast<AffineParallelOp>(op))
+        normalizeAffineParallel(affineParallel);
+      else if (auto affineFor = dyn_cast<AffineForOp>(op))
+        normalizeAffineFor(affineFor);
+    });
+  }
 };
 
 } // namespace
 
-std::unique_ptr<OperationPass<FuncOp>>
-mlir::createAffineParallelNormalizePass() {
-  return std::make_unique<AffineParallelNormalizePass>();
+std::unique_ptr<OperationPass<FuncOp>> mlir::createAffineLoopNormalizePass() {
+  return std::make_unique<AffineLoopNormalizePass>();
 }

mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineTransforms
   AffineDataCopyGeneration.cpp
   AffineLoopInvariantCodeMotion.cpp
+  AffineLoopNormalize.cpp
   AffineParallelize.cpp
-  AffineParallelNormalize.cpp
   LoopTiling.cpp
   LoopUnroll.cpp
   LoopUnrollAndJam.cpp
   SuperVectorize.cpp
   SimplifyAffineStructures.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine

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

-// RUN: mlir-opt %s -affine-parallel-normalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -affine-loop-normalize -split-input-file | FileCheck %s
 
 // 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 @normalize_parallel() {
   %cst = constant 1.0 : f32
   %0 = alloc() : memref<2x4xf32>
   // CHECK: affine.parallel (%[[i0:.*]], %[[j0:.*]]) = (0, 0) to (4, 2)
   affine.parallel (%i, %j) = (0, 1) to (10, 5) step (3, 2) {
     // CHECK: %[[i1:.*]] = affine.apply [[$MAP0]](%[[i0]])
     // CHECK: %[[j1:.*]] = affine.apply [[$MAP1]](%[[j0]])
     // CHECK: affine.parallel (%[[k0:.*]]) = (0) to (%[[j1]] - %[[i1]])
     affine.parallel (%k) = (%i) to (%j) {
       // CHECK: %[[k1:.*]] = affine.apply [[$MAP2]](%[[i1]], %[[k0]])
       // CHECK: affine.store %{{.*}}, %{{.*}}[%[[i1]], %[[k1]]] : memref<2x4xf32>
       affine.store %cst, %0[%i, %k] : memref<2x4xf32>
     }
   }
   return
 }
+
+// -----
+
+// Check that single iteration loop is removed and its body is promoted to the
+// parent block.
+
+// CHECK-LABEL: func @single_iteration_loop
+func @single_iteration_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