Utils.h

Go to the documentation of this file.

//===- Utils.h - Affine dialect utilities -----------------------*- 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 header file declares a set of utilities for the affine dialect ops.
//
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_DIALECT_AFFINE_UTILS_H
#define MLIR_DIALECT_AFFINE_UTILS_H
 
#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/IR/OpDefinition.h"
#include <optional>
 
namespace mlir {
class DominanceInfo;
class Operation;
class PostDominanceInfo;
class ImplicitLocOpBuilder;
 
namespace func {
class FuncOp;
} // namespace func
 
namespace memref {
class AllocOp;
class AllocaOp;
class ReinterpretCastOp;
} // namespace memref
 
namespace affine {
class AffineForOp;
class AffineIfOp;
class AffineParallelOp;
 
using ReductionLoopMap = DenseMap<Operation *, SmallVector<LoopReduction, 2>>;
 
/// Replaces a parallel affine.for op with a 1-d affine.parallel op. `forOp`'s
/// body is taken by the affine.parallel op and the former is erased.
/// (mlir::isLoopParallel can be used to detect a parallel affine.for op.) The
/// reductions specified in `parallelReductions` are also parallelized.
/// Parallelization will fail in the presence of loop iteration arguments that
/// are not listed in `parallelReductions`. `resOp` if non-null is set to the
/// newly created affine.parallel op.
LogicalResult affineParallelize(AffineForOp forOp,
                                ArrayRef<LoopReduction> parallelReductions = {},
                                AffineParallelOp *resOp = nullptr);
 
/// Hoists out affine.if/else to as high as possible, i.e., past all invariant
/// affine.fors/parallel's. Returns success if any hoisting happened; folded` is
/// set to true if the op was folded or erased. This hoisting could lead to
/// significant code expansion in some cases.
LogicalResult hoistAffineIfOp(AffineIfOp ifOp, bool *folded = nullptr);
 
/// Holds parameters to perform n-D vectorization on a single loop nest.
/// For example, for the following loop nest:
///
/// func @vec2d(%in: memref<64x128x512xf32>, %out: memref<64x128x512xf32>) {
///   affine.for %i0 = 0 to 64 {
///     affine.for %i1 = 0 to 128 {
///       affine.for %i2 = 0 to 512 {
///         %ld = affine.load %in[%i0, %i1, %i2] : memref<64x128x512xf32>
///         affine.store %ld, %out[%i0, %i1, %i2] : memref<64x128x512xf32>
///       }
///     }
///   }
///   return
/// }
///
/// and VectorizationStrategy = 'vectorSizes = {8, 4}', 'loopToVectorDim =
/// {{i1->0}, {i2->1}}', SuperVectorizer will generate:
///
///  func @vec2d(%arg0: memref<64x128x512xf32>, %arg1: memref<64x128x512xf32>) {
///    affine.for %arg2 = 0 to 64 {
///      affine.for %arg3 = 0 to 128 step 8 {
///        affine.for %arg4 = 0 to 512 step 4 {
///          %cst = arith.constant 0.000000e+00 : f32
///          %0 = vector.transfer_read %arg0[%arg2, %arg3, %arg4], %cst : ...
///          vector.transfer_write %0, %arg1[%arg2, %arg3, %arg4] : ...
///        }
///      }
///    }
///    return
///  }
// TODO: Hoist to a VectorizationStrategy.cpp when appropriate.
struct VectorizationStrategy {
  // Vectorization factors to apply to each target vector dimension.
  // Each factor will be applied to a different loop.
  SmallVector<int64_t, 8> vectorSizes;
  // Maps each AffineForOp vectorization candidate with its vector dimension.
  // The candidate will be vectorized using the vectorization factor in
  // 'vectorSizes' for that dimension.
  DenseMap<Operation *, unsigned> loopToVectorDim;
  // Maps loops that implement vectorizable reductions to the corresponding
  // reduction descriptors.
  ReductionLoopMap reductionLoops;
};
 
/// Vectorize affine loops that are children of parentOp (including itself)
void vectorizeChildAffineLoops(Operation *parentOp, bool vectorizeReductions,
                               ArrayRef<int64_t> vectorSizes,
                               ArrayRef<int64_t> fastestVaryingPattern);
 
/// Replace affine store and load accesses by scalars by forwarding stores to
/// loads and eliminate invariant affine loads; consequently, eliminate dead
/// allocs.
void affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
                         PostDominanceInfo &postDomInfo,
                         AliasAnalysis &analysis);
 
/// Vectorizes affine loops in 'loops' using the n-D vectorization factors in
/// 'vectorSizes'. By default, each vectorization factor is applied
/// inner-to-outer to the loops of each loop nest. 'fastestVaryingPattern' can
/// be optionally used to provide a different loop vectorization order.
/// If `reductionLoops` is not empty, the given reduction loops may be
/// vectorized along the reduction dimension.
/// TODO: Vectorizing reductions is supported only for 1-D vectorization.
void vectorizeAffineLoops(
    Operation *parentOp,
    llvm::DenseSet<Operation *, DenseMapInfo<Operation *>> &loops,
    ArrayRef<int64_t> vectorSizes, ArrayRef<int64_t> fastestVaryingPattern,
    const ReductionLoopMap &reductionLoops = ReductionLoopMap());
 
/// External utility to vectorize affine loops from a single loop nest using an
/// n-D vectorization strategy (see doc in VectorizationStrategy definition).
/// Loops are provided in a 2D vector container. The first dimension represents
/// the nesting level relative to the loops to be vectorized. The second
/// dimension contains the loops. This means that:
///   a) every loop in 'loops[i]' must have a parent loop in 'loops[i-1]',
///   b) a loop in 'loops[i]' may or may not have a child loop in 'loops[i+1]'.
///
/// For example, for the following loop nest:
///
///   func @vec2d(%in0: memref<64x128x512xf32>, %in1: memref<64x128x128xf32>,
///               %out0: memref<64x128x512xf32>,
///               %out1: memref<64x128x128xf32>) {
///     affine.for %i0 = 0 to 64 {
///       affine.for %i1 = 0 to 128 {
///         affine.for %i2 = 0 to 512 {
///           %ld = affine.load %in0[%i0, %i1, %i2] : memref<64x128x512xf32>
///           affine.store %ld, %out0[%i0, %i1, %i2] : memref<64x128x512xf32>
///         }
///         affine.for %i3 = 0 to 128 {
///           %ld = affine.load %in1[%i0, %i1, %i3] : memref<64x128x128xf32>
///           affine.store %ld, %out1[%i0, %i1, %i3] : memref<64x128x128xf32>
///         }
///       }
///     }
///     return
///   }
///
/// loops = {{%i0}, {%i2, %i3}}, to vectorize the outermost and the two
/// innermost loops;
/// loops = {{%i1}, {%i2, %i3}}, to vectorize the middle and the two innermost
/// loops;
/// loops = {{%i2}}, to vectorize only the first innermost loop;
/// loops = {{%i3}}, to vectorize only the second innermost loop;
/// loops = {{%i1}}, to vectorize only the middle loop.
LogicalResult
vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
                        const VectorizationStrategy &strategy);
 
