SCF.h

Go to the documentation of this file.

//===- SCFOps.h - Structured Control Flow -----------------------*- C++ -*-===//
//
// 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 defines structured control flow operations.
//
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_DIALECT_SCF_SCF_H
#define MLIR_DIALECT_SCF_SCF_H
 
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
#include "mlir/Dialect/SCF/IR/SCFDialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
 
namespace mlir {
namespace scf {
void buildTerminatedBody(OpBuilder &builder, Location loc);
} // namespace scf
} // namespace mlir
 
#define GET_OP_CLASSES
#include "mlir/Dialect/SCF/IR/SCFOps.h.inc"
 
namespace mlir {
namespace scf {
 
/// Returns the loop parent of an induction variable. If the provided value is
/// not an induction variable, then return nullptr.
ForOp getForInductionVarOwner(Value val);
 
/// Returns the parallel loop parent of an induction variable. If the provided
/// value is not an induction variable, then return nullptr.
ParallelOp getParallelForInductionVarOwner(Value val);
 
/// Returns the ForallOp parent of an thread index variable.
/// If the provided value is not a thread index variable, then return nullptr.
ForallOp getForallOpThreadIndexOwner(Value val);
 
/// Return true if ops a and b (or their ancestors) are in mutually exclusive
/// regions/blocks of an IfOp.
// TODO: Consider moving this functionality to RegionBranchOpInterface.
bool insideMutuallyExclusiveBranches(Operation *a, Operation *b);
 
/// Promotes the loop body of a scf::ForallOp to its containing block.
/// Only tensor.parallel_insert_slice combining ops are supported.
/// Returns failure without modifying the IR if a combining op is unsupported.
LogicalResult promote(RewriterBase &rewriter, scf::ForallOp forallOp);
 
/// An owning vector of values, handy to return from functions.
using ValueVector = SmallVector<Value>;
using LoopVector = SmallVector<scf::ForOp>;
struct LoopNest {
  LoopVector loops;
  ValueVector results;
};
 
/// Creates a perfect nest of "for" loops, i.e. all loops but the innermost
/// contain only another loop and a terminator. The lower, upper bounds and
/// steps are provided as `lbs`, `ubs` and `steps`, which are expected to be of
/// the same size. `iterArgs` points to the initial values of the loop iteration
/// arguments, which will be forwarded through the nest to the innermost loop.
/// The body of the loop is populated using `bodyBuilder`, which accepts an
/// ordered list of induction variables of all loops, followed by a list of
/// iteration arguments of the innermost loop, in the same order as provided to
/// `iterArgs`. This function is expected to return as many values as
/// `iterArgs`, of the same type and in the same order, that will be treated as
/// yielded from the loop body and forwarded back through the loop nest. If the
/// function is not provided, the loop nest is not expected to have iteration
/// arguments, the body of the innermost loop will be left empty, containing
/// only the zero-operand terminator. Returns the LoopNest containing the list
/// of perfectly nest scf::ForOp build during the call.
/// If bound arrays are empty, the body builder will be called
/// once to construct the IR outside of the loop with an empty list of induction
/// variables.
LoopNest buildLoopNest(
    OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
    ValueRange steps, ValueRange iterArgs,
    function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
        bodyBuilder = nullptr);
 
/// A convenience version for building loop nests without iteration arguments
/// (like for reductions). Does not take the initial value of reductions or
/// expect the body building functions to return their current value.
/// The built nested scf::For are captured in `capturedLoops` when non-null.
LoopNest buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
                       ValueRange ubs, ValueRange steps,
                       function_ref<void(OpBuilder &, Location, ValueRange)>
                           bodyBuilder = nullptr);
 
/// Perform a replacement of one iter OpOperand of an scf.for to the
/// `replacement` value with a different type. A callback is used to insert
/// cast ops inside the block to account for type differences.
using ValueTypeCastFnTy =
    llvm::function_ref<Value(OpBuilder &, Location loc, Type, Value)>;
SmallVector<Value> replaceAndCastForOpIterArg(RewriterBase &rewriter,
                                              scf::ForOp forOp,
                                              OpOperand &operand,
                                              Value replacement,
                                              const ValueTypeCastFnTy &castFn);
 
/// A simple, conservative analysis to determine if the loop is shape
/// preserving. I.e., the arg-th yielded value has the same shape as the
/// corresponding region iter_arg: it is the iter_arg, or derives from it
/// through tied results of destination-style ops and of shape-preserving loops.
bool isShapePreserving(ForOp forOp, int64_t arg);
 
/// Helper function to compute the difference between two values. This is used
/// by the loop implementations to compute the trip count.
std::optional<llvm::APSInt> computeUbMinusLb(Value lb, Value ub, bool isSigned);
 
} // namespace scf
} // namespace mlir
#endif // MLIR_DIALECT_SCF_SCF_H