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

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.

Differential Revision: https://reviews.llvm.org/D90267

Author: dcaballe

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

@@ -35,8 +35,8 @@ createAffineLoopInvariantCodeMotionPass();
 /// 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

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

@@ -120,10 +120,9 @@ def AffineParallelize : FunctionPass<"affine-parallelize"> {
   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"> {

mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp renamed to mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp

@@ -1,24 +1,26 @@
-//===- 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.
@@ -77,20 +79,36 @@ void normalizeAffineParallel(AffineParallelOp op) {
   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

@@ -1,8 +1,8 @@
 add_mlir_dialect_library(MLIRAffineTransforms
   AffineDataCopyGeneration.cpp
   AffineLoopInvariantCodeMotion.cpp
+  AffineLoopNormalize.cpp
   AffineParallelize.cpp
-  AffineParallelNormalize.cpp
   LoopTiling.cpp
   LoopUnroll.cpp
   LoopUnrollAndJam.cpp

mlir/test/Dialect/Affine/affine-parallel-normalize.mlir renamed to mlir/test/Dialect/Affine/affine-loop-normalize.mlir

@@ -1,4 +1,4 @@
-// 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.
 
@@ -23,3 +23,22 @@ func @normalize_parallel() {
   }
   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