/// Normalize a affine.parallel op so that lower bounds are 0 and steps are 1.
/// As currently implemented, this transformation cannot fail and will return
/// early if the op is already in a normalized form.
void normalizeAffineParallel(AffineParallelOp op);
 
/// Normalize an affine.for op. An affine.for op is normalized by converting the
/// lower bound to zero and loop step to one. The upper bound is set to the trip
/// count of the loop. Original loops must have a lower bound with only a single
/// result. There is no such restriction on upper bounds. Returns success if the
/// loop has been normalized (or is already in the normal form). If
/// `promoteSingleIter` is true, the loop is simply promoted if it has a single
/// iteration.
LogicalResult normalizeAffineFor(AffineForOp op,
                                 bool promoteSingleIter = false);
 
/// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
/// been replaced by either:
///  - `min` if `positivePath` is true when we reach an occurrence of `dim`
///  - `max` if `positivePath` is true when we reach an occurrence of `dim`
/// `positivePath` is negated each time we hit a multiplicative or divisive
/// binary op with a constant negative coefficient.
AffineExpr substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
                        AffineExpr max, bool positivePath = true);
 
/// Replaces all "dereferencing" uses of `oldMemRef` with `newMemRef` while
/// optionally remapping the old memref's indices using the supplied affine map,
/// `indexRemap`. The new memref could be of a different shape or rank.
/// `extraIndices` provides any additional access indices to be added to the
/// start.
///
/// `indexRemap` remaps indices of the old memref access to a new set of indices
/// that are used to index the memref. Additional input operands to indexRemap
/// can be optionally provided in `extraOperands`, and they occupy the start
/// of its input list. `indexRemap`'s dimensional inputs are expected to
/// correspond to memref's indices, and its symbolic inputs if any should be
/// provided in `symbolOperands`.
//
/// If `userFilterFn` is specified, restrict replacement to only those users
/// that pass the specified filter (i.e., the filter returns true).
///
/// 'allowNonDereferencingOps', if set, allows replacement of non-dereferencing
/// uses of a memref without any requirement for access index rewrites as long
/// as the user operation has the MemRefsNormalizable trait. The default value
/// of this flag is false.
///
/// 'replaceInDeallocOp', if set, lets DeallocOp, a non-dereferencing user, to
/// also be a candidate for replacement. The default value of this flag is
/// false.
///
/// Returns true on success and false if the replacement is not possible,
/// whenever a memref is used as an operand in a non-dereferencing context and
/// 'allowNonDereferencingOps' is false, except for dealloc's on the memref
/// which are left untouched. See comments at function definition for an
/// example.
//
//  Ex: to replace load %A[%i, %j] with load %Abuf[%t mod 2, %ii - %i, %j]:
//  The SSA value corresponding to '%t mod 2' should be in 'extraIndices', and
//  index remap will perform (%i, %j) -> (%ii - %i, %j), i.e., indexRemap = (d0,
//  d1, d2) -> (d0 - d1, d2), and %ii will be the extra operand. Without any
//  extra operands, note that 'indexRemap' would just be applied to existing
//  indices (%i, %j).
//
//  TODO: allow extraIndices to be added at any position.
LogicalResult replaceAllMemRefUsesWith(
    Value oldMemRef, Value newMemRef, ArrayRef<Value> extraIndices = {},
    AffineMap indexRemap = AffineMap(), ArrayRef<Value> extraOperands = {},
    ArrayRef<Value> symbolOperands = {},
    llvm::function_ref<bool(Operation *)> userFilterFn = nullptr,
    bool allowNonDereferencingOps = false, bool replaceInDeallocOp = false);
 
