//===- Utils.cpp ---- Utilities for affine dialect transformation ---------===// // // 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 implements miscellaneous transformation utilities for the Affine // dialect. // //===----------------------------------------------------------------------===// #include "mlir/Dialect/Affine/Utils.h" #include "mlir/Dialect/Affine/Analysis/Utils.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineValueMap.h" #include "mlir/Dialect/Affine/LoopUtils.h" #include "mlir/Dialect/Arith/Utils/Utils.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Utils/IndexingUtils.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/Dominance.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/ImplicitLocOpBuilder.h" #include "mlir/IR/IntegerSet.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "llvm/Support/LogicalResult.h" #include <optional> #define DEBUG_TYPE … usingnamespacemlir; usingnamespaceaffine; usingnamespacepresburger; namespace { /// Visit affine expressions recursively and build the sequence of operations /// that correspond to it. Visitation functions return an Value of the /// expression subtree they visited or `nullptr` on error. class AffineApplyExpander : public AffineExprVisitor<AffineApplyExpander, Value> { … }; } // namespace /// Create a sequence of operations that implement the `expr` applied to the /// given dimension and symbol values. mlir::Value mlir::affine::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>> mlir::affine::expandAffineMap(OpBuilder &builder, Location loc, AffineMap affineMap, ValueRange operands) { … } /// Promotes the `then` or the `else` block of `ifOp` (depending on whether /// `elseBlock` is false or true) into `ifOp`'s containing block, and discards /// the rest of the op. static void promoteIfBlock(AffineIfOp ifOp, bool elseBlock) { … } /// Returns the outermost affine.for/parallel op that the `ifOp` is invariant /// on. The `ifOp` could be hoisted and placed right before such an operation. /// This method assumes that the ifOp has been canonicalized (to be correct and /// effective). static Operation *getOutermostInvariantForOp(AffineIfOp ifOp) { … } /// A helper for the mechanics of mlir::hoistAffineIfOp. Hoists `ifOp` just over /// `hoistOverOp`. Returns the new hoisted op if any hoisting happened, /// otherwise the same `ifOp`. static AffineIfOp hoistAffineIfOp(AffineIfOp ifOp, Operation *hoistOverOp) { … } LogicalResult mlir::affine::affineParallelize(AffineForOp forOp, ArrayRef<LoopReduction> parallelReductions, AffineParallelOp *resOp) { … } // Returns success if any hoisting happened. LogicalResult mlir::affine::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) { … } // Return the min expr after replacing the given dim. AffineExpr mlir::affine::substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min, AffineExpr max, bool positivePath) { … } void mlir::affine::normalizeAffineParallel(AffineParallelOp op) { … } LogicalResult mlir::affine::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) { … } /// Returns true if the memory operation of `destAccess` depends on `srcAccess` /// inside of the innermost common surrounding affine loop between the two /// accesses. static bool mustReachAtInnermost(const MemRefAccess &srcAccess, const MemRefAccess &destAccess) { … } /// Returns true if `srcMemOp` may have an effect on `destMemOp` within the /// scope of the outermost `minSurroundingLoops` loops that surround them. /// `srcMemOp` and `destMemOp` are expected to be affine read/write ops. static bool mayHaveEffect(Operation *srcMemOp, Operation *destMemOp, unsigned minSurroundingLoops) { … } template <typename EffectType, typename T> bool mlir::affine::hasNoInterveningEffect( Operation *start, T memOp, llvm::function_ref<bool(Value, Value)> mayAlias) { … } /// Attempt to eliminate loadOp by replacing it with a value stored into memory /// which the load is guaranteed to retrieve. This check involves three /// components: 1) The store and load must be on the same location 2) The store /// must dominate (and therefore must always occur prior to) the load 3) No /// other operations will overwrite the memory loaded between the given load /// and store. If such a value exists, the replaced `loadOp` will be added to /// `loadOpsToErase` and its memref will be added to `memrefsToErase`. static void forwardStoreToLoad( AffineReadOpInterface loadOp, SmallVectorImpl<Operation *> &loadOpsToErase, SmallPtrSetImpl<Value> &memrefsToErase, DominanceInfo &domInfo, llvm::function_ref<bool(Value, Value)> mayAlias) { … } template bool mlir::affine::hasNoInterveningEffect<mlir::MemoryEffects::Read, affine::AffineReadOpInterface>( mlir::Operation *, affine::AffineReadOpInterface, llvm::function_ref<bool(Value, Value)>); // This attempts to find stores which have no impact on the final result. // A writing op writeA will be eliminated if there exists an op writeB if // 1) writeA and writeB have mathematically equivalent affine access functions. // 2) writeB postdominates writeA. // 3) There is no potential read between writeA and writeB. static void findUnusedStore(AffineWriteOpInterface writeA, SmallVectorImpl<Operation *> &opsToErase, PostDominanceInfo &postDominanceInfo, llvm::function_ref<bool(Value, Value)> mayAlias) { … } // The load to load forwarding / redundant load elimination is similar to the // store to load forwarding. // loadA will be be replaced with loadB if: // 1) loadA and loadB have mathematically equivalent affine access functions. // 2) loadB dominates loadA. // 3) There is no write between loadA and loadB. static void loadCSE(AffineReadOpInterface loadA, SmallVectorImpl<Operation *> &loadOpsToErase, DominanceInfo &domInfo, llvm::function_ref<bool(Value, Value)> mayAlias) { … } // The store to load forwarding and load CSE rely on three conditions: // // 1) store/load providing a replacement value and load being replaced need to // have mathematically equivalent affine access functions (checked after full // composition of load/store operands); this implies that they access the same // single memref element for all iterations of the common surrounding loop, // // 2) the store/load op should dominate the load op, // // 3) no operation that may write to memory read by the load being replaced can // occur after executing the instruction (load or store) providing the // replacement value and before the load being replaced (thus potentially // allowing overwriting the memory read by the load). // // The above conditions are simple to check, sufficient, and powerful for most // cases in practice - they are sufficient, but not necessary --- since they // don't reason about loops that are guaranteed to execute at least once or // multiple sources to forward from. // // TODO: more forwarding can be done when support for // loop/conditional live-out SSA values is available. // TODO: do general dead store elimination for memref's. This pass // currently only eliminates the stores only if no other loads/uses (other // than dealloc) remain. // void mlir::affine::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo, PostDominanceInfo &postDomInfo, AliasAnalysis &aliasAnalysis) { … } // Private helper function to transform memref.load with reduced rank. // This function will modify the indices of the memref.load to match the // newMemRef. LogicalResult transformMemRefLoadWithReducedRank( Operation *op, Value oldMemRef, Value newMemRef, unsigned memRefOperandPos, ArrayRef<Value> extraIndices, ArrayRef<Value> extraOperands, ArrayRef<Value> symbolOperands, AffineMap indexRemap) { … } // Perform the replacement in `op`. LogicalResult mlir::affine::replaceAllMemRefUsesWith( Value oldMemRef, Value newMemRef, Operation *op, ArrayRef<Value> extraIndices, AffineMap indexRemap, ArrayRef<Value> extraOperands, ArrayRef<Value> symbolOperands, bool allowNonDereferencingOps) { … } LogicalResult mlir::affine::replaceAllMemRefUsesWith( Value oldMemRef, Value newMemRef, ArrayRef<Value> extraIndices, AffineMap indexRemap, ArrayRef<Value> extraOperands, ArrayRef<Value> symbolOperands, Operation *domOpFilter, Operation *postDomOpFilter, bool allowNonDereferencingOps, bool replaceInDeallocOp) { … } /// Given an operation, inserts one or more single result affine /// apply operations, results of which are exclusively used by this operation /// 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"(%idx, %A, ...) /// "compute"(%idx) /// /// After /// /// affine.for %i = 0 to #map(%N) /// %idx = affine.apply (d0) -> (d0 mod 2) (%i) /// "send"(%idx, %A, ...) /// %idx_ = affine.apply (d0) -> (d0 mod 2) (%i) /// "compute"(%idx_) /// /// This allows applying different transformations on send and compute (for eg. /// different shifts/delays). /// /// Returns nullptr either if none of opInst's operands were the result of an /// affine.apply and thus there was no affine computation slice to create, or if /// all the affine.apply op's supplying operands to this opInst did not have any /// uses besides this opInst; otherwise returns the list of affine.apply /// operations created in output argument `sliceOps`. void mlir::affine::createAffineComputationSlice( Operation *opInst, SmallVectorImpl<AffineApplyOp> *sliceOps) { … } /// Enum to set patterns of affine expr in tiled-layout map. /// TileFloorDiv: <dim expr> div <tile size> /// TileMod: <dim expr> mod <tile size> /// TileNone: None of the above /// Example: /// #tiled_2d_128x256 = affine_map<(d0, d1) /// -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)> /// "d0 div 128" and "d1 div 256" ==> TileFloorDiv /// "d0 mod 128" and "d1 mod 256" ==> TileMod enum TileExprPattern { … }; /// Check if `map` is a tiled layout. In the tiled layout, specific k dimensions /// being floordiv'ed by respective tile sizes appeare in a mod with the same /// tile sizes, and no other expression involves those k dimensions. This /// function stores a vector of tuples (`tileSizePos`) including AffineExpr for /// tile size, positions of corresponding `floordiv` and `mod`. If it is not a /// tiled layout, an empty vector is returned. static LogicalResult getTileSizePos( AffineMap map, SmallVectorImpl<std::tuple<AffineExpr, unsigned, unsigned>> &tileSizePos) { … } /// Check if `dim` dimension of memrefType with `layoutMap` becomes dynamic /// after normalization. Dimensions that include dynamic dimensions in the map /// output will become dynamic dimensions. Return true if `dim` is dynamic /// dimension. /// /// Example: /// #map0 = affine_map<(d0, d1) -> (d0, d1 floordiv 32, d1 mod 32)> /// /// If d1 is dynamic dimension, 2nd and 3rd dimension of map output are dynamic. /// memref<4x?xf32, #map0> ==> memref<4x?x?xf32> static bool isNormalizedMemRefDynamicDim(unsigned dim, AffineMap layoutMap, SmallVectorImpl<unsigned> &inMemrefTypeDynDims) { … } /// Create affine expr to calculate dimension size for a tiled-layout map. static AffineExpr createDimSizeExprForTiledLayout(AffineExpr oldMapOutput, TileExprPattern pat) { … } /// Create new maps to calculate each dimension size of `newMemRefType`, and /// create `newDynamicSizes` from them by using AffineApplyOp. /// /// Steps for normalizing dynamic memrefs for a tiled layout map /// Example: /// #map0 = affine_map<(d0, d1) -> (d0, d1 floordiv 32, d1 mod 32)> /// %0 = dim %arg0, %c1 :memref<4x?xf32> /// %1 = alloc(%0) : memref<4x?xf32, #map0> /// /// (Before this function) /// 1. Check if `map`(#map0) is a tiled layout using `getTileSizePos()`. Only /// single layout map is supported. /// /// 2. Create normalized memrefType using `isNormalizedMemRefDynamicDim()`. It /// is memref<4x?x?xf32> in the above example. /// /// (In this function) /// 3. Create new maps to calculate each dimension of the normalized memrefType /// using `createDimSizeExprForTiledLayout()`. In the tiled layout, the /// dimension size can be calculated by replacing "floordiv <tile size>" with /// "ceildiv <tile size>" and "mod <tile size>" with "<tile size>". /// - New map in the above example /// #map0 = affine_map<(d0, d1) -> (d0)> /// #map1 = affine_map<(d0, d1) -> (d1 ceildiv 32)> /// #map2 = affine_map<(d0, d1) -> (32)> /// /// 4. Create AffineApplyOp to apply the new maps. The output of AffineApplyOp /// is used in dynamicSizes of new AllocOp. /// %0 = dim %arg0, %c1 : memref<4x?xf32> /// %c4 = arith.constant 4 : index /// %1 = affine.apply #map1(%c4, %0) /// %2 = affine.apply #map2(%c4, %0) static void createNewDynamicSizes(MemRefType oldMemRefType, MemRefType newMemRefType, AffineMap map, memref::AllocOp *allocOp, OpBuilder b, SmallVectorImpl<Value> &newDynamicSizes) { … } // TODO: Currently works for static memrefs with a single layout map. LogicalResult mlir::affine::normalizeMemRef(memref::AllocOp *allocOp) { … } MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) { … } DivModValue mlir::affine::getDivMod(OpBuilder &b, Location loc, Value lhs, Value rhs) { … } /// Create IR that computes the product of all elements in the set. static FailureOr<OpFoldResult> getIndexProduct(OpBuilder &b, Location loc, ArrayRef<Value> set) { … } FailureOr<SmallVector<Value>> mlir::affine::delinearizeIndex(OpBuilder &b, Location loc, Value linearIndex, ArrayRef<Value> basis) { … } OpFoldResult mlir::affine::linearizeIndex(ArrayRef<OpFoldResult> multiIndex, ArrayRef<OpFoldResult> basis, ImplicitLocOpBuilder &builder) { … }
Codebase Browser
llvm