[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

diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h
index 580fbf5..57c9b88 100644 (file)
--- a/mlir/include/mlir/Dialect/Affine/Passes.h
+++ b/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

diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td
index 4359ea0..ace2726 100644 (file)
--- a/mlir/include/mlir/Dialect/Affine/Passes.td
+++ b/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"> {

diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
similarity index 73%
rename from mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp
rename to mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
index 053b97f..2ad403e 100644 (file)
--- a/mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
@@ -1,4 +1,4 @@
-//===- 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.
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file implements a normalizer for affine parallel loops.
+// This file implements a normalizer for affine loop-like ops.
 //
 //===----------------------------------------------------------------------===//
 
@@ -14,11 +14,13 @@
 #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>();
 }

diff --git a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
index 899f362..d866a8f 100644 (file)
--- a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
+++ b/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

diff --git a/mlir/test/Dialect/Affine/affine-parallel-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
similarity index 64%
rename from mlir/test/Dialect/Affine/affine-parallel-normalize.mlir
rename to mlir/test/Dialect/Affine/affine-loop-normalize.mlir
index 4473283..ab8a0b7 100644 (file)
--- a/mlir/test/Dialect/Affine/affine-parallel-normalize.mlir
+++ b/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