/// Performs the same replacement as the other version above but only for the
/// dereferencing uses of `oldMemRef` in `op`, except in cases where
/// 'allowNonDereferencingOps' is set to true where we replace the
/// non-dereferencing uses as well.
LogicalResult replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                       Operation *op,
                                       ArrayRef<Value> extraIndices = {},
                                       AffineMap indexRemap = AffineMap(),
                                       ArrayRef<Value> extraOperands = {},
                                       ArrayRef<Value> symbolOperands = {},
                                       bool allowNonDereferencingOps = false);
 
/// Rewrites the memref defined by alloc or reinterpret_cast op to have an
/// identity layout map and updates all its indexing uses. Returns failure if
/// any of its uses escape (while leaving the IR in a valid state).
template <typename AllocLikeOp>
LogicalResult normalizeMemRef(AllocLikeOp op);
extern template LogicalResult
normalizeMemRef<memref::AllocaOp>(memref::AllocaOp op);
extern template LogicalResult
normalizeMemRef<memref::AllocOp>(memref::AllocOp op);
LogicalResult normalizeMemRef(memref::ReinterpretCastOp op);
 
/// Normalizes `memrefType` so that the affine layout map of the memref is
/// transformed to an identity map with a new shape being computed for the
/// normalized memref type and returns it. The old memref type is simplify
/// returned if the normalization failed.
MemRefType normalizeMemRefType(MemRefType memrefType);
 
/// Given an operation, inserts one or more single result affine apply
/// operations, results of which are exclusively used by this operation.
/// The operands of these newly created affine apply ops are
/// guaranteed to be loop iterators or terminal symbols of a function.
///
/// Before
///
/// affine.for %i = 0 to #map(%N)
///   %idx = affine.apply (d0) -> (d0 mod 2) (%i)
///   send %A[%idx], ...
///   %v = "compute"(%idx, ...)
///
/// After
///
/// affine.for %i = 0 to #map(%N)
///   %idx = affine.apply (d0) -> (d0 mod 2) (%i)
///   send %A[%idx], ...
///   %idx_ = affine.apply (d0) -> (d0 mod 2) (%i)
///   %v = "compute"(%idx_, ...)
 
/// This allows the application of different transformations on send and
/// compute (for eg. different shifts/delays)
///
/// Fills `sliceOps` with the list of affine.apply operations.
/// In the following cases, `sliceOps` remains empty:
///   1. If none of opInst's operands were the result of an affine.apply
///      (i.e., there was no affine computation slice to create).
///   2. If all the affine.apply op's supplying operands to this opInst did not
///      have any uses other than those in this opInst.
void createAffineComputationSlice(Operation *opInst,
                                  SmallVectorImpl<AffineApplyOp> *sliceOps);
 
/// Emit code that computes the given affine expression using standard
/// arithmetic operations applied to the provided dimension and symbol values.
Value expandAffineExpr(OpBuilder &builder, Location loc, AffineExpr expr,
                       ValueRange dimValues, ValueRange symbolValues);
 
/// Create a sequence of operations that implement the `affineMap` applied to
/// the given `operands` (as it it were an AffineApplyOp).
std::optional<SmallVector<Value, 8>> expandAffineMap(OpBuilder &builder,
                                                     Location loc,
                                                     AffineMap affineMap,
                                                     ValueRange operands);
 
