diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h --- a/mlir/include/mlir/Dialect/Affine/Passes.h +++ b/mlir/include/mlir/Dialect/Affine/Passes.h @@ -35,8 +35,8 @@ /// ops. std::unique_ptr> createAffineParallelizePass(); -/// Normalize affine.parallel ops so that lower bounds are 0 and steps are 1. -std::unique_ptr> createAffineParallelNormalizePass(); +/// Apply normalization transformations to affine loop-like ops. +std::unique_ptr> 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 --- a/mlir/include/mlir/Dialect/Affine/Passes.td +++ b/mlir/include/mlir/Dialect/Affine/Passes.td @@ -120,10 +120,9 @@ 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 rename from mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp rename to mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp --- 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 steps = op.getSteps(); // No need to do any work if the parallel op is already normalized. @@ -77,20 +79,36 @@ 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 { +struct AffineLoopNormalizePass + : public AffineLoopNormalizeBase { - void runOnFunction() override { getFunction().walk(normalizeAffineParallel); } + void runOnFunction() override { + getFunction().walk([](Operation *op) { + if (auto affineParallel = dyn_cast(op)) + normalizeAffineParallel(affineParallel); + else if (auto affineFor = dyn_cast(op)) + normalizeAffineFor(affineFor); + }); + } }; } // namespace -std::unique_ptr> -mlir::createAffineParallelNormalizePass() { - return std::make_unique(); +std::unique_ptr> mlir::createAffineLoopNormalizePass() { + return std::make_unique(); } diff --git a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt --- 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 rename from mlir/test/Dialect/Affine/affine-parallel-normalize.mlir rename to mlir/test/Dialect/Affine/affine-loop-normalize.mlir --- 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 @@ } 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