//===- AffineOps.cpp - MLIR Affine Operations -----------------------------===// // // 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 // //===----------------------------------------------------------------------===// #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineValueMap.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/UB/IR/UBOps.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/IntegerSet.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/OpDefinition.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Interfaces/ShapedOpInterfaces.h" #include "mlir/Interfaces/ValueBoundsOpInterface.h" #include "mlir/Transforms/InliningUtils.h" #include "llvm/ADT/ScopeExit.h" #include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/SmallVectorExtras.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" #include "llvm/Support/MathExtras.h" #include <numeric> #include <optional> usingnamespacemlir; usingnamespacemlir::affine; divideCeilSigned; divideFloorSigned; mod; #define DEBUG_TYPE … #include "mlir/Dialect/Affine/IR/AffineOpsDialect.cpp.inc" /// A utility function to check if a value is defined at the top level of /// `region` or is an argument of `region`. A value of index type defined at the /// top level of a `AffineScope` region is always a valid symbol for all /// uses in that region. bool mlir::affine::isTopLevelValue(Value value, Region *region) { … } /// Checks if `value` known to be a legal affine dimension or symbol in `src` /// region remains legal if the operation that uses it is inlined into `dest` /// with the given value mapping. `legalityCheck` is either `isValidDim` or /// `isValidSymbol`, depending on the value being required to remain a valid /// dimension or symbol. static bool remainsLegalAfterInline(Value value, Region *src, Region *dest, const IRMapping &mapping, function_ref<bool(Value, Region *)> legalityCheck) { … } /// Checks if all values known to be legal affine dimensions or symbols in `src` /// remain so if their respective users are inlined into `dest`. static bool remainsLegalAfterInline(ValueRange values, Region *src, Region *dest, const IRMapping &mapping, function_ref<bool(Value, Region *)> legalityCheck) { … } /// Checks if an affine read or write operation remains legal after inlining /// from `src` to `dest`. template <typename OpTy> static bool remainsLegalAfterInline(OpTy op, Region *src, Region *dest, const IRMapping &mapping) { … } /// Checks if an affine apply operation remains legal after inlining from `src` /// to `dest`. // Use "unused attribute" marker to silence clang-tidy warning stemming from // the inability to see through "llvm::TypeSwitch". template <> bool LLVM_ATTRIBUTE_UNUSED remainsLegalAfterInline(AffineApplyOp op, Region *src, Region *dest, const IRMapping &mapping) { … } //===----------------------------------------------------------------------===// // AffineDialect Interfaces //===----------------------------------------------------------------------===// namespace { /// This class defines the interface for handling inlining with affine /// operations. struct AffineInlinerInterface : public DialectInlinerInterface { … }; } // namespace //===----------------------------------------------------------------------===// // AffineDialect //===----------------------------------------------------------------------===// void AffineDialect::initialize() { … } /// Materialize a single constant operation from a given attribute value with /// the desired resultant type. Operation *AffineDialect::materializeConstant(OpBuilder &builder, Attribute value, Type type, Location loc) { … } /// A utility function to check if a value is defined at the top level of an /// op with trait `AffineScope`. If the value is defined in an unlinked region, /// conservatively assume it is not top-level. A value of index type defined at /// the top level is always a valid symbol. bool mlir::affine::isTopLevelValue(Value value) { … } /// Returns the closest region enclosing `op` that is held by an operation with /// trait `AffineScope`; `nullptr` if there is no such region. Region *mlir::affine::getAffineScope(Operation *op) { … } // A Value can be used as a dimension id iff it meets one of the following // conditions: // *) It is valid as a symbol. // *) It is an induction variable. // *) It is the result of affine apply operation with dimension id arguments. bool mlir::affine::isValidDim(Value value) { … } // Value can be used as a dimension id iff it meets one of the following // conditions: // *) It is valid as a symbol. // *) It is an induction variable. // *) It is the result of an affine apply operation with dimension id operands. bool mlir::affine::isValidDim(Value value, Region *region) { … } /// Returns true if the 'index' dimension of the `memref` defined by /// `memrefDefOp` is a statically shaped one or defined using a valid symbol /// for `region`. template <typename AnyMemRefDefOp> static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index, Region *region) { … } /// Returns true if the result of the dim op is a valid symbol for `region`. static bool isDimOpValidSymbol(ShapedDimOpInterface dimOp, Region *region) { … } // A value can be used as a symbol (at all its use sites) iff it meets one of // the following conditions: // *) It is a constant. // *) Its defining op or block arg appearance is immediately enclosed by an op // with `AffineScope` trait. // *) It is the result of an affine.apply operation with symbol operands. // *) It is a result of the dim op on a memref whose corresponding size is a // valid symbol. bool mlir::affine::isValidSymbol(Value value) { … } /// A value can be used as a symbol for `region` iff it meets one of the /// following conditions: /// *) It is a constant. /// *) It is the result of an affine apply operation with symbol arguments. /// *) It is a result of the dim op on a memref whose corresponding size is /// a valid symbol. /// *) It is defined at the top level of 'region' or is its argument. /// *) It dominates `region`'s parent op. /// If `region` is null, conservatively assume the symbol definition scope does /// not exist and only accept the values that would be symbols regardless of /// the surrounding region structure, i.e. the first three cases above. bool mlir::affine::isValidSymbol(Value value, Region *region) { … } // Returns true if 'value' is a valid index to an affine operation (e.g. // affine.load, affine.store, affine.dma_start, affine.dma_wait) where // `region` provides the polyhedral symbol scope. Returns false otherwise. static bool isValidAffineIndexOperand(Value value, Region *region) { … } /// Prints dimension and symbol list. static void printDimAndSymbolList(Operation::operand_iterator begin, Operation::operand_iterator end, unsigned numDims, OpAsmPrinter &printer) { … } /// Parses dimension and symbol list and returns true if parsing failed. ParseResult mlir::affine::parseDimAndSymbolList( OpAsmParser &parser, SmallVectorImpl<Value> &operands, unsigned &numDims) { … } /// Utility function to verify that a set of operands are valid dimension and /// symbol identifiers. The operands should be laid out such that the dimension /// operands are before the symbol operands. This function returns failure if /// there was an invalid operand. An operation is provided to emit any necessary /// errors. template <typename OpTy> static LogicalResult verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands, unsigned numDims) { … } //===----------------------------------------------------------------------===// // AffineApplyOp //===----------------------------------------------------------------------===// AffineValueMap AffineApplyOp::getAffineValueMap() { … } ParseResult AffineApplyOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineApplyOp::print(OpAsmPrinter &p) { … } LogicalResult AffineApplyOp::verify() { … } // The result of the affine apply operation can be used as a dimension id if all // its operands are valid dimension ids. bool AffineApplyOp::isValidDim() { … } // The result of the affine apply operation can be used as a dimension id if all // its operands are valid dimension ids with the parent operation of `region` // defining the polyhedral scope for symbols. bool AffineApplyOp::isValidDim(Region *region) { … } // The result of the affine apply operation can be used as a symbol if all its // operands are symbols. bool AffineApplyOp::isValidSymbol() { … } // The result of the affine apply operation can be used as a symbol in `region` // if all its operands are symbols in `region`. bool AffineApplyOp::isValidSymbol(Region *region) { … } OpFoldResult AffineApplyOp::fold(FoldAdaptor adaptor) { … } /// Returns the largest known divisor of `e`. Exploits information from the /// values in `operands`. static int64_t getLargestKnownDivisor(AffineExpr e, ArrayRef<Value> operands) { … } /// Check if `e` is known to be: 0 <= `e` < `k`. Handles the simple cases of `e` /// being an affine dim expression or a constant. static bool isNonNegativeBoundedBy(AffineExpr e, ArrayRef<Value> operands, int64_t k) { … } /// Check if expression `e` is of the form d*e_1 + e_2 where 0 <= e_2 < d. /// Set `div` to `d`, `quotientTimesDiv` to e_1 and `rem` to e_2 if the /// expression is in that form. static bool isQTimesDPlusR(AffineExpr e, ArrayRef<Value> operands, int64_t &div, AffineExpr "ientTimesDiv, AffineExpr &rem) { … } /// Gets the constant lower bound on an `iv`. static std::optional<int64_t> getLowerBound(Value iv) { … } /// Gets the constant upper bound on an affine.for `iv`. static std::optional<int64_t> getUpperBound(Value iv) { … } /// Determine a constant upper bound for `expr` if one exists while exploiting /// values in `operands`. Note that the upper bound is an inclusive one. `expr` /// is guaranteed to be less than or equal to it. static std::optional<int64_t> getUpperBound(AffineExpr expr, unsigned numDims, unsigned numSymbols, ArrayRef<Value> operands) { … } /// Determine a constant lower bound for `expr` if one exists while exploiting /// values in `operands`. Note that the upper bound is an inclusive one. `expr` /// is guaranteed to be less than or equal to it. static std::optional<int64_t> getLowerBound(AffineExpr expr, unsigned numDims, unsigned numSymbols, ArrayRef<Value> operands) { … } /// Simplify `expr` while exploiting information from the values in `operands`. static void simplifyExprAndOperands(AffineExpr &expr, unsigned numDims, unsigned numSymbols, ArrayRef<Value> operands) { … } /// Simplify the expressions in `map` while making use of lower or upper bounds /// of its operands. If `isMax` is true, the map is to be treated as a max of /// its result expressions, and min otherwise. Eg: min (d0, d1) -> (8, 4 * d0 + /// d1) can be simplified to (8) if the operands are respectively lower bounded /// by 2 and 0 (the second expression can't be lower than 8). static void simplifyMinOrMaxExprWithOperands(AffineMap &map, ArrayRef<Value> operands, bool isMax) { … } /// Simplify the map while exploiting information on the values in `operands`. // Use "unused attribute" marker to silence warning stemming from the inability // to see through the template expansion. static void LLVM_ATTRIBUTE_UNUSED simplifyMapWithOperands(AffineMap &map, ArrayRef<Value> operands) { … } /// Replace all occurrences of AffineExpr at position `pos` in `map` by the /// defining AffineApplyOp expression and operands. /// When `dimOrSymbolPosition < dims.size()`, AffineDimExpr@[pos] is replaced. /// When `dimOrSymbolPosition >= dims.size()`, /// AffineSymbolExpr@[pos - dims.size()] is replaced. /// Mutate `map`,`dims` and `syms` in place as follows: /// 1. `dims` and `syms` are only appended to. /// 2. `map` dim and symbols are gradually shifted to higher positions. /// 3. Old `dim` and `sym` entries are replaced by nullptr /// This avoids the need for any bookkeeping. static LogicalResult replaceDimOrSym(AffineMap *map, unsigned dimOrSymbolPosition, SmallVectorImpl<Value> &dims, SmallVectorImpl<Value> &syms) { … } /// Iterate over `operands` and fold away all those produced by an AffineApplyOp /// iteratively. Perform canonicalization of map and operands as well as /// AffineMap simplification. `map` and `operands` are mutated in place. static void composeAffineMapAndOperands(AffineMap *map, SmallVectorImpl<Value> *operands) { … } void mlir::affine::fullyComposeAffineMapAndOperands( AffineMap *map, SmallVectorImpl<Value> *operands) { … } AffineApplyOp mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } AffineApplyOp mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e, ArrayRef<OpFoldResult> operands) { … } /// Composes the given affine map with the given list of operands, pulling in /// the maps from any affine.apply operations that supply the operands. static void composeMultiResultAffineMap(AffineMap &map, SmallVectorImpl<Value> &operands) { … } OpFoldResult mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } OpFoldResult mlir::affine::makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineExpr expr, ArrayRef<OpFoldResult> operands) { … } SmallVector<OpFoldResult> mlir::affine::makeComposedFoldedMultiResultAffineApply( OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } template <typename OpTy> static OpTy makeComposedMinMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } AffineMinOp mlir::affine::makeComposedAffineMin(OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } template <typename OpTy> static OpFoldResult makeComposedFoldedMinMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } OpFoldResult mlir::affine::makeComposedFoldedAffineMin(OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } OpFoldResult mlir::affine::makeComposedFoldedAffineMax(OpBuilder &b, Location loc, AffineMap map, ArrayRef<OpFoldResult> operands) { … } // A symbol may appear as a dim in affine.apply operations. This function // canonicalizes dims that are valid symbols into actual symbols. template <class MapOrSet> static void canonicalizePromotedSymbols(MapOrSet *mapOrSet, SmallVectorImpl<Value> *operands) { … } // Works for either an affine map or an integer set. template <class MapOrSet> static void canonicalizeMapOrSetAndOperands(MapOrSet *mapOrSet, SmallVectorImpl<Value> *operands) { … } void mlir::affine::canonicalizeMapAndOperands( AffineMap *map, SmallVectorImpl<Value> *operands) { … } void mlir::affine::canonicalizeSetAndOperands( IntegerSet *set, SmallVectorImpl<Value> *operands) { … } namespace { /// Simplify AffineApply, AffineLoad, and AffineStore operations by composing /// maps that supply results into them. /// template <typename AffineOpTy> struct SimplifyAffineOp : public OpRewritePattern<AffineOpTy> { … }; // Specialize the template to account for the different build signatures for // affine load, store, and apply ops. template <> void SimplifyAffineOp<AffineLoadOp>::replaceAffineOp( PatternRewriter &rewriter, AffineLoadOp load, AffineMap map, ArrayRef<Value> mapOperands) const { … } template <> void SimplifyAffineOp<AffinePrefetchOp>::replaceAffineOp( PatternRewriter &rewriter, AffinePrefetchOp prefetch, AffineMap map, ArrayRef<Value> mapOperands) const { … } template <> void SimplifyAffineOp<AffineStoreOp>::replaceAffineOp( PatternRewriter &rewriter, AffineStoreOp store, AffineMap map, ArrayRef<Value> mapOperands) const { … } template <> void SimplifyAffineOp<AffineVectorLoadOp>::replaceAffineOp( PatternRewriter &rewriter, AffineVectorLoadOp vectorload, AffineMap map, ArrayRef<Value> mapOperands) const { … } template <> void SimplifyAffineOp<AffineVectorStoreOp>::replaceAffineOp( PatternRewriter &rewriter, AffineVectorStoreOp vectorstore, AffineMap map, ArrayRef<Value> mapOperands) const { … } // Generic version for ops that don't have extra operands. template <typename AffineOpTy> void SimplifyAffineOp<AffineOpTy>::replaceAffineOp( PatternRewriter &rewriter, AffineOpTy op, AffineMap map, ArrayRef<Value> mapOperands) const { … } } // namespace void AffineApplyOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // AffineDmaStartOp //===----------------------------------------------------------------------===// // TODO: Check that map operands are loop IVs or symbols. void AffineDmaStartOp::build(OpBuilder &builder, OperationState &result, Value srcMemRef, AffineMap srcMap, ValueRange srcIndices, Value destMemRef, AffineMap dstMap, ValueRange destIndices, Value tagMemRef, AffineMap tagMap, ValueRange tagIndices, Value numElements, Value stride, Value elementsPerStride) { … } void AffineDmaStartOp::print(OpAsmPrinter &p) { … } // Parse AffineDmaStartOp. // Ex: // affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %size, // %stride, %num_elt_per_stride // : memref<3076 x f32, 0>, memref<1024 x f32, 2>, memref<1 x i32> // ParseResult AffineDmaStartOp::parse(OpAsmParser &parser, OperationState &result) { … } LogicalResult AffineDmaStartOp::verifyInvariantsImpl() { … } LogicalResult AffineDmaStartOp::fold(ArrayRef<Attribute> cstOperands, SmallVectorImpl<OpFoldResult> &results) { … } void AffineDmaStartOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } //===----------------------------------------------------------------------===// // AffineDmaWaitOp //===----------------------------------------------------------------------===// // TODO: Check that map operands are loop IVs or symbols. void AffineDmaWaitOp::build(OpBuilder &builder, OperationState &result, Value tagMemRef, AffineMap tagMap, ValueRange tagIndices, Value numElements) { … } void AffineDmaWaitOp::print(OpAsmPrinter &p) { … } // Parse AffineDmaWaitOp. // Eg: // affine.dma_wait %tag[%index], %num_elements // : memref<1 x i32, (d0) -> (d0), 4> // ParseResult AffineDmaWaitOp::parse(OpAsmParser &parser, OperationState &result) { … } LogicalResult AffineDmaWaitOp::verifyInvariantsImpl() { … } LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands, SmallVectorImpl<OpFoldResult> &results) { … } void AffineDmaWaitOp::getEffects( SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) { … } //===----------------------------------------------------------------------===// // AffineForOp //===----------------------------------------------------------------------===// /// 'bodyBuilder' is used to build the body of affine.for. If iterArgs and /// bodyBuilder are empty/null, we include default terminator op. void AffineForOp::build(OpBuilder &builder, OperationState &result, ValueRange lbOperands, AffineMap lbMap, ValueRange ubOperands, AffineMap ubMap, int64_t step, ValueRange iterArgs, BodyBuilderFn bodyBuilder) { … } void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb, int64_t ub, int64_t step, ValueRange iterArgs, BodyBuilderFn bodyBuilder) { … } LogicalResult AffineForOp::verifyRegions() { … } /// Parse a for operation loop bounds. static ParseResult parseBound(bool isLower, OperationState &result, OpAsmParser &p) { … } ParseResult AffineForOp::parse(OpAsmParser &parser, OperationState &result) { … } static void printBound(AffineMapAttr boundMap, Operation::operand_range boundOperands, const char *prefix, OpAsmPrinter &p) { … } unsigned AffineForOp::getNumIterOperands() { … } std::optional<MutableArrayRef<OpOperand>> AffineForOp::getYieldedValuesMutable() { … } void AffineForOp::print(OpAsmPrinter &p) { … } /// Fold the constant bounds of a loop. static LogicalResult foldLoopBounds(AffineForOp forOp) { … } /// Canonicalize the bounds of the given loop. static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) { … } namespace { /// Returns constant trip count in trivial cases. static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) { … } /// This is a pattern to fold trivially empty loop bodies. /// TODO: This should be moved into the folding hook. struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> { … }; } // namespace void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) { … } void AffineForOp::getSuccessorRegions( RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) { … } /// Returns true if the affine.for has zero iterations in trivial cases. static bool hasTrivialZeroTripCount(AffineForOp op) { … } LogicalResult AffineForOp::fold(FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) { … } AffineBound AffineForOp::getLowerBound() { … } AffineBound AffineForOp::getUpperBound() { … } void AffineForOp::setLowerBound(ValueRange lbOperands, AffineMap map) { … } void AffineForOp::setUpperBound(ValueRange ubOperands, AffineMap map) { … } bool AffineForOp::hasConstantLowerBound() { … } bool AffineForOp::hasConstantUpperBound() { … } int64_t AffineForOp::getConstantLowerBound() { … } int64_t AffineForOp::getConstantUpperBound() { … } void AffineForOp::setConstantLowerBound(int64_t value) { … } void AffineForOp::setConstantUpperBound(int64_t value) { … } AffineForOp::operand_range AffineForOp::getControlOperands() { … } bool AffineForOp::matchingBoundOperandList() { … } SmallVector<Region *> AffineForOp::getLoopRegions() { … } std::optional<SmallVector<Value>> AffineForOp::getLoopInductionVars() { … } std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopLowerBounds() { … } std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopSteps() { … } std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopUpperBounds() { … } FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields( RewriterBase &rewriter, ValueRange newInitOperands, bool replaceInitOperandUsesInLoop, const NewYieldValuesFn &newYieldValuesFn) { … } Speculation::Speculatability AffineForOp::getSpeculatability() { … } /// Returns true if the provided value is the induction variable of a /// AffineForOp. bool mlir::affine::isAffineForInductionVar(Value val) { … } bool mlir::affine::isAffineParallelInductionVar(Value val) { … } bool mlir::affine::isAffineInductionVar(Value val) { … } AffineForOp mlir::affine::getForInductionVarOwner(Value val) { … } AffineParallelOp mlir::affine::getAffineParallelInductionVarOwner(Value val) { … } /// Extracts the induction variables from a list of AffineForOps and returns /// them. void mlir::affine::extractForInductionVars(ArrayRef<AffineForOp> forInsts, SmallVectorImpl<Value> *ivs) { … } void mlir::affine::extractInductionVars(ArrayRef<mlir::Operation *> affineOps, SmallVectorImpl<mlir::Value> &ivs) { … } /// Builds an affine loop nest, using "loopCreatorFn" to create individual loop /// operations. template <typename BoundListTy, typename LoopCreatorTy> static void buildAffineLoopNestImpl( OpBuilder &builder, Location loc, BoundListTy lbs, BoundListTy ubs, ArrayRef<int64_t> steps, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn, LoopCreatorTy &&loopCreatorFn) { … } /// Creates an affine loop from the bounds known to be constants. static AffineForOp buildAffineLoopFromConstants(OpBuilder &builder, Location loc, int64_t lb, int64_t ub, int64_t step, AffineForOp::BodyBuilderFn bodyBuilderFn) { … } /// Creates an affine loop from the bounds that may or may not be constants. static AffineForOp buildAffineLoopFromValues(OpBuilder &builder, Location loc, Value lb, Value ub, int64_t step, AffineForOp::BodyBuilderFn bodyBuilderFn) { … } void mlir::affine::buildAffineLoopNest( OpBuilder &builder, Location loc, ArrayRef<int64_t> lbs, ArrayRef<int64_t> ubs, ArrayRef<int64_t> steps, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) { … } void mlir::affine::buildAffineLoopNest( OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs, ArrayRef<int64_t> steps, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) { … } //===----------------------------------------------------------------------===// // AffineIfOp //===----------------------------------------------------------------------===// namespace { /// Remove else blocks that have nothing other than a zero value yield. struct SimplifyDeadElse : public OpRewritePattern<AffineIfOp> { … }; /// Removes affine.if cond if the condition is always true or false in certain /// trivial cases. Promotes the then/else block in the parent operation block. struct AlwaysTrueOrFalseIf : public OpRewritePattern<AffineIfOp> { … }; } // namespace /// AffineIfOp has two regions -- `then` and `else`. The flow of data should be /// as follows: AffineIfOp -> `then`/`else` -> AffineIfOp void AffineIfOp::getSuccessorRegions( RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) { … } LogicalResult AffineIfOp::verify() { … } ParseResult AffineIfOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineIfOp::print(OpAsmPrinter &p) { … } IntegerSet AffineIfOp::getIntegerSet() { … } void AffineIfOp::setIntegerSet(IntegerSet newSet) { … } void AffineIfOp::setConditional(IntegerSet set, ValueRange operands) { … } void AffineIfOp::build(OpBuilder &builder, OperationState &result, TypeRange resultTypes, IntegerSet set, ValueRange args, bool withElseRegion) { … } void AffineIfOp::build(OpBuilder &builder, OperationState &result, IntegerSet set, ValueRange args, bool withElseRegion) { … } /// Compose any affine.apply ops feeding into `operands` of the integer set /// `set` by composing the maps of such affine.apply ops with the integer /// set constraints. static void composeSetAndOperands(IntegerSet &set, SmallVectorImpl<Value> &operands) { … } /// Canonicalize an affine if op's conditional (integer set + operands). LogicalResult AffineIfOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) { … } void AffineIfOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // AffineLoadOp //===----------------------------------------------------------------------===// void AffineLoadOp::build(OpBuilder &builder, OperationState &result, AffineMap map, ValueRange operands) { … } void AffineLoadOp::build(OpBuilder &builder, OperationState &result, Value memref, AffineMap map, ValueRange mapOperands) { … } void AffineLoadOp::build(OpBuilder &builder, OperationState &result, Value memref, ValueRange indices) { … } ParseResult AffineLoadOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineLoadOp::print(OpAsmPrinter &p) { … } /// Verify common indexing invariants of affine.load, affine.store, /// affine.vector_load and affine.vector_store. static LogicalResult verifyMemoryOpIndexing(Operation *op, AffineMapAttr mapAttr, Operation::operand_range mapOperands, MemRefType memrefType, unsigned numIndexOperands) { … } LogicalResult AffineLoadOp::verify() { … } void AffineLoadOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } OpFoldResult AffineLoadOp::fold(FoldAdaptor adaptor) { … } //===----------------------------------------------------------------------===// // AffineStoreOp //===----------------------------------------------------------------------===// void AffineStoreOp::build(OpBuilder &builder, OperationState &result, Value valueToStore, Value memref, AffineMap map, ValueRange mapOperands) { … } // Use identity map. void AffineStoreOp::build(OpBuilder &builder, OperationState &result, Value valueToStore, Value memref, ValueRange indices) { … } ParseResult AffineStoreOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineStoreOp::print(OpAsmPrinter &p) { … } LogicalResult AffineStoreOp::verify() { … } void AffineStoreOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } LogicalResult AffineStoreOp::fold(FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) { … } //===----------------------------------------------------------------------===// // AffineMinMaxOpBase //===----------------------------------------------------------------------===// template <typename T> static LogicalResult verifyAffineMinMaxOp(T op) { … } template <typename T> static void printAffineMinMaxOp(OpAsmPrinter &p, T op) { … } template <typename T> static ParseResult parseAffineMinMaxOp(OpAsmParser &parser, OperationState &result) { … } /// Fold an affine min or max operation with the given operands. The operand /// list may contain nulls, which are interpreted as the operand not being a /// constant. template <typename T> static OpFoldResult foldMinMaxOp(T op, ArrayRef<Attribute> operands) { … } /// Remove duplicated expressions in affine min/max ops. template <typename T> struct DeduplicateAffineMinMaxExpressions : public OpRewritePattern<T> { … }; /// Merge an affine min/max op to its consumers if its consumer is also an /// affine min/max op. /// /// This pattern requires the producer affine min/max op is bound to a /// dimension/symbol that is used as a standalone expression in the consumer /// affine op's map. /// /// For example, a pattern like the following: /// /// %0 = affine.min affine_map<()[s0] -> (s0 + 16, s0 * 8)> ()[%sym1] /// %1 = affine.min affine_map<(d0)[s0] -> (s0 + 4, d0)> (%0)[%sym2] /// /// Can be turned into: /// /// %1 = affine.min affine_map< /// ()[s0, s1] -> (s0 + 4, s1 + 16, s1 * 8)> ()[%sym2, %sym1] template <typename T> struct MergeAffineMinMaxOp : public OpRewritePattern<T> { … }; /// Canonicalize the result expression order of an affine map and return success /// if the order changed. /// /// The function flattens the map's affine expressions to coefficient arrays and /// sorts them in lexicographic order. A coefficient array contains a multiplier /// for every dimension/symbol and a constant term. The canonicalization fails /// if a result expression is not pure or if the flattening requires local /// variables that, unlike dimensions and symbols, have no global order. static LogicalResult canonicalizeMapExprAndTermOrder(AffineMap &map) { … } /// Canonicalize the affine map result expression order of an affine min/max /// operation. /// /// The pattern calls `canonicalizeMapExprAndTermOrder` to order the result /// expressions and replaces the operation if the order changed. /// /// For example, the following operation: /// /// %0 = affine.min affine_map<(d0, d1) -> (d0 + d1, d1 + 16, 32)> (%i0, %i1) /// /// Turns into: /// /// %0 = affine.min affine_map<(d0, d1) -> (32, d1 + 16, d0 + d1)> (%i0, %i1) template <typename T> struct CanonicalizeAffineMinMaxOpExprAndTermOrder : public OpRewritePattern<T> { … }; template <typename T> struct CanonicalizeSingleResultAffineMinMaxOp : public OpRewritePattern<T> { … }; //===----------------------------------------------------------------------===// // AffineMinOp //===----------------------------------------------------------------------===// // // %0 = affine.min (d0) -> (1000, d0 + 512) (%i0) // OpFoldResult AffineMinOp::fold(FoldAdaptor adaptor) { … } void AffineMinOp::getCanonicalizationPatterns(RewritePatternSet &patterns, MLIRContext *context) { … } LogicalResult AffineMinOp::verify() { … } ParseResult AffineMinOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineMinOp::print(OpAsmPrinter &p) { … } //===----------------------------------------------------------------------===// // AffineMaxOp //===----------------------------------------------------------------------===// // // %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) // OpFoldResult AffineMaxOp::fold(FoldAdaptor adaptor) { … } void AffineMaxOp::getCanonicalizationPatterns(RewritePatternSet &patterns, MLIRContext *context) { … } LogicalResult AffineMaxOp::verify() { … } ParseResult AffineMaxOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineMaxOp::print(OpAsmPrinter &p) { … } //===----------------------------------------------------------------------===// // AffinePrefetchOp //===----------------------------------------------------------------------===// // // affine.prefetch %0[%i, %j + 5], read, locality<3>, data : memref<400x400xi32> // ParseResult AffinePrefetchOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffinePrefetchOp::print(OpAsmPrinter &p) { … } LogicalResult AffinePrefetchOp::verify() { … } void AffinePrefetchOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } LogicalResult AffinePrefetchOp::fold(FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) { … } //===----------------------------------------------------------------------===// // AffineParallelOp //===----------------------------------------------------------------------===// void AffineParallelOp::build(OpBuilder &builder, OperationState &result, TypeRange resultTypes, ArrayRef<arith::AtomicRMWKind> reductions, ArrayRef<int64_t> ranges) { … } void AffineParallelOp::build(OpBuilder &builder, OperationState &result, TypeRange resultTypes, ArrayRef<arith::AtomicRMWKind> reductions, ArrayRef<AffineMap> lbMaps, ValueRange lbArgs, ArrayRef<AffineMap> ubMaps, ValueRange ubArgs, ArrayRef<int64_t> steps) { … } SmallVector<Region *> AffineParallelOp::getLoopRegions() { … } unsigned AffineParallelOp::getNumDims() { … } AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() { … } AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() { … } AffineMap AffineParallelOp::getLowerBoundMap(unsigned pos) { … } AffineMap AffineParallelOp::getUpperBoundMap(unsigned pos) { … } AffineValueMap AffineParallelOp::getLowerBoundsValueMap() { … } AffineValueMap AffineParallelOp::getUpperBoundsValueMap() { … } std::optional<SmallVector<int64_t, 8>> AffineParallelOp::getConstantRanges() { … } Block *AffineParallelOp::getBody() { … } OpBuilder AffineParallelOp::getBodyBuilder() { … } void AffineParallelOp::setLowerBounds(ValueRange lbOperands, AffineMap map) { … } void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) { … } void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) { … } // check whether resultType match op or not in affine.parallel static bool isResultTypeMatchAtomicRMWKind(Type resultType, arith::AtomicRMWKind op) { … } LogicalResult AffineParallelOp::verify() { … } LogicalResult AffineValueMap::canonicalize() { … } /// Canonicalize the bounds of the given loop. static LogicalResult canonicalizeLoopBounds(AffineParallelOp op) { … } LogicalResult AffineParallelOp::fold(FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) { … } /// Prints a lower(upper) bound of an affine parallel loop with max(min) /// conditions in it. `mapAttr` is a flat list of affine expressions and `group` /// identifies which of the those expressions form max/min groups. `operands` /// are the SSA values of dimensions and symbols and `keyword` is either "min" /// or "max". static void printMinMaxBound(OpAsmPrinter &p, AffineMapAttr mapAttr, DenseIntElementsAttr group, ValueRange operands, StringRef keyword) { … } void AffineParallelOp::print(OpAsmPrinter &p) { … } /// Given a list of lists of parsed operands, populates `uniqueOperands` with /// unique operands. Also populates `replacements with affine expressions of /// `kind` that can be used to update affine maps previously accepting a /// `operands` to accept `uniqueOperands` instead. static ParseResult deduplicateAndResolveOperands( OpAsmParser &parser, ArrayRef<SmallVector<OpAsmParser::UnresolvedOperand>> operands, SmallVectorImpl<Value> &uniqueOperands, SmallVectorImpl<AffineExpr> &replacements, AffineExprKind kind) { … } namespace { enum class MinMaxKind { … }; } // namespace /// Parses an affine map that can contain a min/max for groups of its results, /// e.g., max(expr-1, expr-2), expr-3, max(expr-4, expr-5, expr-6). Populates /// `result` attributes with the map (flat list of expressions) and the grouping /// (list of integers that specify how many expressions to put into each /// min/max) attributes. Deduplicates repeated operands. /// /// parallel-bound ::= `(` parallel-group-list `)` /// parallel-group-list ::= parallel-group (`,` parallel-group-list)? /// parallel-group ::= simple-group | min-max-group /// simple-group ::= expr-of-ssa-ids /// min-max-group ::= ( `min` | `max` ) `(` expr-of-ssa-ids-list `)` /// expr-of-ssa-ids-list ::= expr-of-ssa-ids (`,` expr-of-ssa-id-list)? /// /// Examples: /// (%0, min(%1 + %2, %3), %4, min(%5 floordiv 32, %6)) /// (%0, max(%1 - 2 * %2)) static ParseResult parseAffineMapWithMinMax(OpAsmParser &parser, OperationState &result, MinMaxKind kind) { … } // // operation ::= `affine.parallel` `(` ssa-ids `)` `=` parallel-bound // `to` parallel-bound steps? region attr-dict? // steps ::= `steps` `(` integer-literals `)` // ParseResult AffineParallelOp::parse(OpAsmParser &parser, OperationState &result) { … } //===----------------------------------------------------------------------===// // AffineYieldOp //===----------------------------------------------------------------------===// LogicalResult AffineYieldOp::verify() { … } //===----------------------------------------------------------------------===// // AffineVectorLoadOp //===----------------------------------------------------------------------===// void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result, VectorType resultType, AffineMap map, ValueRange operands) { … } void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result, VectorType resultType, Value memref, AffineMap map, ValueRange mapOperands) { … } void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result, VectorType resultType, Value memref, ValueRange indices) { … } void AffineVectorLoadOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { … } ParseResult AffineVectorLoadOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineVectorLoadOp::print(OpAsmPrinter &p) { … } /// Verify common invariants of affine.vector_load and affine.vector_store. static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType, VectorType vectorType) { … } LogicalResult AffineVectorLoadOp::verify() { … } //===----------------------------------------------------------------------===// // AffineVectorStoreOp //===----------------------------------------------------------------------===// void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result, Value valueToStore, Value memref, AffineMap map, ValueRange mapOperands) { … } // Use identity map. void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result, Value valueToStore, Value memref, ValueRange indices) { … } void AffineVectorStoreOp::getCanonicalizationPatterns( RewritePatternSet &results, MLIRContext *context) { … } ParseResult AffineVectorStoreOp::parse(OpAsmParser &parser, OperationState &result) { … } void AffineVectorStoreOp::print(OpAsmPrinter &p) { … } LogicalResult AffineVectorStoreOp::verify() { … } //===----------------------------------------------------------------------===// // DelinearizeIndexOp //===----------------------------------------------------------------------===// LogicalResult AffineDelinearizeIndexOp::inferReturnTypes( MLIRContext *context, std::optional<::mlir::Location> location, ValueRange operands, DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, SmallVectorImpl<Type> &inferredReturnTypes) { … } void AffineDelinearizeIndexOp::build(OpBuilder &builder, OperationState &result, Value linearIndex, ArrayRef<OpFoldResult> basis) { … } LogicalResult AffineDelinearizeIndexOp::verify() { … } namespace { // Drops delinearization indices that correspond to unit-extent basis struct DropUnitExtentBasis : public OpRewritePattern<affine::AffineDelinearizeIndexOp> { … }; /// Drop delinearization pattern related to loops in the following way /// /// ``` /// <loop>(%iv) = (%c0) to (%ub) step (%c1) { /// %0 = affine.delinearize_index %iv into (%ub) : index /// <some_use>(%0) /// } /// ``` /// /// can be canonicalized to /// /// ``` /// <loop>(%iv) = (%c0) to (%ub) step (%c1) { /// <some_use>(%iv) /// } /// ``` struct DropDelinearizeOfSingleLoop : public OpRewritePattern<affine::AffineDelinearizeIndexOp> { … }; } // namespace void affine::AffineDelinearizeIndexOp::getCanonicalizationPatterns( RewritePatternSet &patterns, MLIRContext *context) { … } //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// #define GET_OP_CLASSES #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
Codebase Browser
llvm