/// Holds the result of (div a, b)  and (mod a, b).
struct DivModValue {
  Value quotient;
  Value remainder;
};
 
/// Create IR to calculate (div lhs, rhs) and (mod lhs, rhs).
DivModValue getDivMod(OpBuilder &b, Location loc, Value lhs, Value rhs);
 
/// Generate the IR to delinearize `linearIndex` given the `basis` and return
/// the multi-index. `hasOuterBound` indicates whether `basis` has an entry
/// given the size of the first multi-index result - if it is true, the function
/// will return `basis.size()` values, otherwise, it will return `basis.size() +
/// 1`.
FailureOr<SmallVector<Value>> delinearizeIndex(OpBuilder &b, Location loc,
                                               Value linearIndex,
                                               ArrayRef<Value> basis,
                                               bool hasOuterBound = true);
 
FailureOr<SmallVector<Value>> delinearizeIndex(OpBuilder &b, Location loc,
                                               Value linearIndex,
                                               ArrayRef<OpFoldResult> basis,
                                               bool hasOuterBound = true);
 
// Generate IR that extracts the linear index from a multi-index according to
// a basis/shape. The basis may contain either `multiIndex.size()` or
// `multiIndex.size() - 1` elements.
OpFoldResult linearizeIndex(ArrayRef<OpFoldResult> multiIndex,
                            ArrayRef<OpFoldResult> basis,
                            ImplicitLocOpBuilder &builder);
 
OpFoldResult linearizeIndex(OpBuilder &builder, Location loc,
                            ArrayRef<OpFoldResult> multiIndex,
                            ArrayRef<OpFoldResult> basis);
 
/// Ensure that all operations that could be executed after `start`
/// (noninclusive) and prior to `memOp` (e.g. on a control flow/op path
/// between the operations) do not have the potential memory effect
/// `EffectType` on `memOp`. `memOp`  is an operation that reads or writes to
/// a memref. For example, if `EffectType` is MemoryEffects::Write, this method
/// will check if there is no write to the memory between `start` and `memOp`
/// that would change the read within `memOp`.
template <typename EffectType, typename T>
bool hasNoInterveningEffect(Operation *start, T memOp,
                            llvm::function_ref<bool(Value, Value)> mayAlias);
 
struct AffineValueExpr {
  explicit AffineValueExpr(AffineExpr e) : e(e) {}
  AffineValueExpr bind(Value v) {
    this->v = v;
    return *this;
  }
  AffineValueExpr bind(OpFoldResult v) {
    this->v = v;
    return *this;
  }
  operator AffineExpr() const { return e; }
  operator OpFoldResult() const { return v; }
  AffineExpr e;
  OpFoldResult v;
};
 
/// Helper struct to build simple AffineValueExprs with minimal type inference
/// support.
struct AffineBuilder {
  AffineBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
  OpFoldResult add(AffineValueExpr lhs, AffineValueExpr rhs) {
    return makeComposedFoldedAffineApply(b, loc, {lhs.e + rhs.e}, {lhs, rhs});
  }
  OpFoldResult sub(AffineValueExpr lhs, AffineValueExpr rhs) {
    return makeComposedFoldedAffineApply(b, loc, {lhs.e - rhs.e}, {lhs, rhs});
  }
  OpFoldResult mul(AffineValueExpr lhs, AffineValueExpr rhs) {
    return makeComposedFoldedAffineApply(b, loc, {lhs.e * rhs.e}, {lhs, rhs});
  }
  OpFoldResult floor(AffineValueExpr lhs, AffineValueExpr rhs) {
    return makeComposedFoldedAffineApply(b, loc, {lhs.e.floorDiv(rhs.e)},
                                         {lhs, rhs});
  }
  OpFoldResult ceil(AffineValueExpr lhs, AffineValueExpr rhs) {
    return makeComposedFoldedAffineApply(b, loc, {lhs.e.ceilDiv(rhs.e)},
                                         {lhs, rhs});
  }
  OpFoldResult min(ArrayRef<OpFoldResult> vals) {
    return makeComposedFoldedAffineMin(
        b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), b.getContext()),
        vals);
  }
  OpFoldResult max(ArrayRef<OpFoldResult> vals) {
    return makeComposedFoldedAffineMax(
        b, loc, AffineMap::getMultiDimIdentityMap(vals.size(), b.getContext()),
        vals);
  }
 
private:
  OpBuilder &b;
  Location loc;
};
 
} // namespace affine
} // namespace mlir
 
#endif // MLIR_DIALECT_AFFINE_UTILS_H