llvm/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

//===- LoopSpecialization.cpp - scf.parallel/SCR.for specialization -------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Specializes parallel loops and for loops for easier unrolling and
// vectorization.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/SCF/Transforms/Passes.h"

#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/DenseMap.h"

namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPPEELING
#define GEN_PASS_DEF_SCFFORLOOPSPECIALIZATION
#define GEN_PASS_DEF_SCFPARALLELLOOPSPECIALIZATION
#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
} // namespace mlir

usingnamespacemlir;
usingnamespacemlir::affine;
ForOp;
ParallelOp;

/// Rewrite a parallel loop with bounds defined by an affine.min with a constant
/// into 2 loops after checking if the bounds are equal to that constant. This
/// is beneficial if the loop will almost always have the constant bound and
/// that version can be fully unrolled and vectorized.
static void specializeParallelLoopForUnrolling(ParallelOp op) { … }

/// Rewrite a for loop with bounds defined by an affine.min with a constant into
/// 2 loops after checking if the bounds are equal to that constant. This is
/// beneficial if the loop will almost always have the constant bound and that
/// version can be fully unrolled and vectorized.
static void specializeForLoopForUnrolling(ForOp op) { … }

/// Rewrite a for loop with bounds/step that potentially do not divide evenly
/// into a for loop where the step divides the iteration space evenly, followed
/// by an scf.if for the last (partial) iteration (if any).
///
/// This function rewrites the given scf.for loop in-place and creates a new
/// scf.if operation for the last iteration. It replaces all uses of the
/// unpeeled loop with the results of the newly generated scf.if.
///
/// The newly generated scf.if operation is returned via `ifOp`. The boundary
/// at which the loop is split (new upper bound) is returned via `splitBound`.
/// The return value indicates whether the loop was rewritten or not.
static LogicalResult peelForLoop(RewriterBase &b, ForOp forOp,
                                 ForOp &partialIteration, Value &splitBound) { … }

static void rewriteAffineOpAfterPeeling(RewriterBase &rewriter, ForOp forOp,
                                        ForOp partialIteration,
                                        Value previousUb) { … }

LogicalResult mlir::scf::peelForLoopAndSimplifyBounds(RewriterBase &rewriter,
                                                      ForOp forOp,
                                                      ForOp &partialIteration) { … }

/// When the `peelFront` option is set as true, the first iteration of the loop
/// is peeled off. This function rewrites the original scf::ForOp as two
/// scf::ForOp Ops, the first scf::ForOp corresponds to the first iteration of
/// the loop which can be canonicalized away in the following optimization. The
/// second loop Op contains the remaining iteration, and the new lower bound is
/// the original lower bound plus the number of steps.
LogicalResult mlir::scf::peelForLoopFirstIteration(RewriterBase &b, ForOp forOp,
                                                   ForOp &firstIteration) { … }

static constexpr char kPeeledLoopLabel[] = …;
static constexpr char kPartialIterationLabel[] = …;

namespace {
struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> { … };
} // namespace

namespace {
struct ParallelLoopSpecialization
    : public impl::SCFParallelLoopSpecializationBase<
          ParallelLoopSpecialization> { … };

struct ForLoopSpecialization
    : public impl::SCFForLoopSpecializationBase<ForLoopSpecialization> { … };

struct ForLoopPeeling : public impl::SCFForLoopPeelingBase<ForLoopPeeling> { … };
} // namespace

std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() { … }

std::unique_ptr<Pass> mlir::createForLoopSpecializationPass() { … }

std::unique_ptr<Pass> mlir::createForLoopPeelingPass() { … }