| //===- AffineOps.h - 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 |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file defines convenience types for working with Affine operations |
| // in the MLIR operation set. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H |
| #define MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H |
| |
| #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h" |
| #include "mlir/Dialect/StandardOps/IR/Ops.h" |
| #include "mlir/IR/AffineMap.h" |
| #include "mlir/Interfaces/LoopLikeInterface.h" |
| |
| namespace mlir { |
| class AffineApplyOp; |
| class AffineBound; |
| class AffineValueMap; |
| |
| /// A utility function to check if a value is defined at the top level of an |
| /// op with trait `AffineScope` or is a region argument for such an op. A value |
| /// of index type defined at the top level is always a valid symbol for all its |
| /// uses. |
| bool isTopLevelValue(Value value); |
| |
| /// AffineDmaStartOp starts a non-blocking DMA operation that transfers data |
| /// from a source memref to a destination memref. The source and destination |
| /// memref need not be of the same dimensionality, but need to have the same |
| /// elemental type. The operands include the source and destination memref's |
| /// each followed by its indices, size of the data transfer in terms of the |
| /// number of elements (of the elemental type of the memref), a tag memref with |
| /// its indices, and optionally at the end, a stride and a |
| /// number_of_elements_per_stride arguments. The tag location is used by an |
| /// AffineDmaWaitOp to check for completion. The indices of the source memref, |
| /// destination memref, and the tag memref have the same restrictions as any |
| /// affine.load/store. In particular, index for each memref dimension must be an |
| /// affine expression of loop induction variables and symbols. |
| /// The optional stride arguments should be of 'index' type, and specify a |
| /// stride for the slower memory space (memory space with a lower memory space |
| /// id), transferring chunks of number_of_elements_per_stride every stride until |
| /// %num_elements are transferred. Either both or no stride arguments should be |
| /// specified. The value of 'num_elements' must be a multiple of |
| /// 'number_of_elements_per_stride'. If the source and destination locations |
| /// overlap the behavior of this operation is not defined. |
| // |
| // For example, an AffineDmaStartOp operation that transfers 256 elements of a |
| // memref '%src' in memory space 0 at indices [%i + 3, %j] to memref '%dst' in |
| // memory space 1 at indices [%k + 7, %l], would be specified as follows: |
| // |
| // %num_elements = arith.constant 256 |
| // %idx = arith.constant 0 : index |
| // %tag = alloc() : memref<1xi32, 4> |
| // affine.dma_start %src[%i + 3, %j], %dst[%k + 7, %l], %tag[%idx], |
| // %num_elements : |
| // memref<40x128xf32, 0>, memref<2x1024xf32, 1>, memref<1xi32, 2> |
| // |
| // If %stride and %num_elt_per_stride are specified, the DMA is expected to |
| // transfer %num_elt_per_stride elements every %stride elements apart from |
| // memory space 0 until %num_elements are transferred. |
| // |
| // affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%idx], %num_elements, |
| // %stride, %num_elt_per_stride : ... |
| // |
| // TODO: add additional operands to allow source and destination striding, and |
| // multiple stride levels (possibly using AffineMaps to specify multiple levels |
| // of striding). |
| class AffineDmaStartOp |
| : public Op<AffineDmaStartOp, OpTrait::MemRefsNormalizable, |
| OpTrait::VariadicOperands, OpTrait::ZeroResult, |
| AffineMapAccessInterface::Trait> { |
| public: |
| using Op::Op; |
| static ArrayRef<StringRef> getAttributeNames() { return {}; } |
| |
| static void 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 = nullptr, Value elementsPerStride = nullptr); |
| |
| /// Returns the operand index of the source memref. |
| unsigned getSrcMemRefOperandIndex() { return 0; } |
| |
| /// Returns the source MemRefType for this DMA operation. |
| Value getSrcMemRef() { return getOperand(getSrcMemRefOperandIndex()); } |
| MemRefType getSrcMemRefType() { |
| return getSrcMemRef().getType().cast<MemRefType>(); |
| } |
| |
| /// Returns the rank (number of indices) of the source MemRefType. |
| unsigned getSrcMemRefRank() { return getSrcMemRefType().getRank(); } |
| |
| /// Returns the affine map used to access the source memref. |
| AffineMap getSrcMap() { return getSrcMapAttr().getValue(); } |
| AffineMapAttr getSrcMapAttr() { |
| return (*this)->getAttr(getSrcMapAttrName()).cast<AffineMapAttr>(); |
| } |
| |
| /// Returns the source memref affine map indices for this DMA operation. |
| operand_range getSrcIndices() { |
| return {operand_begin() + getSrcMemRefOperandIndex() + 1, |
| operand_begin() + getSrcMemRefOperandIndex() + 1 + |
| getSrcMap().getNumInputs()}; |
| } |
| |
| /// Returns the memory space of the source memref. |
| unsigned getSrcMemorySpace() { |
| return getSrcMemRef().getType().cast<MemRefType>().getMemorySpaceAsInt(); |
| } |
| |
| /// Returns the operand index of the destination memref. |
| unsigned getDstMemRefOperandIndex() { |
| return getSrcMemRefOperandIndex() + 1 + getSrcMap().getNumInputs(); |
| } |
| |
| /// Returns the destination MemRefType for this DMA operation. |
| Value getDstMemRef() { return getOperand(getDstMemRefOperandIndex()); } |
| MemRefType getDstMemRefType() { |
| return getDstMemRef().getType().cast<MemRefType>(); |
| } |
| |
| /// Returns the rank (number of indices) of the destination MemRefType. |
| unsigned getDstMemRefRank() { |
| return getDstMemRef().getType().cast<MemRefType>().getRank(); |
| } |
| |
| /// Returns the memory space of the source memref. |
| unsigned getDstMemorySpace() { |
| return getDstMemRef().getType().cast<MemRefType>().getMemorySpaceAsInt(); |
| } |
| |
| /// Returns the affine map used to access the destination memref. |
| AffineMap getDstMap() { return getDstMapAttr().getValue(); } |
| AffineMapAttr getDstMapAttr() { |
| return (*this)->getAttr(getDstMapAttrName()).cast<AffineMapAttr>(); |
| } |
| |
| /// Returns the destination memref indices for this DMA operation. |
| operand_range getDstIndices() { |
| return {operand_begin() + getDstMemRefOperandIndex() + 1, |
| operand_begin() + getDstMemRefOperandIndex() + 1 + |
| getDstMap().getNumInputs()}; |
| } |
| |
| /// Returns the operand index of the tag memref. |
| unsigned getTagMemRefOperandIndex() { |
| return getDstMemRefOperandIndex() + 1 + getDstMap().getNumInputs(); |
| } |
| |
| /// Returns the Tag MemRef for this DMA operation. |
| Value getTagMemRef() { return getOperand(getTagMemRefOperandIndex()); } |
| MemRefType getTagMemRefType() { |
| return getTagMemRef().getType().cast<MemRefType>(); |
| } |
| |
| /// Returns the rank (number of indices) of the tag MemRefType. |
| unsigned getTagMemRefRank() { |
| return getTagMemRef().getType().cast<MemRefType>().getRank(); |
| } |
| |
| /// Returns the affine map used to access the tag memref. |
| AffineMap getTagMap() { return getTagMapAttr().getValue(); } |
| AffineMapAttr getTagMapAttr() { |
| return (*this)->getAttr(getTagMapAttrName()).cast<AffineMapAttr>(); |
| } |
| |
| /// Returns the tag memref indices for this DMA operation. |
| operand_range getTagIndices() { |
| return {operand_begin() + getTagMemRefOperandIndex() + 1, |
| operand_begin() + getTagMemRefOperandIndex() + 1 + |
| getTagMap().getNumInputs()}; |
| } |
| |
| /// Returns the number of elements being transferred by this DMA operation. |
| Value getNumElements() { |
| return getOperand(getTagMemRefOperandIndex() + 1 + |
| getTagMap().getNumInputs()); |
| } |
| |
| /// Impelements the AffineMapAccessInterface. |
| /// Returns the AffineMapAttr associated with 'memref'. |
| NamedAttribute getAffineMapAttrForMemRef(Value memref) { |
| if (memref == getSrcMemRef()) |
| return {StringAttr::get(getContext(), getSrcMapAttrName()), |
| getSrcMapAttr()}; |
| if (memref == getDstMemRef()) |
| return {StringAttr::get(getContext(), getDstMapAttrName()), |
| getDstMapAttr()}; |
| assert(memref == getTagMemRef() && |
| "DmaStartOp expected source, destination or tag memref"); |
| return {StringAttr::get(getContext(), getTagMapAttrName()), |
| getTagMapAttr()}; |
| } |
| |
| /// Returns true if this is a DMA from a faster memory space to a slower one. |
| bool isDestMemorySpaceFaster() { |
| return (getSrcMemorySpace() < getDstMemorySpace()); |
| } |
| |
| /// Returns true if this is a DMA from a slower memory space to a faster one. |
| bool isSrcMemorySpaceFaster() { |
| // Assumes that a lower number is for a slower memory space. |
| return (getDstMemorySpace() < getSrcMemorySpace()); |
| } |
| |
| /// Given a DMA start operation, returns the operand position of either the |
| /// source or destination memref depending on the one that is at the higher |
| /// level of the memory hierarchy. Asserts failure if neither is true. |
| unsigned getFasterMemPos() { |
| assert(isSrcMemorySpaceFaster() || isDestMemorySpaceFaster()); |
| return isSrcMemorySpaceFaster() ? 0 : getDstMemRefOperandIndex(); |
| } |
| |
| static StringRef getSrcMapAttrName() { return "src_map"; } |
| static StringRef getDstMapAttrName() { return "dst_map"; } |
| static StringRef getTagMapAttrName() { return "tag_map"; } |
| |
| static StringRef getOperationName() { return "affine.dma_start"; } |
| static ParseResult parse(OpAsmParser &parser, OperationState &result); |
| void print(OpAsmPrinter &p); |
| LogicalResult verify(); |
| LogicalResult fold(ArrayRef<Attribute> cstOperands, |
| SmallVectorImpl<OpFoldResult> &results); |
| |
| /// Returns true if this DMA operation is strided, returns false otherwise. |
| bool isStrided() { |
| return getNumOperands() != |
| getTagMemRefOperandIndex() + 1 + getTagMap().getNumInputs() + 1; |
| } |
| |
| /// Returns the stride value for this DMA operation. |
| Value getStride() { |
| if (!isStrided()) |
| return nullptr; |
| return getOperand(getNumOperands() - 1 - 1); |
| } |
| |
| /// Returns the number of elements to transfer per stride for this DMA op. |
| Value getNumElementsPerStride() { |
| if (!isStrided()) |
| return nullptr; |
| return getOperand(getNumOperands() - 1); |
| } |
| }; |
| |
| /// AffineDmaWaitOp blocks until the completion of a DMA operation associated |
| /// with the tag element '%tag[%index]'. %tag is a memref, and %index has to be |
| /// an index with the same restrictions as any load/store index. In particular, |
| /// index for each memref dimension must be an affine expression of loop |
| /// induction variables and symbols. %num_elements is the number of elements |
| /// associated with the DMA operation. For example: |
| // |
| // affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %num_elements : |
| // memref<2048xf32, 0>, memref<256xf32, 1>, memref<1xi32, 2> |
| // ... |
| // ... |
| // affine.dma_wait %tag[%index], %num_elements : memref<1xi32, 2> |
| // |
| class AffineDmaWaitOp |
| : public Op<AffineDmaWaitOp, OpTrait::MemRefsNormalizable, |
| OpTrait::VariadicOperands, OpTrait::ZeroResult, |
| AffineMapAccessInterface::Trait> { |
| public: |
| using Op::Op; |
| static ArrayRef<StringRef> getAttributeNames() { return {}; } |
| |
| static void build(OpBuilder &builder, OperationState &result, Value tagMemRef, |
| AffineMap tagMap, ValueRange tagIndices, Value numElements); |
| |
| static StringRef getOperationName() { return "affine.dma_wait"; } |
| |
| /// Returns the Tag MemRef associated with the DMA operation being waited on. |
| Value getTagMemRef() { return getOperand(0); } |
| MemRefType getTagMemRefType() { |
| return getTagMemRef().getType().cast<MemRefType>(); |
| } |
| |
| /// Returns the affine map used to access the tag memref. |
| AffineMap getTagMap() { return getTagMapAttr().getValue(); } |
| AffineMapAttr getTagMapAttr() { |
| return (*this)->getAttr(getTagMapAttrName()).cast<AffineMapAttr>(); |
| } |
| |
| /// Returns the tag memref index for this DMA operation. |
| operand_range getTagIndices() { |
| return {operand_begin() + 1, |
| operand_begin() + 1 + getTagMap().getNumInputs()}; |
| } |
| |
| /// Returns the rank (number of indices) of the tag memref. |
| unsigned getTagMemRefRank() { |
| return getTagMemRef().getType().cast<MemRefType>().getRank(); |
| } |
| |
| /// Impelements the AffineMapAccessInterface. Returns the AffineMapAttr |
| /// associated with 'memref'. |
| NamedAttribute getAffineMapAttrForMemRef(Value memref) { |
| assert(memref == getTagMemRef()); |
| return {StringAttr::get(getContext(), getTagMapAttrName()), |
| getTagMapAttr()}; |
| } |
| |
| /// Returns the number of elements transferred by the associated DMA op. |
| Value getNumElements() { return getOperand(1 + getTagMap().getNumInputs()); } |
| |
| static StringRef getTagMapAttrName() { return "tag_map"; } |
| static ParseResult parse(OpAsmParser &parser, OperationState &result); |
| void print(OpAsmPrinter &p); |
| LogicalResult verify(); |
| LogicalResult fold(ArrayRef<Attribute> cstOperands, |
| SmallVectorImpl<OpFoldResult> &results); |
| }; |
| |
| /// Returns true if the given Value can be used as a dimension id in the region |
| /// of the closest surrounding op that has the trait `AffineScope`. |
| bool isValidDim(Value value); |
| |
| /// Returns true if the given Value can be used as a dimension id in `region`, |
| /// i.e., for all its uses in `region`. |
| bool isValidDim(Value value, Region *region); |
| |
| /// Returns true if the given value can be used as a symbol in the region of the |
| /// closest surrounding op that has the trait `AffineScope`. |
| bool isValidSymbol(Value value); |
| |
| /// Returns true if the given Value can be used as a symbol for `region`, i.e., |
| /// for all its uses in `region`. |
| bool isValidSymbol(Value value, Region *region); |
| |
| /// Parses dimension and symbol list. `numDims` is set to the number of |
| /// dimensions in the list parsed. |
| ParseResult parseDimAndSymbolList(OpAsmParser &parser, |
| SmallVectorImpl<Value> &operands, |
| unsigned &numDims); |
| |
| /// Modifies both `map` and `operands` in-place so as to: |
| /// 1. drop duplicate operands |
| /// 2. drop unused dims and symbols from map |
| /// 3. promote valid symbols to symbolic operands in case they appeared as |
| /// dimensional operands |
| /// 4. propagate constant operands and drop them |
| void canonicalizeMapAndOperands(AffineMap *map, |
| SmallVectorImpl<Value> *operands); |
| |
| /// Canonicalizes an integer set the same way canonicalizeMapAndOperands does |
| /// for affine maps. |
| void canonicalizeSetAndOperands(IntegerSet *set, |
| SmallVectorImpl<Value> *operands); |
| |
| /// Returns a composed AffineApplyOp by composing `map` and `operands` with |
| /// other AffineApplyOps supplying those operands. The operands of the resulting |
| /// AffineApplyOp do not change the length of AffineApplyOp chains. |
| AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, |
| ValueRange operands); |
| /// Variant of `makeComposedAffineApply` which infers the AffineMap from `e`. |
| AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineExpr e, |
| ValueRange values); |
| |
| /// Given an affine map `map` and its input `operands`, this method composes |
| /// into `map`, maps of AffineApplyOps whose results are the values in |
| /// `operands`, iteratively until no more of `operands` are the result of an |
| /// AffineApplyOp. When this function returns, `map` becomes the composed affine |
| /// map, and each Value in `operands` is guaranteed to be either a loop IV or a |
| /// terminal symbol, i.e., a symbol defined at the top level or a block/function |
| /// argument. |
| void fullyComposeAffineMapAndOperands(AffineMap *map, |
| SmallVectorImpl<Value> *operands); |
| } // namespace mlir |
| #include "mlir/Dialect/Affine/IR/AffineOpsDialect.h.inc" |
| |
| #define GET_OP_CLASSES |
| #include "mlir/Dialect/Affine/IR/AffineOps.h.inc" |
| |
| namespace mlir { |
| /// Returns true if the provided value is the induction variable of a |
| /// AffineForOp. |
| bool isForInductionVar(Value val); |
| |
| /// Returns the loop parent of an induction variable. If the provided value is |
| /// not an induction variable, then return nullptr. |
| AffineForOp getForInductionVarOwner(Value val); |
| |
| /// Extracts the induction variables from a list of AffineForOps and places them |
| /// in the output argument `ivs`. |
| void extractForInductionVars(ArrayRef<AffineForOp> forInsts, |
| SmallVectorImpl<Value> *ivs); |
| |
| /// Builds a perfect nest of affine.for loops, i.e., each loop except the |
| /// innermost one contains only another loop and a terminator. The loops iterate |
| /// from "lbs" to "ubs" with "steps". The body of the innermost loop is |
| /// populated by calling "bodyBuilderFn" and providing it with an OpBuilder, a |
| /// Location and a list of loop induction variables. |
| void buildAffineLoopNest(OpBuilder &builder, Location loc, |
| ArrayRef<int64_t> lbs, ArrayRef<int64_t> ubs, |
| ArrayRef<int64_t> steps, |
| function_ref<void(OpBuilder &, Location, ValueRange)> |
| bodyBuilderFn = nullptr); |
| void buildAffineLoopNest(OpBuilder &builder, Location loc, ValueRange lbs, |
| ValueRange ubs, ArrayRef<int64_t> steps, |
| function_ref<void(OpBuilder &, Location, ValueRange)> |
| bodyBuilderFn = nullptr); |
| |
| /// Replace `loop` with a new loop where `newIterOperands` are appended with |
| /// new initialization values and `newYieldedValues` are added as new yielded |
| /// values. The returned ForOp has `newYieldedValues.size()` new result values. |
| /// Additionally, if `replaceLoopResults` is true, all uses of |
| /// `loop.getResults()` are replaced with the first `loop.getNumResults()` |
| /// return values of the original loop respectively. The original loop is |
| /// deleted and the new loop returned. |
| /// Prerequisite: `newIterOperands.size() == newYieldedValues.size()`. |
| AffineForOp replaceForOpWithNewYields(OpBuilder &b, AffineForOp loop, |
| ValueRange newIterOperands, |
| ValueRange newYieldedValues, |
| ValueRange newIterArgs, |
| bool replaceLoopResults = true); |
| |
| /// AffineBound represents a lower or upper bound in the for operation. |
| /// This class does not own the underlying operands. Instead, it refers |
| /// to the operands stored in the AffineForOp. Its life span should not exceed |
| /// that of the for operation it refers to. |
| class AffineBound { |
| public: |
| AffineForOp getAffineForOp() { return op; } |
| AffineMap getMap() { return map; } |
| |
| unsigned getNumOperands() { return opEnd - opStart; } |
| Value getOperand(unsigned idx) { return op.getOperand(opStart + idx); } |
| |
| using operand_iterator = AffineForOp::operand_iterator; |
| using operand_range = AffineForOp::operand_range; |
| |
| operand_iterator operand_begin() { return op.operand_begin() + opStart; } |
| operand_iterator operand_end() { return op.operand_begin() + opEnd; } |
| operand_range getOperands() { return {operand_begin(), operand_end()}; } |
| |
| private: |
| // 'affine.for' operation that contains this bound. |
| AffineForOp op; |
| // Start and end positions of this affine bound operands in the list of |
| // the containing 'affine.for' operation operands. |
| unsigned opStart, opEnd; |
| // Affine map for this bound. |
| AffineMap map; |
| |
| AffineBound(AffineForOp op, unsigned opStart, unsigned opEnd, AffineMap map) |
| : op(op), opStart(opStart), opEnd(opEnd), map(map) {} |
| |
| friend class AffineForOp; |
| }; |
| |
| } // end namespace mlir |
| |
| #endif |