From ffde975e215e8ccba2b96a05f66a5756bebc8b64 Mon Sep 17 00:00:00 2001 From: River Riddle Date: Tue, 20 Aug 2019 15:36:08 -0700 Subject: NFC: Move AffineOps dialect to the Dialect sub-directory. PiperOrigin-RevId: 264482571 --- .../Linalg/Linalg1/include/linalg1/Common.h | 2 +- mlir/examples/Linalg/Linalg4/lib/Transforms.cpp | 2 +- mlir/include/mlir/AffineOps/AffineOps.h | 598 ------- mlir/include/mlir/AffineOps/AffineOps.td | 259 --- mlir/include/mlir/AffineOps/AffineOpsBase.td | 44 - mlir/include/mlir/AffineOps/CMakeLists.txt | 4 - mlir/include/mlir/CMakeLists.txt | 1 - mlir/include/mlir/Dialect/AffineOps/AffineOps.h | 598 +++++++ mlir/include/mlir/Dialect/AffineOps/AffineOps.td | 259 +++ .../mlir/Dialect/AffineOps/AffineOpsBase.td | 44 + mlir/include/mlir/Dialect/AffineOps/CMakeLists.txt | 4 + mlir/include/mlir/Dialect/CMakeLists.txt | 1 + .../mlir/Dialect/Linalg/IR/LinalgLibraryOps.td | 2 +- mlir/include/mlir/EDSC/Builders.h | 2 +- mlir/lib/AffineOps/AffineOps.cpp | 1764 -------------------- mlir/lib/AffineOps/CMakeLists.txt | 10 - mlir/lib/AffineOps/DialectRegistration.cpp | 22 - mlir/lib/Analysis/AffineAnalysis.cpp | 2 +- mlir/lib/Analysis/AffineStructures.cpp | 2 +- mlir/lib/Analysis/LoopAnalysis.cpp | 2 +- mlir/lib/Analysis/MemRefBoundCheck.cpp | 2 +- mlir/lib/Analysis/NestedMatcher.cpp | 2 +- mlir/lib/Analysis/SliceAnalysis.cpp | 2 +- mlir/lib/Analysis/TestMemRefDependenceCheck.cpp | 2 +- mlir/lib/Analysis/TestParallelismDetection.cpp | 2 +- mlir/lib/Analysis/Utils.cpp | 2 +- mlir/lib/Analysis/VectorAnalysis.cpp | 2 +- mlir/lib/CMakeLists.txt | 1 - mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp | 2 +- mlir/lib/Conversion/LoopsToGPU/LoopsToGPUPass.cpp | 2 +- mlir/lib/Dialect/AffineOps/AffineOps.cpp | 1764 ++++++++++++++++++++ mlir/lib/Dialect/AffineOps/CMakeLists.txt | 10 + mlir/lib/Dialect/AffineOps/DialectRegistration.cpp | 22 + mlir/lib/Dialect/CMakeLists.txt | 1 + .../lib/Dialect/Linalg/Transforms/LowerToLoops.cpp | 12 +- mlir/lib/Transforms/AffineDataCopyGeneration.cpp | 2 +- mlir/lib/Transforms/LoopFusion.cpp | 2 +- mlir/lib/Transforms/LoopInvariantCodeMotion.cpp | 2 +- mlir/lib/Transforms/LoopTiling.cpp | 2 +- mlir/lib/Transforms/LoopUnroll.cpp | 2 +- mlir/lib/Transforms/LoopUnrollAndJam.cpp | 2 +- mlir/lib/Transforms/LowerAffine.cpp | 2 +- mlir/lib/Transforms/MaterializeVectors.cpp | 2 +- mlir/lib/Transforms/MemRefDataFlowOpt.cpp | 2 +- mlir/lib/Transforms/PipelineDataTransfer.cpp | 2 +- mlir/lib/Transforms/Utils/LoopFusionUtils.cpp | 2 +- mlir/lib/Transforms/Utils/LoopUtils.cpp | 2 +- mlir/lib/Transforms/Utils/Utils.cpp | 2 +- mlir/lib/Transforms/Vectorize.cpp | 2 +- mlir/test/EDSC/builder-api-test.cpp | 2 +- mlir/test/lib/Transforms/TestConstantFold.cpp | 2 +- mlir/test/lib/Transforms/TestLoopFusion.cpp | 2 +- .../test/lib/Transforms/TestVectorizationUtils.cpp | 2 +- 53 files changed, 2743 insertions(+), 2743 deletions(-) delete mode 100644 mlir/include/mlir/AffineOps/AffineOps.h delete mode 100644 mlir/include/mlir/AffineOps/AffineOps.td delete mode 100644 mlir/include/mlir/AffineOps/AffineOpsBase.td delete mode 100644 mlir/include/mlir/AffineOps/CMakeLists.txt create mode 100644 mlir/include/mlir/Dialect/AffineOps/AffineOps.h create mode 100644 mlir/include/mlir/Dialect/AffineOps/AffineOps.td create mode 100644 mlir/include/mlir/Dialect/AffineOps/AffineOpsBase.td create mode 100644 mlir/include/mlir/Dialect/AffineOps/CMakeLists.txt delete mode 100644 mlir/lib/AffineOps/AffineOps.cpp delete mode 100644 mlir/lib/AffineOps/CMakeLists.txt delete mode 100644 mlir/lib/AffineOps/DialectRegistration.cpp create mode 100644 mlir/lib/Dialect/AffineOps/AffineOps.cpp create mode 100644 mlir/lib/Dialect/AffineOps/CMakeLists.txt create mode 100644 mlir/lib/Dialect/AffineOps/DialectRegistration.cpp diff --git a/mlir/examples/Linalg/Linalg1/include/linalg1/Common.h b/mlir/examples/Linalg/Linalg1/include/linalg1/Common.h index 29ff9bd2d3e..8bedf513907 100644 --- a/mlir/examples/Linalg/Linalg1/include/linalg1/Common.h +++ b/mlir/examples/Linalg/Linalg1/include/linalg1/Common.h @@ -18,9 +18,9 @@ #ifndef LINALG1_COMMON_H_ #define LINALG1_COMMON_H_ -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Analysis/Verifier.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/EDSC/Builders.h" #include "mlir/EDSC/Helpers.h" diff --git a/mlir/examples/Linalg/Linalg4/lib/Transforms.cpp b/mlir/examples/Linalg/Linalg4/lib/Transforms.cpp index 15e544a773c..68143dcc146 100644 --- a/mlir/examples/Linalg/Linalg4/lib/Transforms.cpp +++ b/mlir/examples/Linalg/Linalg4/lib/Transforms.cpp @@ -23,7 +23,7 @@ #include "linalg3/Intrinsics.h" #include "linalg3/TensorOps.h" -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/EDSC/Helpers.h" #include "mlir/IR/OpImplementation.h" #include "mlir/Transforms/LoopUtils.h" diff --git a/mlir/include/mlir/AffineOps/AffineOps.h b/mlir/include/mlir/AffineOps/AffineOps.h deleted file mode 100644 index 59f7fc782e6..00000000000 --- a/mlir/include/mlir/AffineOps/AffineOps.h +++ /dev/null @@ -1,598 +0,0 @@ -//===- AffineOps.h - MLIR Affine Operations -------------------------------===// -// -// Copyright 2019 The MLIR Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// ============================================================================= -// -// This file defines convenience types for working with Affine operations -// in the MLIR operation set. -// -//===----------------------------------------------------------------------===// - -#ifndef MLIR_AFFINEOPS_AFFINEOPS_H -#define MLIR_AFFINEOPS_AFFINEOPS_H - -#include "mlir/IR/AffineMap.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/Dialect.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/StandardTypes.h" - -namespace mlir { -class AffineBound; -class AffineValueMap; -class AffineTerminatorOp; -class FlatAffineConstraints; -class OpBuilder; - -/// A utility function to check if a value is defined at the top level of a -/// function. A value defined at the top level is always a valid symbol. -bool isTopLevelSymbol(Value *value); - -class AffineOpsDialect : public Dialect { -public: - AffineOpsDialect(MLIRContext *context); - static StringRef getDialectNamespace() { return "affine"; } -}; - -/// The "affine.apply" operation applies an affine map to a list of operands, -/// yielding a single result. The operand list must be the same size as the -/// number of arguments to the affine mapping. All operands and the result are -/// of type 'Index'. This operation requires a single affine map attribute named -/// "map". For example: -/// -/// %y = "affine.apply" (%x) { map: (d0) -> (d0 + 1) } : -/// (index) -> (index) -/// -/// equivalently: -/// -/// #map42 = (d0)->(d0+1) -/// %y = affine.apply #map42(%x) -/// -class AffineApplyOp : public Op { -public: - using Op::Op; - - /// Builds an affine apply op with the specified map and operands. - static void build(Builder *builder, OperationState *result, AffineMap map, - ArrayRef operands); - - /// Returns the affine map to be applied by this operation. - AffineMap getAffineMap() { - return getAttrOfType("map").getValue(); - } - - /// Returns true if the result of this operation can be used as dimension id. - bool isValidDim(); - - /// Returns true if the result of this operation is a symbol. - bool isValidSymbol(); - - static StringRef getOperationName() { return "affine.apply"; } - - // Hooks to customize behavior of this op. - static ParseResult parse(OpAsmParser *parser, OperationState *result); - void print(OpAsmPrinter *p); - LogicalResult verify(); - OpFoldResult fold(ArrayRef operands); - - static void getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context); -}; - -/// 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), tranferring 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'. -// -// For example, a DmaStartOp 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 = constant 256 -// %idx = 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(mlir-team): add additional operands to allow source and destination -// striding, and multiple stride levels (possibly using AffineMaps to specify -// multiple levels of striding). -// TODO(andydavis) Consider replacing src/dst memref indices with view memrefs. -class AffineDmaStartOp : public Op { -public: - using Op::Op; - - static void build(Builder *builder, OperationState *result, Value *srcMemRef, - AffineMap srcMap, ArrayRef srcIndices, - Value *destMemRef, AffineMap dstMap, - ArrayRef destIndices, Value *tagMemRef, - AffineMap tagMap, ArrayRef tagIndices, - Value *numElements, Value *stride = nullptr, - Value *elementsPerStride = nullptr); - - /// Returns the operand index of the src memref. - unsigned getSrcMemRefOperandIndex() { return 0; } - - /// Returns the source MemRefType for this DMA operation. - Value *getSrcMemRef() { return getOperand(getSrcMemRefOperandIndex()); } - MemRefType getSrcMemRefType() { - return getSrcMemRef()->getType().cast(); - } - - /// Returns the rank (number of indices) of the source MemRefType. - unsigned getSrcMemRefRank() { return getSrcMemRefType().getRank(); } - - /// Returns the affine map used to access the src memref. - AffineMap getSrcMap() { return getSrcMapAttr().getValue(); } - AffineMapAttr getSrcMapAttr() { - return getAttr(getSrcMapAttrName()).cast(); - } - - /// 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 src memref. - unsigned getSrcMemorySpace() { - return getSrcMemRef()->getType().cast().getMemorySpace(); - } - - /// Returns the operand index of the dst memref. - unsigned getDstMemRefOperandIndex() { - return getSrcMemRefOperandIndex() + 1 + getSrcMap().getNumInputs(); - } - - /// Returns the destination MemRefType for this DMA operations. - Value *getDstMemRef() { return getOperand(getDstMemRefOperandIndex()); } - MemRefType getDstMemRefType() { - return getDstMemRef()->getType().cast(); - } - - /// Returns the rank (number of indices) of the destination MemRefType. - unsigned getDstMemRefRank() { - return getDstMemRef()->getType().cast().getRank(); - } - - /// Returns the memory space of the src memref. - unsigned getDstMemorySpace() { - return getDstMemRef()->getType().cast().getMemorySpace(); - } - - /// Returns the affine map used to access the dst memref. - AffineMap getDstMap() { return getDstMapAttr().getValue(); } - AffineMapAttr getDstMapAttr() { - return getAttr(getDstMapAttrName()).cast(); - } - - /// 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(); - } - - /// Returns the rank (number of indices) of the tag MemRefType. - unsigned getTagMemRefRank() { - return getTagMemRef()->getType().cast().getRank(); - } - - /// Returns the affine map used to access the tag memref. - AffineMap getTagMap() { return getTagMapAttr().getValue(); } - AffineMapAttr getTagMapAttr() { - return getAttr(getTagMapAttrName()).cast(); - } - - /// 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()); - } - - /// Returns the AffineMapAttr associated with 'memref'. - NamedAttribute getAffineMapAttrForMemRef(Value *memref) { - if (memref == getSrcMemRef()) - return {Identifier::get(getSrcMapAttrName(), getContext()), - getSrcMapAttr()}; - else if (memref == getDstMemRef()) - return {Identifier::get(getDstMapAttrName(), getContext()), - getDstMapAttr()}; - assert(memref == getTagMemRef() && - "DmaStartOp expected source, destination or tag memref"); - return {Identifier::get(getTagMapAttrName(), getContext()), - 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(); - static void getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context); - - /// 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 { -public: - using Op::Op; - - static void build(Builder *builder, OperationState *result, Value *tagMemRef, - AffineMap tagMap, ArrayRef 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(); - } - - /// Returns the affine map used to access the tag memref. - AffineMap getTagMap() { return getTagMapAttr().getValue(); } - AffineMapAttr getTagMapAttr() { - return getAttr(getTagMapAttrName()).cast(); - } - - // 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().getRank(); - } - - /// Returns the AffineMapAttr associated with 'memref'. - NamedAttribute getAffineMapAttrForMemRef(Value *memref) { - assert(memref == getTagMemRef()); - return {Identifier::get(getTagMapAttrName(), getContext()), - getTagMapAttr()}; - } - - /// Returns the number of elements transferred in 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(); - static void getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context); -}; - -/// The "affine.load" op reads an element from a memref, where the index -/// for each memref dimension is an affine expression of loop induction -/// variables and symbols. The output of 'affine.load' is a new value with the -/// same type as the elements of the memref. An affine expression of loop IVs -/// and symbols must be specified for each dimension of the memref. The keyword -/// 'symbol' can be used to indicate SSA identifiers which are symbolic. -// -// Example 1: -// -// %1 = affine.load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32> -// -// Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'. -// -// %1 = affine.load %0[%i0 + symbol(%n), %i1 + symbol(%m)] -// : memref<100x100xf32> -// -class AffineLoadOp : public Op::Impl> { -public: - using Op::Op; - - /// Builds an affine load op with the specified map and operands. - static void build(Builder *builder, OperationState *result, AffineMap map, - ArrayRef operands); - /// Builds an affine load op an identify map and operands. - static void build(Builder *builder, OperationState *result, Value *memref, - ArrayRef indices = {}); - - /// Returns the operand index of the memref. - unsigned getMemRefOperandIndex() { return 0; } - - /// Get memref operand. - Value *getMemRef() { return getOperand(getMemRefOperandIndex()); } - void setMemRef(Value *value) { setOperand(getMemRefOperandIndex(), value); } - MemRefType getMemRefType() { - return getMemRef()->getType().cast(); - } - - /// Get affine map operands. - operand_range getIndices() { return llvm::drop_begin(getOperands(), 1); } - - /// Returns the affine map used to index the memref for this operation. - AffineMap getAffineMap() { return getAffineMapAttr().getValue(); } - AffineMapAttr getAffineMapAttr() { - return getAttr(getMapAttrName()).cast(); - } - - /// Returns the AffineMapAttr associated with 'memref'. - NamedAttribute getAffineMapAttrForMemRef(Value *memref) { - assert(memref == getMemRef()); - return {Identifier::get(getMapAttrName(), getContext()), - getAffineMapAttr()}; - } - - static StringRef getMapAttrName() { return "map"; } - static StringRef getOperationName() { return "affine.load"; } - - // Hooks to customize behavior of this op. - static ParseResult parse(OpAsmParser *parser, OperationState *result); - void print(OpAsmPrinter *p); - LogicalResult verify(); - static void getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context); -}; - -/// The "affine.store" op writes an element to a memref, where the index -/// for each memref dimension is an affine expression of loop induction -/// variables and symbols. The 'affine.store' op stores a new value which is the -/// same type as the elements of the memref. An affine expression of loop IVs -/// and symbols must be specified for each dimension of the memref. The keyword -/// 'symbol' can be used to indicate SSA identifiers which are symbolic. -// -// Example 1: -// -// affine.store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32> -// -// Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'. -// -// affine.store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] -// : memref<100x100xf32> -// -class AffineStoreOp : public Op::Impl> { -public: - using Op::Op; - - /// Builds an affine store operation with the specified map and operands. - static void build(Builder *builder, OperationState *result, - Value *valueToStore, AffineMap map, - ArrayRef operands); - /// Builds an affine store operation with an identity map and operands. - static void build(Builder *builder, OperationState *result, - Value *valueToStore, Value *memref, - ArrayRef operands); - - /// Get value to be stored by store operation. - Value *getValueToStore() { return getOperand(0); } - - /// Returns the operand index of the memref. - unsigned getMemRefOperandIndex() { return 1; } - - /// Get memref operand. - Value *getMemRef() { return getOperand(getMemRefOperandIndex()); } - void setMemRef(Value *value) { setOperand(getMemRefOperandIndex(), value); } - - MemRefType getMemRefType() { - return getMemRef()->getType().cast(); - } - - /// Get affine map operands. - operand_range getIndices() { return llvm::drop_begin(getOperands(), 2); } - - /// Returns the affine map used to index the memref for this operation. - AffineMap getAffineMap() { return getAffineMapAttr().getValue(); } - AffineMapAttr getAffineMapAttr() { - return getAttr(getMapAttrName()).cast(); - } - - /// Returns the AffineMapAttr associated with 'memref'. - NamedAttribute getAffineMapAttrForMemRef(Value *memref) { - assert(memref == getMemRef()); - return {Identifier::get(getMapAttrName(), getContext()), - getAffineMapAttr()}; - } - - static StringRef getMapAttrName() { return "map"; } - static StringRef getOperationName() { return "affine.store"; } - - // Hooks to customize behavior of this op. - static ParseResult parse(OpAsmParser *parser, OperationState *result); - void print(OpAsmPrinter *p); - LogicalResult verify(); - static void getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context); -}; - -/// Returns true if the given Value can be used as a dimension id. -bool isValidDim(Value *value); - -/// Returns true if the given Value can be used as a symbol. -bool isValidSymbol(Value *value); - -/// Modifies both `map` and `operands` in-place so as to: -/// 1. drop duplicate operands -/// 2. drop unused dims and symbols from map -void canonicalizeMapAndOperands(AffineMap *map, - llvm::SmallVectorImpl *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, - llvm::ArrayRef operands); - -/// 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, - llvm::SmallVectorImpl *operands); - -#define GET_OP_CLASSES -#include "mlir/AffineOps/AffineOps.h.inc" - -/// Returns 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 forInsts, - SmallVectorImpl *ivs); - -/// 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; } - - /// Returns an AffineValueMap representing this bound. - AffineValueMap getAsAffineValueMap(); - - unsigned getNumOperands() { return opEnd - opStart; } - Value *getOperand(unsigned idx) { - return op.getOperation()->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 diff --git a/mlir/include/mlir/AffineOps/AffineOps.td b/mlir/include/mlir/AffineOps/AffineOps.td deleted file mode 100644 index c517ed0244d..00000000000 --- a/mlir/include/mlir/AffineOps/AffineOps.td +++ /dev/null @@ -1,259 +0,0 @@ -//===- AffineOps.td - Affine operation definitions ---------*- tablegen -*-===// -// -// Copyright 2019 The MLIR Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// ============================================================================= -// -// Defines MLIR affine operations. -// -//===----------------------------------------------------------------------===// - -#ifdef AFFINE_OPS -#else -#define AFFINE_OPS - -#ifdef OP_BASE -#else -include "mlir/IR/OpBase.td" -#endif // OP_BASE - -include "mlir/AffineOps/AffineOpsBase.td" - -def Affine_Dialect : Dialect { - let name = "affine"; - let cppNamespace = ""; -} - -// Base class for Affine dialect ops. -class Affine_Op traits = []> : - Op { - // For every affine op, there needs to be a: - // * void print(OpAsmPrinter *p, ${C++ class of Op} op) - // * LogicalResult verify(${C++ class of Op} op) - // * ParseResult parse${C++ class of Op}(OpAsmParser *parser, - // OperationState *result) - // functions. - let printer = [{ return ::print(p, *this); }]; - let verifier = [{ return ::verify(*this); }]; - let parser = [{ return ::parse$cppClass(parser, result); }]; -} - -// Require regions to have affine terminator. -def ImplicitAffineTerminator - : SingleBlockImplicitTerminator<"AffineTerminatorOp">; - -def AffineForOp : Affine_Op<"for", [ImplicitAffineTerminator]> { - let summary = "for operation"; - let description = [{ - The "affine.for" operation represents an affine loop nest, defining an SSA - value for its induction variable. It has one region capturing the loop body. - The induction variable is represented as a argument of this region. This SSA - value always has type index, which is the size of the machine word. The - stride, represented by step, is a positive constant integer which defaults - to "1" if not present. The lower and upper bounds specify a half-open range: - the range includes the lower bound but does not include the upper bound. - - The body region must contain exactly one block that terminates with - "affine.terminator". Calling AffineForOp::build will create such region - and insert the terminator, so will the parsing even in cases if it is absent - from the custom format. - - The lower and upper bounds of a for operation are represented as an - application of an affine mapping to a list of SSA values passed to the map. - The same restrictions hold for these SSA values as for all bindings of SSA - values to dimensions and symbols. The affine mappings for the bounds may - return multiple results, in which case the max/min keywords are required - (for the lower/upper bound respectively), and the bound is the - maximum/minimum of the returned values. - - Example: - - affine.for %i = 1 to 10 { - ... - } - - }]; - let arguments = (ins Variadic); - let regions = (region SizedRegion<1>:$region); - - let skipDefaultBuilders = 1; - let builders = [ - OpBuilder<"Builder *builder, OperationState *result, " - "int64_t lowerBound, int64_t upperBound, int64_t step = 1">, - OpBuilder<"Builder *builder, OperationState *result, " - "ArrayRef lbOperands, AffineMap lbMap, " - "ArrayRef ubOperands, AffineMap ubMap, " - "int64_t step = 1"> - ]; - - let extraClassDeclaration = [{ - static StringRef getStepAttrName() { return "step"; } - static StringRef getLowerBoundAttrName() { return "lower_bound"; } - static StringRef getUpperBoundAttrName() { return "upper_bound"; } - - Block *getBody() { return ®ion().front(); } - Value *getInductionVar() { return getBody()->getArgument(0); } - OpBuilder getBodyBuilder() { - return OpBuilder(getBody(), std::prev(getBody()->end())); - } - - // TODO: provide iterators for the lower and upper bound operands - // if the current access via getLowerBound(), getUpperBound() is too slow. - - /// Returns operands for the lower bound map. - operand_range getLowerBoundOperands(); - - /// Returns operands for the upper bound map. - operand_range getUpperBoundOperands(); - - /// Returns information about the lower bound as a single object. - AffineBound getLowerBound(); - - /// Returns information about the upper bound as a single object. - AffineBound getUpperBound(); - - /// Returns loop step. - int64_t getStep() { - return getAttr(getStepAttrName()).cast().getInt(); - } - - /// Returns affine map for the lower bound. - AffineMap getLowerBoundMap() { return getLowerBoundMapAttr().getValue(); } - AffineMapAttr getLowerBoundMapAttr() { - return getAttr(getLowerBoundAttrName()).cast(); - } - /// Returns affine map for the upper bound. The upper bound is exclusive. - AffineMap getUpperBoundMap() { return getUpperBoundMapAttr().getValue(); } - AffineMapAttr getUpperBoundMapAttr() { - return getAttr(getUpperBoundAttrName()).cast(); - } - - /// Set lower bound. The new bound must have the same number of operands as - /// the current bound map. Otherwise, 'replaceForLowerBound' should be used. - void setLowerBound(ArrayRef operands, AffineMap map); - /// Set upper bound. The new bound must not have more operands than the - /// current bound map. Otherwise, 'replaceForUpperBound' should be used. - void setUpperBound(ArrayRef operands, AffineMap map); - - /// Set the lower bound map without changing operands. - void setLowerBoundMap(AffineMap map); - - /// Set the upper bound map without changing operands. - void setUpperBoundMap(AffineMap map); - - /// Set loop step. - void setStep(int64_t step) { - assert(step > 0 && "step has to be a positive integer constant"); - auto *context = getLowerBoundMap().getContext(); - setAttr(Identifier::get(getStepAttrName(), context), - IntegerAttr::get(IndexType::get(context), step)); - } - - /// Returns true if the lower bound is constant. - bool hasConstantLowerBound(); - /// Returns true if the upper bound is constant. - bool hasConstantUpperBound(); - /// Returns true if both bounds are constant. - bool hasConstantBounds() { - return hasConstantLowerBound() && hasConstantUpperBound(); - } - /// Returns the value of the constant lower bound. - /// Fails assertion if the bound is non-constant. - int64_t getConstantLowerBound(); - /// Returns the value of the constant upper bound. The upper bound is - /// exclusive. Fails assertion if the bound is non-constant. - int64_t getConstantUpperBound(); - /// Sets the lower bound to the given constant value. - void setConstantLowerBound(int64_t value); - /// Sets the upper bound to the given constant value. - void setConstantUpperBound(int64_t value); - - /// Returns true if both the lower and upper bound have the same operand - /// lists (same operands in the same order). - bool matchingBoundOperandList(); - }]; - - let hasCanonicalizer = 1; -} - -def AffineIfOp : Affine_Op<"if", [ImplicitAffineTerminator]> { - let summary = "if-then-else operation"; - let description = [{ - The "if" operation represents an if-then-else construct for conditionally - executing two regions of code. The operands to an if operation are an - IntegerSet condition and a set of symbol/dimension operands to the - condition set. The operation produces no results. For example: - - affine.if #set(%i) { - ... - } else { - ... - } - - The 'else' blocks to the if operation are optional, and may be omitted. For - example: - - affine.if #set(%i) { - ... - } - }]; - let arguments = (ins Variadic); - let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion); - - let skipDefaultBuilders = 1; - let builders = [ - OpBuilder<"Builder *builder, OperationState *result, " - "Value *cond, bool withElseRegion"> - ]; - - let extraClassDeclaration = [{ - static StringRef getConditionAttrName() { return "condition"; } - - IntegerSet getIntegerSet(); - void setIntegerSet(IntegerSet newSet); - - OpBuilder getThenBodyBuilder() { - assert(!thenRegion().empty() && "Unexpected empty 'then' region."); - Block &body = thenRegion().front(); - return OpBuilder(&body, std::prev(body.end())); - } - OpBuilder getElseBodyBuilder() { - assert(!elseRegion().empty() && "Unexpected empty 'else' region."); - Block &body = elseRegion().front(); - return OpBuilder(&body, std::prev(body.end())); - } - }]; -} - -def AffineTerminatorOp : - Affine_Op<"terminator", [Terminator]> { - let summary = "affine terminator operation"; - let description = [{ - Affine terminator is a special terminator operation for blocks inside affine - loops and branches. It unconditionally transmits the control flow to the - successor of the operation enclosing the region. - - This operation does _not_ have a custom syntax. However, affine control - operations omit the terminator in their custom syntax for brevity. - }]; - - // No custom parsing/printing form. - let parser = ?; - let printer = ?; - - // Fully specified by traits. - let verifier = ?; -} - -#endif // AFFINE_OPS diff --git a/mlir/include/mlir/AffineOps/AffineOpsBase.td b/mlir/include/mlir/AffineOps/AffineOpsBase.td deleted file mode 100644 index 2ac1d379c12..00000000000 --- a/mlir/include/mlir/AffineOps/AffineOpsBase.td +++ /dev/null @@ -1,44 +0,0 @@ -//===- AffineOpsBase.td - Affine operation definitions -----*- tablegen -*-===// -// -// Copyright 2019 The MLIR Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// ============================================================================= -// -// Defines base support for MLIR affine operations. -// -//===----------------------------------------------------------------------===// - -#ifdef AFFINE_OPS_BASE -#else -#define AFFINE_OPS_BASE - -#ifdef OP_BASE -#else -include "mlir/IR/OpBase.td" -#endif // OP_BASE - -// Attributes containing affine maps. -def AffineMapAttr : Attr< - CPred<"$_self.isa()">, "AffineMap attribute"> { - let storageType = [{ AffineMapAttr }]; - let returnType = [{ AffineMap }]; - let constBuilderCall = "$_builder.getAffineMapAttr($0)"; -} - -def AffineMapArrayAttr : TypedArrayAttrBase { - let constBuilderCall = "$_builder.getAffineMapArrayAttr($0)"; -} - -#endif // AFFINE_OPS_BASE diff --git a/mlir/include/mlir/AffineOps/CMakeLists.txt b/mlir/include/mlir/AffineOps/CMakeLists.txt deleted file mode 100644 index 6c5a58c957b..00000000000 --- a/mlir/include/mlir/AffineOps/CMakeLists.txt +++ /dev/null @@ -1,4 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS AffineOps.td) -mlir_tablegen(AffineOps.h.inc -gen-op-decls) -mlir_tablegen(AffineOps.cpp.inc -gen-op-defs) -add_public_tablegen_target(MLIRAffineOpsIncGen) diff --git a/mlir/include/mlir/CMakeLists.txt b/mlir/include/mlir/CMakeLists.txt index 043db03641f..b393ea2c0e8 100644 --- a/mlir/include/mlir/CMakeLists.txt +++ b/mlir/include/mlir/CMakeLists.txt @@ -1,3 +1,2 @@ -add_subdirectory(AffineOps) add_subdirectory(Dialect) add_subdirectory(EDSC) diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.h b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h new file mode 100644 index 00000000000..a6af20eca0b --- /dev/null +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h @@ -0,0 +1,598 @@ +//===- AffineOps.h - MLIR Affine Operations -------------------------------===// +// +// Copyright 2019 The MLIR Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================= +// +// This file defines convenience types for working with Affine operations +// in the MLIR operation set. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_AFFINEOPS_AFFINEOPS_H +#define MLIR_DIALECT_AFFINEOPS_AFFINEOPS_H + +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/StandardTypes.h" + +namespace mlir { +class AffineBound; +class AffineValueMap; +class AffineTerminatorOp; +class FlatAffineConstraints; +class OpBuilder; + +/// A utility function to check if a value is defined at the top level of a +/// function. A value defined at the top level is always a valid symbol. +bool isTopLevelSymbol(Value *value); + +class AffineOpsDialect : public Dialect { +public: + AffineOpsDialect(MLIRContext *context); + static StringRef getDialectNamespace() { return "affine"; } +}; + +/// The "affine.apply" operation applies an affine map to a list of operands, +/// yielding a single result. The operand list must be the same size as the +/// number of arguments to the affine mapping. All operands and the result are +/// of type 'Index'. This operation requires a single affine map attribute named +/// "map". For example: +/// +/// %y = "affine.apply" (%x) { map: (d0) -> (d0 + 1) } : +/// (index) -> (index) +/// +/// equivalently: +/// +/// #map42 = (d0)->(d0+1) +/// %y = affine.apply #map42(%x) +/// +class AffineApplyOp : public Op { +public: + using Op::Op; + + /// Builds an affine apply op with the specified map and operands. + static void build(Builder *builder, OperationState *result, AffineMap map, + ArrayRef operands); + + /// Returns the affine map to be applied by this operation. + AffineMap getAffineMap() { + return getAttrOfType("map").getValue(); + } + + /// Returns true if the result of this operation can be used as dimension id. + bool isValidDim(); + + /// Returns true if the result of this operation is a symbol. + bool isValidSymbol(); + + static StringRef getOperationName() { return "affine.apply"; } + + // Hooks to customize behavior of this op. + static ParseResult parse(OpAsmParser *parser, OperationState *result); + void print(OpAsmPrinter *p); + LogicalResult verify(); + OpFoldResult fold(ArrayRef operands); + + static void getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context); +}; + +/// 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), tranferring 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'. +// +// For example, a DmaStartOp 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 = constant 256 +// %idx = 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(mlir-team): add additional operands to allow source and destination +// striding, and multiple stride levels (possibly using AffineMaps to specify +// multiple levels of striding). +// TODO(andydavis) Consider replacing src/dst memref indices with view memrefs. +class AffineDmaStartOp : public Op { +public: + using Op::Op; + + static void build(Builder *builder, OperationState *result, Value *srcMemRef, + AffineMap srcMap, ArrayRef srcIndices, + Value *destMemRef, AffineMap dstMap, + ArrayRef destIndices, Value *tagMemRef, + AffineMap tagMap, ArrayRef tagIndices, + Value *numElements, Value *stride = nullptr, + Value *elementsPerStride = nullptr); + + /// Returns the operand index of the src memref. + unsigned getSrcMemRefOperandIndex() { return 0; } + + /// Returns the source MemRefType for this DMA operation. + Value *getSrcMemRef() { return getOperand(getSrcMemRefOperandIndex()); } + MemRefType getSrcMemRefType() { + return getSrcMemRef()->getType().cast(); + } + + /// Returns the rank (number of indices) of the source MemRefType. + unsigned getSrcMemRefRank() { return getSrcMemRefType().getRank(); } + + /// Returns the affine map used to access the src memref. + AffineMap getSrcMap() { return getSrcMapAttr().getValue(); } + AffineMapAttr getSrcMapAttr() { + return getAttr(getSrcMapAttrName()).cast(); + } + + /// 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 src memref. + unsigned getSrcMemorySpace() { + return getSrcMemRef()->getType().cast().getMemorySpace(); + } + + /// Returns the operand index of the dst memref. + unsigned getDstMemRefOperandIndex() { + return getSrcMemRefOperandIndex() + 1 + getSrcMap().getNumInputs(); + } + + /// Returns the destination MemRefType for this DMA operations. + Value *getDstMemRef() { return getOperand(getDstMemRefOperandIndex()); } + MemRefType getDstMemRefType() { + return getDstMemRef()->getType().cast(); + } + + /// Returns the rank (number of indices) of the destination MemRefType. + unsigned getDstMemRefRank() { + return getDstMemRef()->getType().cast().getRank(); + } + + /// Returns the memory space of the src memref. + unsigned getDstMemorySpace() { + return getDstMemRef()->getType().cast().getMemorySpace(); + } + + /// Returns the affine map used to access the dst memref. + AffineMap getDstMap() { return getDstMapAttr().getValue(); } + AffineMapAttr getDstMapAttr() { + return getAttr(getDstMapAttrName()).cast(); + } + + /// 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(); + } + + /// Returns the rank (number of indices) of the tag MemRefType. + unsigned getTagMemRefRank() { + return getTagMemRef()->getType().cast().getRank(); + } + + /// Returns the affine map used to access the tag memref. + AffineMap getTagMap() { return getTagMapAttr().getValue(); } + AffineMapAttr getTagMapAttr() { + return getAttr(getTagMapAttrName()).cast(); + } + + /// 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()); + } + + /// Returns the AffineMapAttr associated with 'memref'. + NamedAttribute getAffineMapAttrForMemRef(Value *memref) { + if (memref == getSrcMemRef()) + return {Identifier::get(getSrcMapAttrName(), getContext()), + getSrcMapAttr()}; + else if (memref == getDstMemRef()) + return {Identifier::get(getDstMapAttrName(), getContext()), + getDstMapAttr()}; + assert(memref == getTagMemRef() && + "DmaStartOp expected source, destination or tag memref"); + return {Identifier::get(getTagMapAttrName(), getContext()), + 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(); + static void getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context); + + /// 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 { +public: + using Op::Op; + + static void build(Builder *builder, OperationState *result, Value *tagMemRef, + AffineMap tagMap, ArrayRef 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(); + } + + /// Returns the affine map used to access the tag memref. + AffineMap getTagMap() { return getTagMapAttr().getValue(); } + AffineMapAttr getTagMapAttr() { + return getAttr(getTagMapAttrName()).cast(); + } + + // 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().getRank(); + } + + /// Returns the AffineMapAttr associated with 'memref'. + NamedAttribute getAffineMapAttrForMemRef(Value *memref) { + assert(memref == getTagMemRef()); + return {Identifier::get(getTagMapAttrName(), getContext()), + getTagMapAttr()}; + } + + /// Returns the number of elements transferred in 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(); + static void getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context); +}; + +/// The "affine.load" op reads an element from a memref, where the index +/// for each memref dimension is an affine expression of loop induction +/// variables and symbols. The output of 'affine.load' is a new value with the +/// same type as the elements of the memref. An affine expression of loop IVs +/// and symbols must be specified for each dimension of the memref. The keyword +/// 'symbol' can be used to indicate SSA identifiers which are symbolic. +// +// Example 1: +// +// %1 = affine.load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32> +// +// Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'. +// +// %1 = affine.load %0[%i0 + symbol(%n), %i1 + symbol(%m)] +// : memref<100x100xf32> +// +class AffineLoadOp : public Op::Impl> { +public: + using Op::Op; + + /// Builds an affine load op with the specified map and operands. + static void build(Builder *builder, OperationState *result, AffineMap map, + ArrayRef operands); + /// Builds an affine load op an identify map and operands. + static void build(Builder *builder, OperationState *result, Value *memref, + ArrayRef indices = {}); + + /// Returns the operand index of the memref. + unsigned getMemRefOperandIndex() { return 0; } + + /// Get memref operand. + Value *getMemRef() { return getOperand(getMemRefOperandIndex()); } + void setMemRef(Value *value) { setOperand(getMemRefOperandIndex(), value); } + MemRefType getMemRefType() { + return getMemRef()->getType().cast(); + } + + /// Get affine map operands. + operand_range getIndices() { return llvm::drop_begin(getOperands(), 1); } + + /// Returns the affine map used to index the memref for this operation. + AffineMap getAffineMap() { return getAffineMapAttr().getValue(); } + AffineMapAttr getAffineMapAttr() { + return getAttr(getMapAttrName()).cast(); + } + + /// Returns the AffineMapAttr associated with 'memref'. + NamedAttribute getAffineMapAttrForMemRef(Value *memref) { + assert(memref == getMemRef()); + return {Identifier::get(getMapAttrName(), getContext()), + getAffineMapAttr()}; + } + + static StringRef getMapAttrName() { return "map"; } + static StringRef getOperationName() { return "affine.load"; } + + // Hooks to customize behavior of this op. + static ParseResult parse(OpAsmParser *parser, OperationState *result); + void print(OpAsmPrinter *p); + LogicalResult verify(); + static void getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context); +}; + +/// The "affine.store" op writes an element to a memref, where the index +/// for each memref dimension is an affine expression of loop induction +/// variables and symbols. The 'affine.store' op stores a new value which is the +/// same type as the elements of the memref. An affine expression of loop IVs +/// and symbols must be specified for each dimension of the memref. The keyword +/// 'symbol' can be used to indicate SSA identifiers which are symbolic. +// +// Example 1: +// +// affine.store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32> +// +// Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'. +// +// affine.store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] +// : memref<100x100xf32> +// +class AffineStoreOp : public Op::Impl> { +public: + using Op::Op; + + /// Builds an affine store operation with the specified map and operands. + static void build(Builder *builder, OperationState *result, + Value *valueToStore, AffineMap map, + ArrayRef operands); + /// Builds an affine store operation with an identity map and operands. + static void build(Builder *builder, OperationState *result, + Value *valueToStore, Value *memref, + ArrayRef operands); + + /// Get value to be stored by store operation. + Value *getValueToStore() { return getOperand(0); } + + /// Returns the operand index of the memref. + unsigned getMemRefOperandIndex() { return 1; } + + /// Get memref operand. + Value *getMemRef() { return getOperand(getMemRefOperandIndex()); } + void setMemRef(Value *value) { setOperand(getMemRefOperandIndex(), value); } + + MemRefType getMemRefType() { + return getMemRef()->getType().cast(); + } + + /// Get affine map operands. + operand_range getIndices() { return llvm::drop_begin(getOperands(), 2); } + + /// Returns the affine map used to index the memref for this operation. + AffineMap getAffineMap() { return getAffineMapAttr().getValue(); } + AffineMapAttr getAffineMapAttr() { + return getAttr(getMapAttrName()).cast(); + } + + /// Returns the AffineMapAttr associated with 'memref'. + NamedAttribute getAffineMapAttrForMemRef(Value *memref) { + assert(memref == getMemRef()); + return {Identifier::get(getMapAttrName(), getContext()), + getAffineMapAttr()}; + } + + static StringRef getMapAttrName() { return "map"; } + static StringRef getOperationName() { return "affine.store"; } + + // Hooks to customize behavior of this op. + static ParseResult parse(OpAsmParser *parser, OperationState *result); + void print(OpAsmPrinter *p); + LogicalResult verify(); + static void getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context); +}; + +/// Returns true if the given Value can be used as a dimension id. +bool isValidDim(Value *value); + +/// Returns true if the given Value can be used as a symbol. +bool isValidSymbol(Value *value); + +/// Modifies both `map` and `operands` in-place so as to: +/// 1. drop duplicate operands +/// 2. drop unused dims and symbols from map +void canonicalizeMapAndOperands(AffineMap *map, + llvm::SmallVectorImpl *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, + llvm::ArrayRef operands); + +/// 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, + llvm::SmallVectorImpl *operands); + +#define GET_OP_CLASSES +#include "mlir/Dialect/AffineOps/AffineOps.h.inc" + +/// Returns 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 forInsts, + SmallVectorImpl *ivs); + +/// 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; } + + /// Returns an AffineValueMap representing this bound. + AffineValueMap getAsAffineValueMap(); + + unsigned getNumOperands() { return opEnd - opStart; } + Value *getOperand(unsigned idx) { + return op.getOperation()->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 diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.td b/mlir/include/mlir/Dialect/AffineOps/AffineOps.td new file mode 100644 index 00000000000..237692c04a7 --- /dev/null +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOps.td @@ -0,0 +1,259 @@ +//===- AffineOps.td - Affine operation definitions ---------*- tablegen -*-===// +// +// Copyright 2019 The MLIR Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================= +// +// Defines MLIR affine operations. +// +//===----------------------------------------------------------------------===// + +#ifdef AFFINE_OPS +#else +#define AFFINE_OPS + +#ifdef OP_BASE +#else +include "mlir/IR/OpBase.td" +#endif // OP_BASE + +include "mlir/Dialect/AffineOps/AffineOpsBase.td" + +def Affine_Dialect : Dialect { + let name = "affine"; + let cppNamespace = ""; +} + +// Base class for Affine dialect ops. +class Affine_Op traits = []> : + Op { + // For every affine op, there needs to be a: + // * void print(OpAsmPrinter *p, ${C++ class of Op} op) + // * LogicalResult verify(${C++ class of Op} op) + // * ParseResult parse${C++ class of Op}(OpAsmParser *parser, + // OperationState *result) + // functions. + let printer = [{ return ::print(p, *this); }]; + let verifier = [{ return ::verify(*this); }]; + let parser = [{ return ::parse$cppClass(parser, result); }]; +} + +// Require regions to have affine terminator. +def ImplicitAffineTerminator + : SingleBlockImplicitTerminator<"AffineTerminatorOp">; + +def AffineForOp : Affine_Op<"for", [ImplicitAffineTerminator]> { + let summary = "for operation"; + let description = [{ + The "affine.for" operation represents an affine loop nest, defining an SSA + value for its induction variable. It has one region capturing the loop body. + The induction variable is represented as a argument of this region. This SSA + value always has type index, which is the size of the machine word. The + stride, represented by step, is a positive constant integer which defaults + to "1" if not present. The lower and upper bounds specify a half-open range: + the range includes the lower bound but does not include the upper bound. + + The body region must contain exactly one block that terminates with + "affine.terminator". Calling AffineForOp::build will create such region + and insert the terminator, so will the parsing even in cases if it is absent + from the custom format. + + The lower and upper bounds of a for operation are represented as an + application of an affine mapping to a list of SSA values passed to the map. + The same restrictions hold for these SSA values as for all bindings of SSA + values to dimensions and symbols. The affine mappings for the bounds may + return multiple results, in which case the max/min keywords are required + (for the lower/upper bound respectively), and the bound is the + maximum/minimum of the returned values. + + Example: + + affine.for %i = 1 to 10 { + ... + } + + }]; + let arguments = (ins Variadic); + let regions = (region SizedRegion<1>:$region); + + let skipDefaultBuilders = 1; + let builders = [ + OpBuilder<"Builder *builder, OperationState *result, " + "int64_t lowerBound, int64_t upperBound, int64_t step = 1">, + OpBuilder<"Builder *builder, OperationState *result, " + "ArrayRef lbOperands, AffineMap lbMap, " + "ArrayRef ubOperands, AffineMap ubMap, " + "int64_t step = 1"> + ]; + + let extraClassDeclaration = [{ + static StringRef getStepAttrName() { return "step"; } + static StringRef getLowerBoundAttrName() { return "lower_bound"; } + static StringRef getUpperBoundAttrName() { return "upper_bound"; } + + Block *getBody() { return ®ion().front(); } + Value *getInductionVar() { return getBody()->getArgument(0); } + OpBuilder getBodyBuilder() { + return OpBuilder(getBody(), std::prev(getBody()->end())); + } + + // TODO: provide iterators for the lower and upper bound operands + // if the current access via getLowerBound(), getUpperBound() is too slow. + + /// Returns operands for the lower bound map. + operand_range getLowerBoundOperands(); + + /// Returns operands for the upper bound map. + operand_range getUpperBoundOperands(); + + /// Returns information about the lower bound as a single object. + AffineBound getLowerBound(); + + /// Returns information about the upper bound as a single object. + AffineBound getUpperBound(); + + /// Returns loop step. + int64_t getStep() { + return getAttr(getStepAttrName()).cast().getInt(); + } + + /// Returns affine map for the lower bound. + AffineMap getLowerBoundMap() { return getLowerBoundMapAttr().getValue(); } + AffineMapAttr getLowerBoundMapAttr() { + return getAttr(getLowerBoundAttrName()).cast(); + } + /// Returns affine map for the upper bound. The upper bound is exclusive. + AffineMap getUpperBoundMap() { return getUpperBoundMapAttr().getValue(); } + AffineMapAttr getUpperBoundMapAttr() { + return getAttr(getUpperBoundAttrName()).cast(); + } + + /// Set lower bound. The new bound must have the same number of operands as + /// the current bound map. Otherwise, 'replaceForLowerBound' should be used. + void setLowerBound(ArrayRef operands, AffineMap map); + /// Set upper bound. The new bound must not have more operands than the + /// current bound map. Otherwise, 'replaceForUpperBound' should be used. + void setUpperBound(ArrayRef operands, AffineMap map); + + /// Set the lower bound map without changing operands. + void setLowerBoundMap(AffineMap map); + + /// Set the upper bound map without changing operands. + void setUpperBoundMap(AffineMap map); + + /// Set loop step. + void setStep(int64_t step) { + assert(step > 0 && "step has to be a positive integer constant"); + auto *context = getLowerBoundMap().getContext(); + setAttr(Identifier::get(getStepAttrName(), context), + IntegerAttr::get(IndexType::get(context), step)); + } + + /// Returns true if the lower bound is constant. + bool hasConstantLowerBound(); + /// Returns true if the upper bound is constant. + bool hasConstantUpperBound(); + /// Returns true if both bounds are constant. + bool hasConstantBounds() { + return hasConstantLowerBound() && hasConstantUpperBound(); + } + /// Returns the value of the constant lower bound. + /// Fails assertion if the bound is non-constant. + int64_t getConstantLowerBound(); + /// Returns the value of the constant upper bound. The upper bound is + /// exclusive. Fails assertion if the bound is non-constant. + int64_t getConstantUpperBound(); + /// Sets the lower bound to the given constant value. + void setConstantLowerBound(int64_t value); + /// Sets the upper bound to the given constant value. + void setConstantUpperBound(int64_t value); + + /// Returns true if both the lower and upper bound have the same operand + /// lists (same operands in the same order). + bool matchingBoundOperandList(); + }]; + + let hasCanonicalizer = 1; +} + +def AffineIfOp : Affine_Op<"if", [ImplicitAffineTerminator]> { + let summary = "if-then-else operation"; + let description = [{ + The "if" operation represents an if-then-else construct for conditionally + executing two regions of code. The operands to an if operation are an + IntegerSet condition and a set of symbol/dimension operands to the + condition set. The operation produces no results. For example: + + affine.if #set(%i) { + ... + } else { + ... + } + + The 'else' blocks to the if operation are optional, and may be omitted. For + example: + + affine.if #set(%i) { + ... + } + }]; + let arguments = (ins Variadic); + let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion); + + let skipDefaultBuilders = 1; + let builders = [ + OpBuilder<"Builder *builder, OperationState *result, " + "Value *cond, bool withElseRegion"> + ]; + + let extraClassDeclaration = [{ + static StringRef getConditionAttrName() { return "condition"; } + + IntegerSet getIntegerSet(); + void setIntegerSet(IntegerSet newSet); + + OpBuilder getThenBodyBuilder() { + assert(!thenRegion().empty() && "Unexpected empty 'then' region."); + Block &body = thenRegion().front(); + return OpBuilder(&body, std::prev(body.end())); + } + OpBuilder getElseBodyBuilder() { + assert(!elseRegion().empty() && "Unexpected empty 'else' region."); + Block &body = elseRegion().front(); + return OpBuilder(&body, std::prev(body.end())); + } + }]; +} + +def AffineTerminatorOp : + Affine_Op<"terminator", [Terminator]> { + let summary = "affine terminator operation"; + let description = [{ + Affine terminator is a special terminator operation for blocks inside affine + loops and branches. It unconditionally transmits the control flow to the + successor of the operation enclosing the region. + + This operation does _not_ have a custom syntax. However, affine control + operations omit the terminator in their custom syntax for brevity. + }]; + + // No custom parsing/printing form. + let parser = ?; + let printer = ?; + + // Fully specified by traits. + let verifier = ?; +} + +#endif // AFFINE_OPS diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOpsBase.td b/mlir/include/mlir/Dialect/AffineOps/AffineOpsBase.td new file mode 100644 index 00000000000..2ac1d379c12 --- /dev/null +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOpsBase.td @@ -0,0 +1,44 @@ +//===- AffineOpsBase.td - Affine operation definitions -----*- tablegen -*-===// +// +// Copyright 2019 The MLIR Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================= +// +// Defines base support for MLIR affine operations. +// +//===----------------------------------------------------------------------===// + +#ifdef AFFINE_OPS_BASE +#else +#define AFFINE_OPS_BASE + +#ifdef OP_BASE +#else +include "mlir/IR/OpBase.td" +#endif // OP_BASE + +// Attributes containing affine maps. +def AffineMapAttr : Attr< + CPred<"$_self.isa()">, "AffineMap attribute"> { + let storageType = [{ AffineMapAttr }]; + let returnType = [{ AffineMap }]; + let constBuilderCall = "$_builder.getAffineMapAttr($0)"; +} + +def AffineMapArrayAttr : TypedArrayAttrBase { + let constBuilderCall = "$_builder.getAffineMapArrayAttr($0)"; +} + +#endif // AFFINE_OPS_BASE diff --git a/mlir/include/mlir/Dialect/AffineOps/CMakeLists.txt b/mlir/include/mlir/Dialect/AffineOps/CMakeLists.txt new file mode 100644 index 00000000000..6c5a58c957b --- /dev/null +++ b/mlir/include/mlir/Dialect/AffineOps/CMakeLists.txt @@ -0,0 +1,4 @@ +set(LLVM_TARGET_DEFINITIONS AffineOps.td) +mlir_tablegen(AffineOps.h.inc -gen-op-decls) +mlir_tablegen(AffineOps.cpp.inc -gen-op-defs) +add_public_tablegen_target(MLIRAffineOpsIncGen) diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt index ce53bfc9a57..9235436995a 100644 --- a/mlir/include/mlir/Dialect/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/CMakeLists.txt @@ -1,3 +1,4 @@ +add_subdirectory(AffineOps) add_subdirectory(FxpMathOps) add_subdirectory(GPU) add_subdirectory(Linalg) diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td index 29977c1c637..47d30cf2836 100644 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td @@ -24,7 +24,7 @@ #else #define LINALG_LIBRARY_OPS -include "mlir/AffineOps/AffineOpsBase.td" +include "mlir/Dialect/AffineOps/AffineOpsBase.td" include "mlir/Dialect/Linalg/IR/LinalgBase.td" class LinalgParametricNativeOpTrait : diff --git a/mlir/include/mlir/EDSC/Builders.h b/mlir/include/mlir/EDSC/Builders.h index 29e2e9e1ea7..51c5c331fe9 100644 --- a/mlir/include/mlir/EDSC/Builders.h +++ b/mlir/include/mlir/EDSC/Builders.h @@ -23,7 +23,7 @@ #ifndef MLIR_EDSC_BUILDERS_H_ #define MLIR_EDSC_BUILDERS_H_ -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/Dialect/VectorOps/VectorOps.h" #include "mlir/IR/Builders.h" diff --git a/mlir/lib/AffineOps/AffineOps.cpp b/mlir/lib/AffineOps/AffineOps.cpp deleted file mode 100644 index b00a11083ec..00000000000 --- a/mlir/lib/AffineOps/AffineOps.cpp +++ /dev/null @@ -1,1764 +0,0 @@ -//===- AffineOps.cpp - MLIR Affine Operations -----------------------------===// -// -// Copyright 2019 The MLIR Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// ============================================================================= - -#include "mlir/AffineOps/AffineOps.h" -#include "mlir/Dialect/StandardOps/Ops.h" -#include "mlir/IR/Block.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/Function.h" -#include "mlir/IR/IntegerSet.h" -#include "mlir/IR/Matchers.h" -#include "mlir/IR/OpImplementation.h" -#include "mlir/IR/PatternMatch.h" -#include "llvm/ADT/SetVector.h" -#include "llvm/ADT/SmallBitVector.h" -#include "llvm/Support/Debug.h" -using namespace mlir; -using llvm::dbgs; - -#define DEBUG_TYPE "affine-analysis" - -//===----------------------------------------------------------------------===// -// AffineOpsDialect -//===----------------------------------------------------------------------===// - -AffineOpsDialect::AffineOpsDialect(MLIRContext *context) - : Dialect(getDialectNamespace(), context) { - addOperations(); -} - -/// A utility function to check if a given region is attached to a function. -static bool isFunctionRegion(Region *region) { - return llvm::isa(region->getParentOp()); -} - -/// A utility function to check if a value is defined at the top level of a -/// function. A value defined at the top level is always a valid symbol. -bool mlir::isTopLevelSymbol(Value *value) { - if (auto *arg = dyn_cast(value)) - return isFunctionRegion(arg->getOwner()->getParent()); - return isFunctionRegion(value->getDefiningOp()->getParentRegion()); -} - -// Value can be used as a dimension id if it is valid as a symbol, or -// it is an induction variable, or it is a result of affine apply operation -// with dimension id arguments. -bool mlir::isValidDim(Value *value) { - // The value must be an index type. - if (!value->getType().isIndex()) - return false; - - if (auto *op = value->getDefiningOp()) { - // Top level operation or constant operation is ok. - if (isFunctionRegion(op->getParentRegion()) || isa(op)) - return true; - // Affine apply operation is ok if all of its operands are ok. - if (auto applyOp = dyn_cast(op)) - return applyOp.isValidDim(); - // The dim op is okay if its operand memref/tensor is defined at the top - // level. - if (auto dimOp = dyn_cast(op)) - return isTopLevelSymbol(dimOp.getOperand()); - return false; - } - // This value is a block argument (which also includes 'affine.for' loop IVs). - return true; -} - -// Value can be used as a symbol if it is a constant, or it is defined at -// the top level, or it is a result of affine apply operation with symbol -// arguments. -bool mlir::isValidSymbol(Value *value) { - // The value must be an index type. - if (!value->getType().isIndex()) - return false; - - if (auto *op = value->getDefiningOp()) { - // Top level operation or constant operation is ok. - if (isFunctionRegion(op->getParentRegion()) || isa(op)) - return true; - // Affine apply operation is ok if all of its operands are ok. - if (auto applyOp = dyn_cast(op)) - return applyOp.isValidSymbol(); - // The dim op is okay if its operand memref/tensor is defined at the top - // level. - if (auto dimOp = dyn_cast(op)) - return isTopLevelSymbol(dimOp.getOperand()); - return false; - } - // Otherwise, check that the value is a top level symbol. - return isTopLevelSymbol(value); -} - -// Returns true if 'value' is a valid index to an affine operation (e.g. -// affine.load, affine.store, affine.dma_start, affine.dma_wait). -// Returns false otherwise. -static bool isValidAffineIndexOperand(Value *value) { - return isValidDim(value) || isValidSymbol(value); -} - -/// Utility function to verify that a set of operands are valid dimension and -/// symbol identifiers. The operands should be layed 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 -static LogicalResult -verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands, - unsigned numDims) { - unsigned opIt = 0; - for (auto *operand : operands) { - if (opIt++ < numDims) { - if (!isValidDim(operand)) - return op.emitOpError("operand cannot be used as a dimension id"); - } else if (!isValidSymbol(operand)) { - return op.emitOpError("operand cannot be used as a symbol"); - } - } - return success(); -} - -//===----------------------------------------------------------------------===// -// AffineApplyOp -//===----------------------------------------------------------------------===// - -void AffineApplyOp::build(Builder *builder, OperationState *result, - AffineMap map, ArrayRef operands) { - result->addOperands(operands); - result->types.append(map.getNumResults(), builder->getIndexType()); - result->addAttribute("map", builder->getAffineMapAttr(map)); -} - -ParseResult AffineApplyOp::parse(OpAsmParser *parser, OperationState *result) { - auto &builder = parser->getBuilder(); - auto affineIntTy = builder.getIndexType(); - - AffineMapAttr mapAttr; - unsigned numDims; - if (parser->parseAttribute(mapAttr, "map", result->attributes) || - parseDimAndSymbolList(parser, result->operands, numDims) || - parser->parseOptionalAttributeDict(result->attributes)) - return failure(); - auto map = mapAttr.getValue(); - - if (map.getNumDims() != numDims || - numDims + map.getNumSymbols() != result->operands.size()) { - return parser->emitError(parser->getNameLoc(), - "dimension or symbol index mismatch"); - } - - result->types.append(map.getNumResults(), affineIntTy); - return success(); -} - -void AffineApplyOp::print(OpAsmPrinter *p) { - *p << "affine.apply " << getAttr("map"); - printDimAndSymbolList(operand_begin(), operand_end(), - getAffineMap().getNumDims(), p); - p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{"map"}); -} - -LogicalResult AffineApplyOp::verify() { - // Check that affine map attribute was specified. - auto affineMapAttr = getAttrOfType("map"); - if (!affineMapAttr) - return emitOpError("requires an affine map"); - - // Check input and output dimensions match. - auto map = affineMapAttr.getValue(); - - // Verify that operand count matches affine map dimension and symbol count. - if (getNumOperands() != map.getNumDims() + map.getNumSymbols()) - return emitOpError( - "operand count and affine map dimension and symbol count must match"); - - // Verify that all operands are of `index` type. - for (Type t : getOperandTypes()) { - if (!t.isIndex()) - return emitOpError("operands must be of type 'index'"); - } - - if (!getResult()->getType().isIndex()) - return emitOpError("result must be of type 'index'"); - - // Verify that the operands are valid dimension and symbol identifiers. - if (failed(verifyDimAndSymbolIdentifiers(*this, getOperands(), - map.getNumDims()))) - return failure(); - - // Verify that the map only produces one result. - if (map.getNumResults() != 1) - return emitOpError("mapping must produce one value"); - - return success(); -} - -// The result of the affine apply operation can be used as a dimension id if it -// is a CFG value or if it is an Value, and all the operands are valid -// dimension ids. -bool AffineApplyOp::isValidDim() { - return llvm::all_of(getOperands(), - [](Value *op) { return mlir::isValidDim(op); }); -} - -// The result of the affine apply operation can be used as a symbol if it is -// a CFG value or if it is an Value, and all the operands are symbols. -bool AffineApplyOp::isValidSymbol() { - return llvm::all_of(getOperands(), - [](Value *op) { return mlir::isValidSymbol(op); }); -} - -OpFoldResult AffineApplyOp::fold(ArrayRef operands) { - auto map = getAffineMap(); - - // Fold dims and symbols to existing values. - auto expr = map.getResult(0); - if (auto dim = expr.dyn_cast()) - return getOperand(dim.getPosition()); - if (auto sym = expr.dyn_cast()) - return getOperand(map.getNumDims() + sym.getPosition()); - - // Otherwise, default to folding the map. - SmallVector result; - if (failed(map.constantFold(operands, result))) - return {}; - return result[0]; -} - -namespace { -/// An `AffineApplyNormalizer` is a helper class that is not visible to the user -/// and supports renumbering operands of AffineApplyOp. This acts as a -/// reindexing map of Value* to positional dims or symbols and allows -/// simplifications such as: -/// -/// ```mlir -/// %1 = affine.apply (d0, d1) -> (d0 - d1) (%0, %0) -/// ``` -/// -/// into: -/// -/// ```mlir -/// %1 = affine.apply () -> (0) -/// ``` -struct AffineApplyNormalizer { - AffineApplyNormalizer(AffineMap map, ArrayRef operands); - - /// Returns the AffineMap resulting from normalization. - AffineMap getAffineMap() { return affineMap; } - - SmallVector getOperands() { - SmallVector res(reorderedDims); - res.append(concatenatedSymbols.begin(), concatenatedSymbols.end()); - return res; - } - -private: - /// Helper function to insert `v` into the coordinate system of the current - /// AffineApplyNormalizer. Returns the AffineDimExpr with the corresponding - /// renumbered position. - AffineDimExpr renumberOneDim(Value *v); - - /// Given an `other` normalizer, this rewrites `other.affineMap` in the - /// coordinate system of the current AffineApplyNormalizer. - /// Returns the rewritten AffineMap and updates the dims and symbols of - /// `this`. - AffineMap renumber(const AffineApplyNormalizer &other); - - /// Maps of Value* to position in `affineMap`. - DenseMap dimValueToPosition; - - /// Ordered dims and symbols matching positional dims and symbols in - /// `affineMap`. - SmallVector reorderedDims; - SmallVector concatenatedSymbols; - - AffineMap affineMap; - - /// Used with RAII to control the depth at which AffineApply are composed - /// recursively. Only accepts depth 1 for now to allow a behavior where a - /// newly composed AffineApplyOp does not increase the length of the chain of - /// AffineApplyOps. Full composition is implemented iteratively on top of - /// this behavior. - static unsigned &affineApplyDepth() { - static thread_local unsigned depth = 0; - return depth; - } - static constexpr unsigned kMaxAffineApplyDepth = 1; - - AffineApplyNormalizer() { affineApplyDepth()++; } - -public: - ~AffineApplyNormalizer() { affineApplyDepth()--; } -}; -} // end anonymous namespace. - -AffineDimExpr AffineApplyNormalizer::renumberOneDim(Value *v) { - DenseMap::iterator iterPos; - bool inserted = false; - std::tie(iterPos, inserted) = - dimValueToPosition.insert(std::make_pair(v, dimValueToPosition.size())); - if (inserted) { - reorderedDims.push_back(v); - } - return getAffineDimExpr(iterPos->second, v->getContext()) - .cast(); -} - -AffineMap AffineApplyNormalizer::renumber(const AffineApplyNormalizer &other) { - SmallVector dimRemapping; - for (auto *v : other.reorderedDims) { - auto kvp = other.dimValueToPosition.find(v); - if (dimRemapping.size() <= kvp->second) - dimRemapping.resize(kvp->second + 1); - dimRemapping[kvp->second] = renumberOneDim(kvp->first); - } - unsigned numSymbols = concatenatedSymbols.size(); - unsigned numOtherSymbols = other.concatenatedSymbols.size(); - SmallVector symRemapping(numOtherSymbols); - for (unsigned idx = 0; idx < numOtherSymbols; ++idx) { - symRemapping[idx] = - getAffineSymbolExpr(idx + numSymbols, other.affineMap.getContext()); - } - concatenatedSymbols.insert(concatenatedSymbols.end(), - other.concatenatedSymbols.begin(), - other.concatenatedSymbols.end()); - auto map = other.affineMap; - return map.replaceDimsAndSymbols(dimRemapping, symRemapping, - dimRemapping.size(), symRemapping.size()); -} - -// Gather the positions of the operands that are produced by an AffineApplyOp. -static llvm::SetVector -indicesFromAffineApplyOp(ArrayRef operands) { - llvm::SetVector res; - for (auto en : llvm::enumerate(operands)) - if (isa_and_nonnull(en.value()->getDefiningOp())) - res.insert(en.index()); - return res; -} - -// Support the special case of a symbol coming from an AffineApplyOp that needs -// to be composed into the current AffineApplyOp. -// This case is handled by rewriting all such symbols into dims for the purpose -// of allowing mathematical AffineMap composition. -// Returns an AffineMap where symbols that come from an AffineApplyOp have been -// rewritten as dims and are ordered after the original dims. -// TODO(andydavis,ntv): This promotion makes AffineMap lose track of which -// symbols are represented as dims. This loss is static but can still be -// recovered dynamically (with `isValidSymbol`). Still this is annoying for the -// semi-affine map case. A dynamic canonicalization of all dims that are valid -// symbols (a.k.a `canonicalizePromotedSymbols`) into symbols helps and even -// results in better simplifications and foldings. But we should evaluate -// whether this behavior is what we really want after using more. -static AffineMap promoteComposedSymbolsAsDims(AffineMap map, - ArrayRef symbols) { - if (symbols.empty()) { - return map; - } - - // Sanity check on symbols. - for (auto *sym : symbols) { - assert(isValidSymbol(sym) && "Expected only valid symbols"); - (void)sym; - } - - // Extract the symbol positions that come from an AffineApplyOp and - // needs to be rewritten as dims. - auto symPositions = indicesFromAffineApplyOp(symbols); - if (symPositions.empty()) { - return map; - } - - // Create the new map by replacing each symbol at pos by the next new dim. - unsigned numDims = map.getNumDims(); - unsigned numSymbols = map.getNumSymbols(); - unsigned numNewDims = 0; - unsigned numNewSymbols = 0; - SmallVector symReplacements(numSymbols); - for (unsigned i = 0; i < numSymbols; ++i) { - symReplacements[i] = - symPositions.count(i) > 0 - ? getAffineDimExpr(numDims + numNewDims++, map.getContext()) - : getAffineSymbolExpr(numNewSymbols++, map.getContext()); - } - assert(numSymbols >= numNewDims); - AffineMap newMap = map.replaceDimsAndSymbols( - {}, symReplacements, numDims + numNewDims, numNewSymbols); - - return newMap; -} - -/// The AffineNormalizer composes AffineApplyOp recursively. Its purpose is to -/// keep a correspondence between the mathematical `map` and the `operands` of -/// a given AffineApplyOp. This correspondence is maintained by iterating over -/// the operands and forming an `auxiliaryMap` that can be composed -/// mathematically with `map`. To keep this correspondence in cases where -/// symbols are produced by affine.apply operations, we perform a local rewrite -/// of symbols as dims. -/// -/// Rationale for locally rewriting symbols as dims: -/// ================================================ -/// The mathematical composition of AffineMap must always concatenate symbols -/// because it does not have enough information to do otherwise. For example, -/// composing `(d0)[s0] -> (d0 + s0)` with itself must produce -/// `(d0)[s0, s1] -> (d0 + s0 + s1)`. -/// -/// The result is only equivalent to `(d0)[s0] -> (d0 + 2 * s0)` when -/// applied to the same mlir::Value* for both s0 and s1. -/// As a consequence mathematical composition of AffineMap always concatenates -/// symbols. -/// -/// When AffineMaps are used in AffineApplyOp however, they may specify -/// composition via symbols, which is ambiguous mathematically. This corner case -/// is handled by locally rewriting such symbols that come from AffineApplyOp -/// into dims and composing through dims. -/// TODO(andydavis, ntv): Composition via symbols comes at a significant code -/// complexity. Alternatively we should investigate whether we want to -/// explicitly disallow symbols coming from affine.apply and instead force the -/// user to compose symbols beforehand. The annoyances may be small (i.e. 1 or 2 -/// extra API calls for such uses, which haven't popped up until now) and the -/// benefit potentially big: simpler and more maintainable code for a -/// non-trivial, recursive, procedure. -AffineApplyNormalizer::AffineApplyNormalizer(AffineMap map, - ArrayRef operands) - : AffineApplyNormalizer() { - static_assert(kMaxAffineApplyDepth > 0, "kMaxAffineApplyDepth must be > 0"); - assert(map.getNumInputs() == operands.size() && - "number of operands does not match the number of map inputs"); - - LLVM_DEBUG(map.print(dbgs() << "\nInput map: ")); - - // Promote symbols that come from an AffineApplyOp to dims by rewriting the - // map to always refer to: - // (dims, symbols coming from AffineApplyOp, other symbols). - // The order of operands can remain unchanged. - // This is a simplification that relies on 2 ordering properties: - // 1. rewritten symbols always appear after the original dims in the map; - // 2. operands are traversed in order and either dispatched to: - // a. auxiliaryExprs (dims and symbols rewritten as dims); - // b. concatenatedSymbols (all other symbols) - // This allows operand order to remain unchanged. - unsigned numDimsBeforeRewrite = map.getNumDims(); - map = promoteComposedSymbolsAsDims(map, - operands.take_back(map.getNumSymbols())); - - LLVM_DEBUG(map.print(dbgs() << "\nRewritten map: ")); - - SmallVector auxiliaryExprs; - bool furtherCompose = (affineApplyDepth() <= kMaxAffineApplyDepth); - // We fully spell out the 2 cases below. In this particular instance a little - // code duplication greatly improves readability. - // Note that the first branch would disappear if we only supported full - // composition (i.e. infinite kMaxAffineApplyDepth). - if (!furtherCompose) { - // 1. Only dispatch dims or symbols. - for (auto en : llvm::enumerate(operands)) { - auto *t = en.value(); - assert(t->getType().isIndex()); - bool isDim = (en.index() < map.getNumDims()); - if (isDim) { - // a. The mathematical composition of AffineMap composes dims. - auxiliaryExprs.push_back(renumberOneDim(t)); - } else { - // b. The mathematical composition of AffineMap concatenates symbols. - // We do the same for symbol operands. - concatenatedSymbols.push_back(t); - } - } - } else { - assert(numDimsBeforeRewrite <= operands.size()); - // 2. Compose AffineApplyOps and dispatch dims or symbols. - for (unsigned i = 0, e = operands.size(); i < e; ++i) { - auto *t = operands[i]; - auto affineApply = dyn_cast_or_null(t->getDefiningOp()); - if (affineApply) { - // a. Compose affine.apply operations. - LLVM_DEBUG(affineApply.getOperation()->print( - dbgs() << "\nCompose AffineApplyOp recursively: ")); - AffineMap affineApplyMap = affineApply.getAffineMap(); - SmallVector affineApplyOperands( - affineApply.getOperands().begin(), affineApply.getOperands().end()); - AffineApplyNormalizer normalizer(affineApplyMap, affineApplyOperands); - - LLVM_DEBUG(normalizer.affineMap.print( - dbgs() << "\nRenumber into current normalizer: ")); - - auto renumberedMap = renumber(normalizer); - - LLVM_DEBUG( - renumberedMap.print(dbgs() << "\nRecursive composition yields: ")); - - auxiliaryExprs.push_back(renumberedMap.getResult(0)); - } else { - if (i < numDimsBeforeRewrite) { - // b. The mathematical composition of AffineMap composes dims. - auxiliaryExprs.push_back(renumberOneDim(t)); - } else { - // c. The mathematical composition of AffineMap concatenates symbols. - // We do the same for symbol operands. - concatenatedSymbols.push_back(t); - } - } - } - } - - // Early exit if `map` is already composed. - if (auxiliaryExprs.empty()) { - affineMap = map; - return; - } - - assert(concatenatedSymbols.size() >= map.getNumSymbols() && - "Unexpected number of concatenated symbols"); - auto numDims = dimValueToPosition.size(); - auto numSymbols = concatenatedSymbols.size() - map.getNumSymbols(); - auto auxiliaryMap = AffineMap::get(numDims, numSymbols, auxiliaryExprs); - - LLVM_DEBUG(map.print(dbgs() << "\nCompose map: ")); - LLVM_DEBUG(auxiliaryMap.print(dbgs() << "\nWith map: ")); - LLVM_DEBUG(map.compose(auxiliaryMap).print(dbgs() << "\nResult: ")); - - // TODO(andydavis,ntv): Disabling simplification results in major speed gains. - // Another option is to cache the results as it is expected a lot of redundant - // work is performed in practice. - affineMap = simplifyAffineMap(map.compose(auxiliaryMap)); - - LLVM_DEBUG(affineMap.print(dbgs() << "\nSimplified result: ")); - LLVM_DEBUG(dbgs() << "\n"); -} - -/// Implements `map` and `operands` composition and simplification to support -/// `makeComposedAffineApply`. This can be called to achieve the same effects -/// on `map` and `operands` without creating an AffineApplyOp that needs to be -/// immediately deleted. -static void composeAffineMapAndOperands(AffineMap *map, - SmallVectorImpl *operands) { - AffineApplyNormalizer normalizer(*map, *operands); - auto normalizedMap = normalizer.getAffineMap(); - auto normalizedOperands = normalizer.getOperands(); - canonicalizeMapAndOperands(&normalizedMap, &normalizedOperands); - *map = normalizedMap; - *operands = normalizedOperands; - assert(*map); -} - -void mlir::fullyComposeAffineMapAndOperands( - AffineMap *map, SmallVectorImpl *operands) { - while (llvm::any_of(*operands, [](Value *v) { - return isa_and_nonnull(v->getDefiningOp()); - })) { - composeAffineMapAndOperands(map, operands); - } -} - -AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc, - AffineMap map, - ArrayRef operands) { - AffineMap normalizedMap = map; - SmallVector normalizedOperands(operands.begin(), operands.end()); - composeAffineMapAndOperands(&normalizedMap, &normalizedOperands); - assert(normalizedMap); - return b.create(loc, normalizedMap, normalizedOperands); -} - -// A symbol may appear as a dim in affine.apply operations. This function -// canonicalizes dims that are valid symbols into actual symbols. -static void -canonicalizePromotedSymbols(AffineMap *map, - llvm::SmallVectorImpl *operands) { - if (!map || operands->empty()) - return; - - assert(map->getNumInputs() == operands->size() && - "map inputs must match number of operands"); - - auto *context = map->getContext(); - SmallVector resultOperands; - resultOperands.reserve(operands->size()); - SmallVector remappedSymbols; - remappedSymbols.reserve(operands->size()); - unsigned nextDim = 0; - unsigned nextSym = 0; - unsigned oldNumSyms = map->getNumSymbols(); - SmallVector dimRemapping(map->getNumDims()); - for (unsigned i = 0, e = map->getNumInputs(); i != e; ++i) { - if (i < map->getNumDims()) { - if (isValidSymbol((*operands)[i])) { - // This is a valid symbols that appears as a dim, canonicalize it. - dimRemapping[i] = getAffineSymbolExpr(oldNumSyms + nextSym++, context); - remappedSymbols.push_back((*operands)[i]); - } else { - dimRemapping[i] = getAffineDimExpr(nextDim++, context); - resultOperands.push_back((*operands)[i]); - } - } else { - resultOperands.push_back((*operands)[i]); - } - } - - resultOperands.append(remappedSymbols.begin(), remappedSymbols.end()); - *operands = resultOperands; - *map = map->replaceDimsAndSymbols(dimRemapping, {}, nextDim, - oldNumSyms + nextSym); - - assert(map->getNumInputs() == operands->size() && - "map inputs must match number of operands"); -} - -void mlir::canonicalizeMapAndOperands( - AffineMap *map, llvm::SmallVectorImpl *operands) { - if (!map || operands->empty()) - return; - - assert(map->getNumInputs() == operands->size() && - "map inputs must match number of operands"); - - canonicalizePromotedSymbols(map, operands); - - // Check to see what dims are used. - llvm::SmallBitVector usedDims(map->getNumDims()); - llvm::SmallBitVector usedSyms(map->getNumSymbols()); - map->walkExprs([&](AffineExpr expr) { - if (auto dimExpr = expr.dyn_cast()) - usedDims[dimExpr.getPosition()] = true; - else if (auto symExpr = expr.dyn_cast()) - usedSyms[symExpr.getPosition()] = true; - }); - - auto *context = map->getContext(); - - SmallVector resultOperands; - resultOperands.reserve(operands->size()); - - llvm::SmallDenseMap seenDims; - SmallVector dimRemapping(map->getNumDims()); - unsigned nextDim = 0; - for (unsigned i = 0, e = map->getNumDims(); i != e; ++i) { - if (usedDims[i]) { - auto it = seenDims.find((*operands)[i]); - if (it == seenDims.end()) { - dimRemapping[i] = getAffineDimExpr(nextDim++, context); - resultOperands.push_back((*operands)[i]); - seenDims.insert(std::make_pair((*operands)[i], dimRemapping[i])); - } else { - dimRemapping[i] = it->second; - } - } - } - llvm::SmallDenseMap seenSymbols; - SmallVector symRemapping(map->getNumSymbols()); - unsigned nextSym = 0; - for (unsigned i = 0, e = map->getNumSymbols(); i != e; ++i) { - if (usedSyms[i]) { - auto it = seenSymbols.find((*operands)[i + map->getNumDims()]); - if (it == seenSymbols.end()) { - symRemapping[i] = getAffineSymbolExpr(nextSym++, context); - resultOperands.push_back((*operands)[i + map->getNumDims()]); - seenSymbols.insert(std::make_pair((*operands)[i + map->getNumDims()], - symRemapping[i])); - } else { - symRemapping[i] = it->second; - } - } - } - *map = - map->replaceDimsAndSymbols(dimRemapping, symRemapping, nextDim, nextSym); - *operands = resultOperands; -} - -namespace { -/// Simplify AffineApply operations. -/// -struct SimplifyAffineApply : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - PatternMatchResult matchAndRewrite(AffineApplyOp apply, - PatternRewriter &rewriter) const override { - auto map = apply.getAffineMap(); - - AffineMap oldMap = map; - SmallVector resultOperands(apply.getOperands()); - composeAffineMapAndOperands(&map, &resultOperands); - if (map == oldMap) - return matchFailure(); - - rewriter.replaceOpWithNewOp(apply, map, resultOperands); - return matchSuccess(); - } -}; -} // end anonymous namespace. - -void AffineApplyOp::getCanonicalizationPatterns( - OwningRewritePatternList &results, MLIRContext *context) { - results.insert(context); -} - -//===----------------------------------------------------------------------===// -// Common canonicalization pattern support logic -//===----------------------------------------------------------------------===// - -namespace { -/// This is a common class used for patterns of the form -/// "someop(memrefcast) -> someop". It folds the source of any memref_cast -/// into the root operation directly. -struct MemRefCastFolder : public RewritePattern { - /// The rootOpName is the name of the root operation to match against. - MemRefCastFolder(StringRef rootOpName, MLIRContext *context) - : RewritePattern(rootOpName, 1, context) {} - - PatternMatchResult match(Operation *op) const override { - for (auto *operand : op->getOperands()) - if (matchPattern(operand, m_Op())) - return matchSuccess(); - - return matchFailure(); - } - - void rewrite(Operation *op, PatternRewriter &rewriter) const override { - for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i) - if (auto *memref = op->getOperand(i)->getDefiningOp()) - if (auto cast = dyn_cast(memref)) - op->setOperand(i, cast.getOperand()); - rewriter.updatedRootInPlace(op); - } -}; - -} // end anonymous namespace. - -//===----------------------------------------------------------------------===// -// AffineDmaStartOp -//===----------------------------------------------------------------------===// - -// TODO(b/133776335) Check that map operands are loop IVs or symbols. -void AffineDmaStartOp::build(Builder *builder, OperationState *result, - Value *srcMemRef, AffineMap srcMap, - ArrayRef srcIndices, Value *destMemRef, - AffineMap dstMap, ArrayRef destIndices, - Value *tagMemRef, AffineMap tagMap, - ArrayRef tagIndices, Value *numElements, - Value *stride, Value *elementsPerStride) { - result->addOperands(srcMemRef); - result->addAttribute(getSrcMapAttrName(), builder->getAffineMapAttr(srcMap)); - result->addOperands(srcIndices); - result->addOperands(destMemRef); - result->addAttribute(getDstMapAttrName(), builder->getAffineMapAttr(dstMap)); - result->addOperands(destIndices); - result->addOperands(tagMemRef); - result->addAttribute(getTagMapAttrName(), builder->getAffineMapAttr(tagMap)); - result->addOperands(tagIndices); - result->addOperands(numElements); - if (stride) { - result->addOperands({stride, elementsPerStride}); - } -} - -void AffineDmaStartOp::print(OpAsmPrinter *p) { - *p << "affine.dma_start " << *getSrcMemRef() << '['; - SmallVector operands(getSrcIndices()); - p->printAffineMapOfSSAIds(getSrcMapAttr(), operands); - *p << "], " << *getDstMemRef() << '['; - operands.assign(getDstIndices().begin(), getDstIndices().end()); - p->printAffineMapOfSSAIds(getDstMapAttr(), operands); - *p << "], " << *getTagMemRef() << '['; - operands.assign(getTagIndices().begin(), getTagIndices().end()); - p->printAffineMapOfSSAIds(getTagMapAttr(), operands); - *p << "], " << *getNumElements(); - if (isStrided()) { - *p << ", " << *getStride(); - *p << ", " << *getNumElementsPerStride(); - } - *p << " : " << getSrcMemRefType() << ", " << getDstMemRefType() << ", " - << getTagMemRefType(); -} - -// 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) { - OpAsmParser::OperandType srcMemRefInfo; - AffineMapAttr srcMapAttr; - SmallVector srcMapOperands; - OpAsmParser::OperandType dstMemRefInfo; - AffineMapAttr dstMapAttr; - SmallVector dstMapOperands; - OpAsmParser::OperandType tagMemRefInfo; - AffineMapAttr tagMapAttr; - SmallVector tagMapOperands; - OpAsmParser::OperandType numElementsInfo; - SmallVector strideInfo; - - SmallVector types; - auto indexType = parser->getBuilder().getIndexType(); - - // Parse and resolve the following list of operands: - // *) dst memref followed by its affine maps operands (in square brackets). - // *) src memref followed by its affine map operands (in square brackets). - // *) tag memref followed by its affine map operands (in square brackets). - // *) number of elements transferred by DMA operation. - if (parser->parseOperand(srcMemRefInfo) || - parser->parseAffineMapOfSSAIds(srcMapOperands, srcMapAttr, - getSrcMapAttrName(), result->attributes) || - parser->parseComma() || parser->parseOperand(dstMemRefInfo) || - parser->parseAffineMapOfSSAIds(dstMapOperands, dstMapAttr, - getDstMapAttrName(), result->attributes) || - parser->parseComma() || parser->parseOperand(tagMemRefInfo) || - parser->parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr, - getTagMapAttrName(), result->attributes) || - parser->parseComma() || parser->parseOperand(numElementsInfo)) - return failure(); - - // Parse optional stride and elements per stride. - if (parser->parseTrailingOperandList(strideInfo)) { - return failure(); - } - if (!strideInfo.empty() && strideInfo.size() != 2) { - return parser->emitError(parser->getNameLoc(), - "expected two stride related operands"); - } - bool isStrided = strideInfo.size() == 2; - - if (parser->parseColonTypeList(types)) - return failure(); - - if (types.size() != 3) - return parser->emitError(parser->getNameLoc(), "expected three types"); - - if (parser->resolveOperand(srcMemRefInfo, types[0], result->operands) || - parser->resolveOperands(srcMapOperands, indexType, result->operands) || - parser->resolveOperand(dstMemRefInfo, types[1], result->operands) || - parser->resolveOperands(dstMapOperands, indexType, result->operands) || - parser->resolveOperand(tagMemRefInfo, types[2], result->operands) || - parser->resolveOperands(tagMapOperands, indexType, result->operands) || - parser->resolveOperand(numElementsInfo, indexType, result->operands)) - return failure(); - - if (isStrided) { - if (parser->resolveOperands(strideInfo, indexType, result->operands)) - return failure(); - } - - // Check that src/dst/tag operand counts match their map.numInputs. - if (srcMapOperands.size() != srcMapAttr.getValue().getNumInputs() || - dstMapOperands.size() != dstMapAttr.getValue().getNumInputs() || - tagMapOperands.size() != tagMapAttr.getValue().getNumInputs()) - return parser->emitError(parser->getNameLoc(), - "memref operand count not equal to map.numInputs"); - return success(); -} - -LogicalResult AffineDmaStartOp::verify() { - if (!getOperand(getSrcMemRefOperandIndex())->getType().isa()) - return emitOpError("expected DMA source to be of memref type"); - if (!getOperand(getDstMemRefOperandIndex())->getType().isa()) - return emitOpError("expected DMA destination to be of memref type"); - if (!getOperand(getTagMemRefOperandIndex())->getType().isa()) - return emitOpError("expected DMA tag to be of memref type"); - - // DMAs from different memory spaces supported. - if (getSrcMemorySpace() == getDstMemorySpace()) { - return emitOpError("DMA should be between different memory spaces"); - } - unsigned numInputsAllMaps = getSrcMap().getNumInputs() + - getDstMap().getNumInputs() + - getTagMap().getNumInputs(); - if (getNumOperands() != numInputsAllMaps + 3 + 1 && - getNumOperands() != numInputsAllMaps + 3 + 1 + 2) { - return emitOpError("incorrect number of operands"); - } - - for (auto *idx : getSrcIndices()) { - if (!idx->getType().isIndex()) - return emitOpError("src index to dma_start must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) - return emitOpError("src index must be a dimension or symbol identifier"); - } - for (auto *idx : getDstIndices()) { - if (!idx->getType().isIndex()) - return emitOpError("dst index to dma_start must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) - return emitOpError("dst index must be a dimension or symbol identifier"); - } - for (auto *idx : getTagIndices()) { - if (!idx->getType().isIndex()) - return emitOpError("tag index to dma_start must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) - return emitOpError("tag index must be a dimension or symbol identifier"); - } - return success(); -} - -void AffineDmaStartOp::getCanonicalizationPatterns( - OwningRewritePatternList &results, MLIRContext *context) { - /// dma_start(memrefcast) -> dma_start - results.insert(getOperationName(), context); -} - -//===----------------------------------------------------------------------===// -// AffineDmaWaitOp -//===----------------------------------------------------------------------===// - -// TODO(b/133776335) Check that map operands are loop IVs or symbols. -void AffineDmaWaitOp::build(Builder *builder, OperationState *result, - Value *tagMemRef, AffineMap tagMap, - ArrayRef tagIndices, Value *numElements) { - result->addOperands(tagMemRef); - result->addAttribute(getTagMapAttrName(), builder->getAffineMapAttr(tagMap)); - result->addOperands(tagIndices); - result->addOperands(numElements); -} - -void AffineDmaWaitOp::print(OpAsmPrinter *p) { - *p << "affine.dma_wait " << *getTagMemRef() << '['; - SmallVector operands(getTagIndices()); - p->printAffineMapOfSSAIds(getTagMapAttr(), operands); - *p << "], "; - p->printOperand(getNumElements()); - *p << " : " << getTagMemRef()->getType(); -} - -// Parse AffineDmaWaitOp. -// Eg: -// affine.dma_wait %tag[%index], %num_elements -// : memref<1 x i32, (d0) -> (d0), 4> -// -ParseResult AffineDmaWaitOp::parse(OpAsmParser *parser, - OperationState *result) { - OpAsmParser::OperandType tagMemRefInfo; - AffineMapAttr tagMapAttr; - SmallVector tagMapOperands; - Type type; - auto indexType = parser->getBuilder().getIndexType(); - OpAsmParser::OperandType numElementsInfo; - - // Parse tag memref, its map operands, and dma size. - if (parser->parseOperand(tagMemRefInfo) || - parser->parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr, - getTagMapAttrName(), result->attributes) || - parser->parseComma() || parser->parseOperand(numElementsInfo) || - parser->parseColonType(type) || - parser->resolveOperand(tagMemRefInfo, type, result->operands) || - parser->resolveOperands(tagMapOperands, indexType, result->operands) || - parser->resolveOperand(numElementsInfo, indexType, result->operands)) - return failure(); - - if (!type.isa()) - return parser->emitError(parser->getNameLoc(), - "expected tag to be of memref type"); - - if (tagMapOperands.size() != tagMapAttr.getValue().getNumInputs()) - return parser->emitError(parser->getNameLoc(), - "tag memref operand count != to map.numInputs"); - return success(); -} - -LogicalResult AffineDmaWaitOp::verify() { - if (!getOperand(0)->getType().isa()) - return emitOpError("expected DMA tag to be of memref type"); - for (auto *idx : getTagIndices()) { - if (!idx->getType().isIndex()) - return emitOpError("index to dma_wait must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) - return emitOpError("index must be a dimension or symbol identifier"); - } - return success(); -} - -void AffineDmaWaitOp::getCanonicalizationPatterns( - OwningRewritePatternList &results, MLIRContext *context) { - /// dma_wait(memrefcast) -> dma_wait - results.insert(getOperationName(), context); -} - -//===----------------------------------------------------------------------===// -// AffineForOp -//===----------------------------------------------------------------------===// - -void AffineForOp::build(Builder *builder, OperationState *result, - ArrayRef lbOperands, AffineMap lbMap, - ArrayRef ubOperands, AffineMap ubMap, - int64_t step) { - assert(((!lbMap && lbOperands.empty()) || - lbOperands.size() == lbMap.getNumInputs()) && - "lower bound operand count does not match the affine map"); - assert(((!ubMap && ubOperands.empty()) || - ubOperands.size() == ubMap.getNumInputs()) && - "upper bound operand count does not match the affine map"); - assert(step > 0 && "step has to be a positive integer constant"); - - // Add an attribute for the step. - result->addAttribute(getStepAttrName(), - builder->getIntegerAttr(builder->getIndexType(), step)); - - // Add the lower bound. - result->addAttribute(getLowerBoundAttrName(), - builder->getAffineMapAttr(lbMap)); - result->addOperands(lbOperands); - - // Add the upper bound. - result->addAttribute(getUpperBoundAttrName(), - builder->getAffineMapAttr(ubMap)); - result->addOperands(ubOperands); - - // Create a region and a block for the body. The argument of the region is - // the loop induction variable. - Region *bodyRegion = result->addRegion(); - Block *body = new Block(); - body->addArgument(IndexType::get(builder->getContext())); - bodyRegion->push_back(body); - ensureTerminator(*bodyRegion, *builder, result->location); - - // Set the operands list as resizable so that we can freely modify the bounds. - result->setOperandListToResizable(); -} - -void AffineForOp::build(Builder *builder, OperationState *result, int64_t lb, - int64_t ub, int64_t step) { - auto lbMap = AffineMap::getConstantMap(lb, builder->getContext()); - auto ubMap = AffineMap::getConstantMap(ub, builder->getContext()); - return build(builder, result, {}, lbMap, {}, ubMap, step); -} - -static LogicalResult verify(AffineForOp op) { - // Check that the body defines as single block argument for the induction - // variable. - auto *body = op.getBody(); - if (body->getNumArguments() != 1 || - !body->getArgument(0)->getType().isIndex()) - return op.emitOpError( - "expected body to have a single index argument for the " - "induction variable"); - - // Verify that there are enough operands for the bounds. - AffineMap lowerBoundMap = op.getLowerBoundMap(), - upperBoundMap = op.getUpperBoundMap(); - if (op.getNumOperands() != - (lowerBoundMap.getNumInputs() + upperBoundMap.getNumInputs())) - return op.emitOpError( - "operand count must match with affine map dimension and symbol count"); - - // Verify that the bound operands are valid dimension/symbols. - /// Lower bound. - if (failed(verifyDimAndSymbolIdentifiers(op, op.getLowerBoundOperands(), - op.getLowerBoundMap().getNumDims()))) - return failure(); - /// Upper bound. - if (failed(verifyDimAndSymbolIdentifiers(op, op.getUpperBoundOperands(), - op.getUpperBoundMap().getNumDims()))) - return failure(); - return success(); -} - -/// Parse a for operation loop bounds. -static ParseResult parseBound(bool isLower, OperationState *result, - OpAsmParser *p) { - // 'min' / 'max' prefixes are generally syntactic sugar, but are required if - // the map has multiple results. - bool failedToParsedMinMax = - failed(p->parseOptionalKeyword(isLower ? "max" : "min")); - - auto &builder = p->getBuilder(); - auto boundAttrName = isLower ? AffineForOp::getLowerBoundAttrName() - : AffineForOp::getUpperBoundAttrName(); - - // Parse ssa-id as identity map. - SmallVector boundOpInfos; - if (p->parseOperandList(boundOpInfos)) - return failure(); - - if (!boundOpInfos.empty()) { - // Check that only one operand was parsed. - if (boundOpInfos.size() > 1) - return p->emitError(p->getNameLoc(), - "expected only one loop bound operand"); - - // TODO: improve error message when SSA value is not an affine integer. - // Currently it is 'use of value ... expects different type than prior uses' - if (p->resolveOperand(boundOpInfos.front(), builder.getIndexType(), - result->operands)) - return failure(); - - // Create an identity map using symbol id. This representation is optimized - // for storage. Analysis passes may expand it into a multi-dimensional map - // if desired. - AffineMap map = builder.getSymbolIdentityMap(); - result->addAttribute(boundAttrName, builder.getAffineMapAttr(map)); - return success(); - } - - // Get the attribute location. - llvm::SMLoc attrLoc = p->getCurrentLocation(); - - Attribute boundAttr; - if (p->parseAttribute(boundAttr, builder.getIndexType(), boundAttrName, - result->attributes)) - return failure(); - - // Parse full form - affine map followed by dim and symbol list. - if (auto affineMapAttr = boundAttr.dyn_cast()) { - unsigned currentNumOperands = result->operands.size(); - unsigned numDims; - if (parseDimAndSymbolList(p, result->operands, numDims)) - return failure(); - - auto map = affineMapAttr.getValue(); - if (map.getNumDims() != numDims) - return p->emitError( - p->getNameLoc(), - "dim operand count and integer set dim count must match"); - - unsigned numDimAndSymbolOperands = - result->operands.size() - currentNumOperands; - if (numDims + map.getNumSymbols() != numDimAndSymbolOperands) - return p->emitError( - p->getNameLoc(), - "symbol operand count and integer set symbol count must match"); - - // If the map has multiple results, make sure that we parsed the min/max - // prefix. - if (map.getNumResults() > 1 && failedToParsedMinMax) { - if (isLower) { - return p->emitError(attrLoc, "lower loop bound affine map with " - "multiple results requires 'max' prefix"); - } - return p->emitError(attrLoc, "upper loop bound affine map with multiple " - "results requires 'min' prefix"); - } - return success(); - } - - // Parse custom assembly form. - if (auto integerAttr = boundAttr.dyn_cast()) { - result->attributes.pop_back(); - result->addAttribute( - boundAttrName, builder.getAffineMapAttr( - builder.getConstantAffineMap(integerAttr.getInt()))); - return success(); - } - - return p->emitError( - p->getNameLoc(), - "expected valid affine map representation for loop bounds"); -} - -ParseResult parseAffineForOp(OpAsmParser *parser, OperationState *result) { - auto &builder = parser->getBuilder(); - OpAsmParser::OperandType inductionVariable; - // Parse the induction variable followed by '='. - if (parser->parseRegionArgument(inductionVariable) || parser->parseEqual()) - return failure(); - - // Parse loop bounds. - if (parseBound(/*isLower=*/true, result, parser) || - parser->parseKeyword("to", " between bounds") || - parseBound(/*isLower=*/false, result, parser)) - return failure(); - - // Parse the optional loop step, we default to 1 if one is not present. - if (parser->parseOptionalKeyword("step")) { - result->addAttribute( - AffineForOp::getStepAttrName(), - builder.getIntegerAttr(builder.getIndexType(), /*value=*/1)); - } else { - llvm::SMLoc stepLoc = parser->getCurrentLocation(); - IntegerAttr stepAttr; - if (parser->parseAttribute(stepAttr, builder.getIndexType(), - AffineForOp::getStepAttrName().data(), - result->attributes)) - return failure(); - - if (stepAttr.getValue().getSExtValue() < 0) - return parser->emitError( - stepLoc, - "expected step to be representable as a positive signed integer"); - } - - // Parse the body region. - Region *body = result->addRegion(); - if (parser->parseRegion(*body, inductionVariable, builder.getIndexType())) - return failure(); - - AffineForOp::ensureTerminator(*body, builder, result->location); - - // Parse the optional attribute list. - if (parser->parseOptionalAttributeDict(result->attributes)) - return failure(); - - // Set the operands list as resizable so that we can freely modify the bounds. - result->setOperandListToResizable(); - return success(); -} - -static void printBound(AffineMapAttr boundMap, - Operation::operand_range boundOperands, - const char *prefix, OpAsmPrinter *p) { - AffineMap map = boundMap.getValue(); - - // Check if this bound should be printed using custom assembly form. - // The decision to restrict printing custom assembly form to trivial cases - // comes from the will to roundtrip MLIR binary -> text -> binary in a - // lossless way. - // Therefore, custom assembly form parsing and printing is only supported for - // zero-operand constant maps and single symbol operand identity maps. - if (map.getNumResults() == 1) { - AffineExpr expr = map.getResult(0); - - // Print constant bound. - if (map.getNumDims() == 0 && map.getNumSymbols() == 0) { - if (auto constExpr = expr.dyn_cast()) { - *p << constExpr.getValue(); - return; - } - } - - // Print bound that consists of a single SSA symbol if the map is over a - // single symbol. - if (map.getNumDims() == 0 && map.getNumSymbols() == 1) { - if (auto symExpr = expr.dyn_cast()) { - p->printOperand(*boundOperands.begin()); - return; - } - } - } else { - // Map has multiple results. Print 'min' or 'max' prefix. - *p << prefix << ' '; - } - - // Print the map and its operands. - *p << boundMap; - printDimAndSymbolList(boundOperands.begin(), boundOperands.end(), - map.getNumDims(), p); -} - -void print(OpAsmPrinter *p, AffineForOp op) { - *p << "affine.for "; - p->printOperand(op.getBody()->getArgument(0)); - *p << " = "; - printBound(op.getLowerBoundMapAttr(), op.getLowerBoundOperands(), "max", p); - *p << " to "; - printBound(op.getUpperBoundMapAttr(), op.getUpperBoundOperands(), "min", p); - - if (op.getStep() != 1) - *p << " step " << op.getStep(); - p->printRegion(op.region(), - /*printEntryBlockArgs=*/false, - /*printBlockTerminators=*/false); - p->printOptionalAttrDict(op.getAttrs(), - /*elidedAttrs=*/{op.getLowerBoundAttrName(), - op.getUpperBoundAttrName(), - op.getStepAttrName()}); -} - -namespace { -/// This is a pattern to fold constant loop bounds. -struct AffineForLoopBoundFolder : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - PatternMatchResult matchAndRewrite(AffineForOp forOp, - PatternRewriter &rewriter) const override { - auto foldLowerOrUpperBound = [&forOp](bool lower) { - // Check to see if each of the operands is the result of a constant. If - // so, get the value. If not, ignore it. - SmallVector operandConstants; - auto boundOperands = - lower ? forOp.getLowerBoundOperands() : forOp.getUpperBoundOperands(); - for (auto *operand : boundOperands) { - Attribute operandCst; - matchPattern(operand, m_Constant(&operandCst)); - operandConstants.push_back(operandCst); - } - - AffineMap boundMap = - lower ? forOp.getLowerBoundMap() : forOp.getUpperBoundMap(); - assert(boundMap.getNumResults() >= 1 && - "bound maps should have at least one result"); - SmallVector foldedResults; - if (failed(boundMap.constantFold(operandConstants, foldedResults))) - return failure(); - - // Compute the max or min as applicable over the results. - assert(!foldedResults.empty() && - "bounds should have at least one result"); - auto maxOrMin = foldedResults[0].cast().getValue(); - for (unsigned i = 1, e = foldedResults.size(); i < e; i++) { - auto foldedResult = foldedResults[i].cast().getValue(); - maxOrMin = lower ? llvm::APIntOps::smax(maxOrMin, foldedResult) - : llvm::APIntOps::smin(maxOrMin, foldedResult); - } - lower ? forOp.setConstantLowerBound(maxOrMin.getSExtValue()) - : forOp.setConstantUpperBound(maxOrMin.getSExtValue()); - return success(); - }; - - // Try to fold the lower bound. - bool folded = false; - if (!forOp.hasConstantLowerBound()) - folded |= succeeded(foldLowerOrUpperBound(/*lower=*/true)); - - // Try to fold the upper bound. - if (!forOp.hasConstantUpperBound()) - folded |= succeeded(foldLowerOrUpperBound(/*lower=*/false)); - - // If any of the bounds were folded we return success. - if (!folded) - return matchFailure(); - rewriter.updatedRootInPlace(forOp); - return matchSuccess(); - } -}; -} // end anonymous namespace - -void AffineForOp::getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context) { - results.insert(context); -} - -AffineBound AffineForOp::getLowerBound() { - auto lbMap = getLowerBoundMap(); - return AffineBound(AffineForOp(*this), 0, lbMap.getNumInputs(), lbMap); -} - -AffineBound AffineForOp::getUpperBound() { - auto lbMap = getLowerBoundMap(); - auto ubMap = getUpperBoundMap(); - return AffineBound(AffineForOp(*this), lbMap.getNumInputs(), getNumOperands(), - ubMap); -} - -void AffineForOp::setLowerBound(ArrayRef lbOperands, AffineMap map) { - assert(lbOperands.size() == map.getNumInputs()); - assert(map.getNumResults() >= 1 && "bound map has at least one result"); - - SmallVector newOperands(lbOperands.begin(), lbOperands.end()); - - auto ubOperands = getUpperBoundOperands(); - newOperands.append(ubOperands.begin(), ubOperands.end()); - getOperation()->setOperands(newOperands); - - setAttr(getLowerBoundAttrName(), AffineMapAttr::get(map)); -} - -void AffineForOp::setUpperBound(ArrayRef ubOperands, AffineMap map) { - assert(ubOperands.size() == map.getNumInputs()); - assert(map.getNumResults() >= 1 && "bound map has at least one result"); - - SmallVector newOperands(getLowerBoundOperands()); - newOperands.append(ubOperands.begin(), ubOperands.end()); - getOperation()->setOperands(newOperands); - - setAttr(getUpperBoundAttrName(), AffineMapAttr::get(map)); -} - -void AffineForOp::setLowerBoundMap(AffineMap map) { - auto lbMap = getLowerBoundMap(); - assert(lbMap.getNumDims() == map.getNumDims() && - lbMap.getNumSymbols() == map.getNumSymbols()); - assert(map.getNumResults() >= 1 && "bound map has at least one result"); - (void)lbMap; - setAttr(getLowerBoundAttrName(), AffineMapAttr::get(map)); -} - -void AffineForOp::setUpperBoundMap(AffineMap map) { - auto ubMap = getUpperBoundMap(); - assert(ubMap.getNumDims() == map.getNumDims() && - ubMap.getNumSymbols() == map.getNumSymbols()); - assert(map.getNumResults() >= 1 && "bound map has at least one result"); - (void)ubMap; - setAttr(getUpperBoundAttrName(), AffineMapAttr::get(map)); -} - -bool AffineForOp::hasConstantLowerBound() { - return getLowerBoundMap().isSingleConstant(); -} - -bool AffineForOp::hasConstantUpperBound() { - return getUpperBoundMap().isSingleConstant(); -} - -int64_t AffineForOp::getConstantLowerBound() { - return getLowerBoundMap().getSingleConstantResult(); -} - -int64_t AffineForOp::getConstantUpperBound() { - return getUpperBoundMap().getSingleConstantResult(); -} - -void AffineForOp::setConstantLowerBound(int64_t value) { - setLowerBound({}, AffineMap::getConstantMap(value, getContext())); -} - -void AffineForOp::setConstantUpperBound(int64_t value) { - setUpperBound({}, AffineMap::getConstantMap(value, getContext())); -} - -AffineForOp::operand_range AffineForOp::getLowerBoundOperands() { - return {operand_begin(), operand_begin() + getLowerBoundMap().getNumInputs()}; -} - -AffineForOp::operand_range AffineForOp::getUpperBoundOperands() { - return {operand_begin() + getLowerBoundMap().getNumInputs(), operand_end()}; -} - -bool AffineForOp::matchingBoundOperandList() { - auto lbMap = getLowerBoundMap(); - auto ubMap = getUpperBoundMap(); - if (lbMap.getNumDims() != ubMap.getNumDims() || - lbMap.getNumSymbols() != ubMap.getNumSymbols()) - return false; - - unsigned numOperands = lbMap.getNumInputs(); - for (unsigned i = 0, e = lbMap.getNumInputs(); i < e; i++) { - // Compare Value *'s. - if (getOperand(i) != getOperand(numOperands + i)) - return false; - } - return true; -} - -/// Returns if the provided value is the induction variable of a AffineForOp. -bool mlir::isForInductionVar(Value *val) { - return getForInductionVarOwner(val) != AffineForOp(); -} - -/// Returns the loop parent of an induction variable. If the provided value is -/// not an induction variable, then return nullptr. -AffineForOp mlir::getForInductionVarOwner(Value *val) { - auto *ivArg = dyn_cast(val); - if (!ivArg || !ivArg->getOwner()) - return AffineForOp(); - auto *containingInst = ivArg->getOwner()->getParent()->getParentOp(); - return dyn_cast(containingInst); -} - -/// Extracts the induction variables from a list of AffineForOps and returns -/// them. -void mlir::extractForInductionVars(ArrayRef forInsts, - SmallVectorImpl *ivs) { - ivs->reserve(forInsts.size()); - for (auto forInst : forInsts) - ivs->push_back(forInst.getInductionVar()); -} - -//===----------------------------------------------------------------------===// -// AffineIfOp -//===----------------------------------------------------------------------===// - -static LogicalResult verify(AffineIfOp op) { - // Verify that we have a condition attribute. - auto conditionAttr = - op.getAttrOfType(op.getConditionAttrName()); - if (!conditionAttr) - return op.emitOpError( - "requires an integer set attribute named 'condition'"); - - // Verify that there are enough operands for the condition. - IntegerSet condition = conditionAttr.getValue(); - if (op.getNumOperands() != condition.getNumOperands()) - return op.emitOpError( - "operand count and condition integer set dimension and " - "symbol count must match"); - - // Verify that the operands are valid dimension/symbols. - if (failed(verifyDimAndSymbolIdentifiers( - op, op.getOperation()->getNonSuccessorOperands(), - condition.getNumDims()))) - return failure(); - - // Verify that the entry of each child region does not have arguments. - for (auto ®ion : op.getOperation()->getRegions()) { - for (auto &b : region) - if (b.getNumArguments() != 0) - return op.emitOpError( - "requires that child entry blocks have no arguments"); - } - return success(); -} - -ParseResult parseAffineIfOp(OpAsmParser *parser, OperationState *result) { - // Parse the condition attribute set. - IntegerSetAttr conditionAttr; - unsigned numDims; - if (parser->parseAttribute(conditionAttr, AffineIfOp::getConditionAttrName(), - result->attributes) || - parseDimAndSymbolList(parser, result->operands, numDims)) - return failure(); - - // Verify the condition operands. - auto set = conditionAttr.getValue(); - if (set.getNumDims() != numDims) - return parser->emitError( - parser->getNameLoc(), - "dim operand count and integer set dim count must match"); - if (numDims + set.getNumSymbols() != result->operands.size()) - return parser->emitError( - parser->getNameLoc(), - "symbol operand count and integer set symbol count must match"); - - // Create the regions for 'then' and 'else'. The latter must be created even - // if it remains empty for the validity of the operation. - result->regions.reserve(2); - Region *thenRegion = result->addRegion(); - Region *elseRegion = result->addRegion(); - - // Parse the 'then' region. - if (parser->parseRegion(*thenRegion, {}, {})) - return failure(); - AffineIfOp::ensureTerminator(*thenRegion, parser->getBuilder(), - result->location); - - // If we find an 'else' keyword then parse the 'else' region. - if (!parser->parseOptionalKeyword("else")) { - if (parser->parseRegion(*elseRegion, {}, {})) - return failure(); - AffineIfOp::ensureTerminator(*elseRegion, parser->getBuilder(), - result->location); - } - - // Parse the optional attribute list. - if (parser->parseOptionalAttributeDict(result->attributes)) - return failure(); - - return success(); -} - -void print(OpAsmPrinter *p, AffineIfOp op) { - auto conditionAttr = - op.getAttrOfType(op.getConditionAttrName()); - *p << "affine.if " << conditionAttr; - printDimAndSymbolList(op.operand_begin(), op.operand_end(), - conditionAttr.getValue().getNumDims(), p); - p->printRegion(op.thenRegion(), - /*printEntryBlockArgs=*/false, - /*printBlockTerminators=*/false); - - // Print the 'else' regions if it has any blocks. - auto &elseRegion = op.elseRegion(); - if (!elseRegion.empty()) { - *p << " else"; - p->printRegion(elseRegion, - /*printEntryBlockArgs=*/false, - /*printBlockTerminators=*/false); - } - - // Print the attribute list. - p->printOptionalAttrDict(op.getAttrs(), - /*elidedAttrs=*/op.getConditionAttrName()); -} - -IntegerSet AffineIfOp::getIntegerSet() { - return getAttrOfType(getConditionAttrName()).getValue(); -} -void AffineIfOp::setIntegerSet(IntegerSet newSet) { - setAttr(getConditionAttrName(), IntegerSetAttr::get(newSet)); -} - -//===----------------------------------------------------------------------===// -// AffineLoadOp -//===----------------------------------------------------------------------===// - -void AffineLoadOp::build(Builder *builder, OperationState *result, - AffineMap map, ArrayRef operands) { - result->addOperands(operands); - if (map) - result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); - auto memrefType = operands[0]->getType().cast(); - result->types.push_back(memrefType.getElementType()); -} - -void AffineLoadOp::build(Builder *builder, OperationState *result, - Value *memref, ArrayRef indices) { - result->addOperands(memref); - result->addOperands(indices); - auto memrefType = memref->getType().cast(); - auto rank = memrefType.getRank(); - // Create identity map for memrefs with at least one dimension or () -> () - // for zero-dimensional memrefs. - auto map = rank ? builder->getMultiDimIdentityMap(rank) - : builder->getEmptyAffineMap(); - result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); - result->types.push_back(memrefType.getElementType()); -} - -ParseResult AffineLoadOp::parse(OpAsmParser *parser, OperationState *result) { - auto &builder = parser->getBuilder(); - auto affineIntTy = builder.getIndexType(); - - MemRefType type; - OpAsmParser::OperandType memrefInfo; - AffineMapAttr mapAttr; - SmallVector mapOperands; - return failure( - parser->parseOperand(memrefInfo) || - parser->parseAffineMapOfSSAIds(mapOperands, mapAttr, getMapAttrName(), - result->attributes) || - parser->parseOptionalAttributeDict(result->attributes) || - parser->parseColonType(type) || - parser->resolveOperand(memrefInfo, type, result->operands) || - parser->resolveOperands(mapOperands, affineIntTy, result->operands) || - parser->addTypeToList(type.getElementType(), result->types)); -} - -void AffineLoadOp::print(OpAsmPrinter *p) { - *p << "affine.load " << *getMemRef() << '['; - AffineMapAttr mapAttr = getAttrOfType(getMapAttrName()); - if (mapAttr) { - SmallVector operands(getIndices()); - p->printAffineMapOfSSAIds(mapAttr, operands); - } - *p << ']'; - p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{getMapAttrName()}); - *p << " : " << getMemRefType(); -} - -LogicalResult AffineLoadOp::verify() { - if (getType() != getMemRefType().getElementType()) - return emitOpError("result type must match element type of memref"); - - auto mapAttr = getAttrOfType(getMapAttrName()); - if (mapAttr) { - AffineMap map = getAttrOfType(getMapAttrName()).getValue(); - if (map.getNumResults() != getMemRefType().getRank()) - return emitOpError("affine.load affine map num results must equal" - " memref rank"); - if (map.getNumInputs() != getNumOperands() - 1) - return emitOpError("expects as many subscripts as affine map inputs"); - } else { - if (getMemRefType().getRank() != getNumOperands() - 1) - return emitOpError( - "expects the number of subscripts to be equal to memref rank"); - } - - for (auto *idx : getIndices()) { - if (!idx->getType().isIndex()) - return emitOpError("index to load must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) - return emitOpError("index must be a dimension or symbol identifier"); - } - return success(); -} - -void AffineLoadOp::getCanonicalizationPatterns( - OwningRewritePatternList &results, MLIRContext *context) { - /// load(memrefcast) -> load - results.insert(getOperationName(), context); -} - -//===----------------------------------------------------------------------===// -// AffineStoreOp -//===----------------------------------------------------------------------===// - -void AffineStoreOp::build(Builder *builder, OperationState *result, - Value *valueToStore, AffineMap map, - ArrayRef operands) { - result->addOperands(valueToStore); - result->addOperands(operands); - if (map) - result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); -} - -void AffineStoreOp::build(Builder *builder, OperationState *result, - Value *valueToStore, Value *memref, - ArrayRef operands) { - result->addOperands(valueToStore); - result->addOperands(memref); - result->addOperands(operands); - auto memrefType = memref->getType().cast(); - auto rank = memrefType.getRank(); - // Create identity map for memrefs with at least one dimension or () -> () - // for zero-dimensional memrefs. - auto map = rank ? builder->getMultiDimIdentityMap(rank) - : builder->getEmptyAffineMap(); - result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); -} - -ParseResult AffineStoreOp::parse(OpAsmParser *parser, OperationState *result) { - auto affineIntTy = parser->getBuilder().getIndexType(); - - MemRefType type; - OpAsmParser::OperandType storeValueInfo; - OpAsmParser::OperandType memrefInfo; - AffineMapAttr mapAttr; - SmallVector mapOperands; - return failure( - parser->parseOperand(storeValueInfo) || parser->parseComma() || - parser->parseOperand(memrefInfo) || - parser->parseAffineMapOfSSAIds(mapOperands, mapAttr, getMapAttrName(), - result->attributes) || - parser->parseOptionalAttributeDict(result->attributes) || - parser->parseColonType(type) || - parser->resolveOperand(storeValueInfo, type.getElementType(), - result->operands) || - parser->resolveOperand(memrefInfo, type, result->operands) || - parser->resolveOperands(mapOperands, affineIntTy, result->operands)); -} - -void AffineStoreOp::print(OpAsmPrinter *p) { - *p << "affine.store " << *getValueToStore(); - *p << ", " << *getMemRef() << '['; - AffineMapAttr mapAttr = getAttrOfType(getMapAttrName()); - if (mapAttr) { - SmallVector operands(getIndices()); - p->printAffineMapOfSSAIds(mapAttr, operands); - } - *p << ']'; - p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{getMapAttrName()}); - *p << " : " << getMemRefType(); -} - -LogicalResult AffineStoreOp::verify() { - // First operand must have same type as memref element type. - if (getValueToStore()->getType() != getMemRefType().getElementType()) - return emitOpError("first operand must have same type memref element type"); - - auto mapAttr = getAttrOfType(getMapAttrName()); - if (mapAttr) { - AffineMap map = mapAttr.getValue(); - if (map.getNumResults() != getMemRefType().getRank()) - return emitOpError("affine.store affine map num results must equal" - " memref rank"); - if (map.getNumInputs() != getNumOperands() - 2) - return emitOpError("expects as many subscripts as affine map inputs"); - } else { - if (getMemRefType().getRank() != getNumOperands() - 2) - return emitOpError( - "expects the number of subscripts to be equal to memref rank"); - } - - for (auto *idx : getIndices()) { - if (!idx->getType().isIndex()) - return emitOpError("index to store must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) - return emitOpError("index must be a dimension or symbol identifier"); - } - return success(); -} - -void AffineStoreOp::getCanonicalizationPatterns( - OwningRewritePatternList &results, MLIRContext *context) { - /// load(memrefcast) -> load - results.insert(getOperationName(), context); -} - -#define GET_OP_CLASSES -#include "mlir/AffineOps/AffineOps.cpp.inc" diff --git a/mlir/lib/AffineOps/CMakeLists.txt b/mlir/lib/AffineOps/CMakeLists.txt deleted file mode 100644 index a8cf24e6c2b..00000000000 --- a/mlir/lib/AffineOps/CMakeLists.txt +++ /dev/null @@ -1,10 +0,0 @@ -add_llvm_library(MLIRAffineOps - AffineOps.cpp - DialectRegistration.cpp - - ADDITIONAL_HEADER_DIRS - ${MLIR_MAIN_INCLUDE_DIR}/mlir/AffineOps - ) -add_dependencies(MLIRAffineOps MLIRAffineOpsIncGen MLIRIR MLIRStandardOps) -target_link_libraries(MLIRAffineOps MLIRIR MLIRStandardOps) - diff --git a/mlir/lib/AffineOps/DialectRegistration.cpp b/mlir/lib/AffineOps/DialectRegistration.cpp deleted file mode 100644 index 0afb32c1bd6..00000000000 --- a/mlir/lib/AffineOps/DialectRegistration.cpp +++ /dev/null @@ -1,22 +0,0 @@ -//===- DialectRegistration.cpp - Register Affine Op dialect ---------------===// -// -// Copyright 2019 The MLIR Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// ============================================================================= - -#include "mlir/AffineOps/AffineOps.h" -using namespace mlir; - -// Static initialization for Affine op dialect registration. -static DialectRegistration StandardOps; diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp index e074e5d4405..92997ad27a7 100644 --- a/mlir/lib/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Analysis/AffineAnalysis.cpp @@ -21,9 +21,9 @@ //===----------------------------------------------------------------------===// #include "mlir/Analysis/AffineAnalysis.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/Builders.h" diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp index b1e818ac02c..70daca9754f 100644 --- a/mlir/lib/Analysis/AffineStructures.cpp +++ b/mlir/lib/Analysis/AffineStructures.cpp @@ -20,7 +20,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Analysis/AffineStructures.h" -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/AffineMap.h" diff --git a/mlir/lib/Analysis/LoopAnalysis.cpp b/mlir/lib/Analysis/LoopAnalysis.cpp index 1e1095743c9..21d47c3c1ea 100644 --- a/mlir/lib/Analysis/LoopAnalysis.cpp +++ b/mlir/lib/Analysis/LoopAnalysis.cpp @@ -21,11 +21,11 @@ #include "mlir/Analysis/LoopAnalysis.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/NestedMatcher.h" #include "mlir/Analysis/VectorAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/Dialect/VectorOps/VectorOps.h" #include "mlir/IR/AffineMap.h" diff --git a/mlir/lib/Analysis/MemRefBoundCheck.cpp b/mlir/lib/Analysis/MemRefBoundCheck.cpp index 85fe3109f6a..849407520da 100644 --- a/mlir/lib/Analysis/MemRefBoundCheck.cpp +++ b/mlir/lib/Analysis/MemRefBoundCheck.cpp @@ -20,11 +20,11 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Passes.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/lib/Analysis/NestedMatcher.cpp b/mlir/lib/Analysis/NestedMatcher.cpp index c7c0db90a7b..9d7d17f836c 100644 --- a/mlir/lib/Analysis/NestedMatcher.cpp +++ b/mlir/lib/Analysis/NestedMatcher.cpp @@ -16,7 +16,7 @@ // ============================================================================= #include "mlir/Analysis/NestedMatcher.h" -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "llvm/ADT/ArrayRef.h" diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp index c240d779c44..2f7eddf5ab3 100644 --- a/mlir/lib/Analysis/SliceAnalysis.cpp +++ b/mlir/lib/Analysis/SliceAnalysis.cpp @@ -20,8 +20,8 @@ //===----------------------------------------------------------------------===// #include "mlir/Analysis/SliceAnalysis.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/VectorAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/LoopOps/LoopOps.h" #include "mlir/IR/Function.h" #include "mlir/IR/Operation.h" diff --git a/mlir/lib/Analysis/TestMemRefDependenceCheck.cpp b/mlir/lib/Analysis/TestMemRefDependenceCheck.cpp index 9ecdcf7c2fe..477121fcc24 100644 --- a/mlir/lib/Analysis/TestMemRefDependenceCheck.cpp +++ b/mlir/lib/Analysis/TestMemRefDependenceCheck.cpp @@ -19,11 +19,11 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Passes.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/lib/Analysis/TestParallelismDetection.cpp b/mlir/lib/Analysis/TestParallelismDetection.cpp index 246cfbe9720..351a6a7a191 100644 --- a/mlir/lib/Analysis/TestParallelismDetection.cpp +++ b/mlir/lib/Analysis/TestParallelismDetection.cpp @@ -19,9 +19,9 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/Passes.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/IR/Builders.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp index d4fc42ceff7..aaefd98d1bd 100644 --- a/mlir/lib/Analysis/Utils.cpp +++ b/mlir/lib/Analysis/Utils.cpp @@ -22,9 +22,9 @@ #include "mlir/Analysis/Utils.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "llvm/ADT/DenseMap.h" diff --git a/mlir/lib/Analysis/VectorAnalysis.cpp b/mlir/lib/Analysis/VectorAnalysis.cpp index f34515f73a0..9846abb7be2 100644 --- a/mlir/lib/Analysis/VectorAnalysis.cpp +++ b/mlir/lib/Analysis/VectorAnalysis.cpp @@ -16,9 +16,9 @@ // ============================================================================= #include "mlir/Analysis/VectorAnalysis.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/LoopAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/Dialect/VectorOps/VectorOps.h" #include "mlir/IR/Builders.h" diff --git a/mlir/lib/CMakeLists.txt b/mlir/lib/CMakeLists.txt index bcb2d21d2da..f34b1e8bead 100644 --- a/mlir/lib/CMakeLists.txt +++ b/mlir/lib/CMakeLists.txt @@ -1,4 +1,3 @@ -add_subdirectory(AffineOps) add_subdirectory(Analysis) add_subdirectory(Conversion) add_subdirectory(Dialect) diff --git a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp index 13ba898dc44..154a8660bee 100644 --- a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp +++ b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPU.cpp @@ -22,7 +22,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Conversion/LoopsToGPU/LoopsToGPU.h" -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/GPU/GPUDialect.h" #include "mlir/Dialect/LoopOps/LoopOps.h" #include "mlir/Dialect/StandardOps/Ops.h" diff --git a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPUPass.cpp b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPUPass.cpp index 4b241e497c6..9dd9fdbbb87 100644 --- a/mlir/lib/Conversion/LoopsToGPU/LoopsToGPUPass.cpp +++ b/mlir/lib/Conversion/LoopsToGPU/LoopsToGPUPass.cpp @@ -16,8 +16,8 @@ // ============================================================================= #include "mlir/Conversion/LoopsToGPU/LoopsToGPUPass.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Conversion/LoopsToGPU/LoopsToGPU.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/LoopOps/LoopOps.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp new file mode 100644 index 00000000000..7db3fa07c52 --- /dev/null +++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp @@ -0,0 +1,1764 @@ +//===- AffineOps.cpp - MLIR Affine Operations -----------------------------===// +// +// Copyright 2019 The MLIR Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================= + +#include "mlir/Dialect/AffineOps/AffineOps.h" +#include "mlir/Dialect/StandardOps/Ops.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Function.h" +#include "mlir/IR/IntegerSet.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/IR/PatternMatch.h" +#include "llvm/ADT/SetVector.h" +#include "llvm/ADT/SmallBitVector.h" +#include "llvm/Support/Debug.h" +using namespace mlir; +using llvm::dbgs; + +#define DEBUG_TYPE "affine-analysis" + +//===----------------------------------------------------------------------===// +// AffineOpsDialect +//===----------------------------------------------------------------------===// + +AffineOpsDialect::AffineOpsDialect(MLIRContext *context) + : Dialect(getDialectNamespace(), context) { + addOperations(); +} + +/// A utility function to check if a given region is attached to a function. +static bool isFunctionRegion(Region *region) { + return llvm::isa(region->getParentOp()); +} + +/// A utility function to check if a value is defined at the top level of a +/// function. A value defined at the top level is always a valid symbol. +bool mlir::isTopLevelSymbol(Value *value) { + if (auto *arg = dyn_cast(value)) + return isFunctionRegion(arg->getOwner()->getParent()); + return isFunctionRegion(value->getDefiningOp()->getParentRegion()); +} + +// Value can be used as a dimension id if it is valid as a symbol, or +// it is an induction variable, or it is a result of affine apply operation +// with dimension id arguments. +bool mlir::isValidDim(Value *value) { + // The value must be an index type. + if (!value->getType().isIndex()) + return false; + + if (auto *op = value->getDefiningOp()) { + // Top level operation or constant operation is ok. + if (isFunctionRegion(op->getParentRegion()) || isa(op)) + return true; + // Affine apply operation is ok if all of its operands are ok. + if (auto applyOp = dyn_cast(op)) + return applyOp.isValidDim(); + // The dim op is okay if its operand memref/tensor is defined at the top + // level. + if (auto dimOp = dyn_cast(op)) + return isTopLevelSymbol(dimOp.getOperand()); + return false; + } + // This value is a block argument (which also includes 'affine.for' loop IVs). + return true; +} + +// Value can be used as a symbol if it is a constant, or it is defined at +// the top level, or it is a result of affine apply operation with symbol +// arguments. +bool mlir::isValidSymbol(Value *value) { + // The value must be an index type. + if (!value->getType().isIndex()) + return false; + + if (auto *op = value->getDefiningOp()) { + // Top level operation or constant operation is ok. + if (isFunctionRegion(op->getParentRegion()) || isa(op)) + return true; + // Affine apply operation is ok if all of its operands are ok. + if (auto applyOp = dyn_cast(op)) + return applyOp.isValidSymbol(); + // The dim op is okay if its operand memref/tensor is defined at the top + // level. + if (auto dimOp = dyn_cast(op)) + return isTopLevelSymbol(dimOp.getOperand()); + return false; + } + // Otherwise, check that the value is a top level symbol. + return isTopLevelSymbol(value); +} + +// Returns true if 'value' is a valid index to an affine operation (e.g. +// affine.load, affine.store, affine.dma_start, affine.dma_wait). +// Returns false otherwise. +static bool isValidAffineIndexOperand(Value *value) { + return isValidDim(value) || isValidSymbol(value); +} + +/// Utility function to verify that a set of operands are valid dimension and +/// symbol identifiers. The operands should be layed 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 +static LogicalResult +verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands, + unsigned numDims) { + unsigned opIt = 0; + for (auto *operand : operands) { + if (opIt++ < numDims) { + if (!isValidDim(operand)) + return op.emitOpError("operand cannot be used as a dimension id"); + } else if (!isValidSymbol(operand)) { + return op.emitOpError("operand cannot be used as a symbol"); + } + } + return success(); +} + +//===----------------------------------------------------------------------===// +// AffineApplyOp +//===----------------------------------------------------------------------===// + +void AffineApplyOp::build(Builder *builder, OperationState *result, + AffineMap map, ArrayRef operands) { + result->addOperands(operands); + result->types.append(map.getNumResults(), builder->getIndexType()); + result->addAttribute("map", builder->getAffineMapAttr(map)); +} + +ParseResult AffineApplyOp::parse(OpAsmParser *parser, OperationState *result) { + auto &builder = parser->getBuilder(); + auto affineIntTy = builder.getIndexType(); + + AffineMapAttr mapAttr; + unsigned numDims; + if (parser->parseAttribute(mapAttr, "map", result->attributes) || + parseDimAndSymbolList(parser, result->operands, numDims) || + parser->parseOptionalAttributeDict(result->attributes)) + return failure(); + auto map = mapAttr.getValue(); + + if (map.getNumDims() != numDims || + numDims + map.getNumSymbols() != result->operands.size()) { + return parser->emitError(parser->getNameLoc(), + "dimension or symbol index mismatch"); + } + + result->types.append(map.getNumResults(), affineIntTy); + return success(); +} + +void AffineApplyOp::print(OpAsmPrinter *p) { + *p << "affine.apply " << getAttr("map"); + printDimAndSymbolList(operand_begin(), operand_end(), + getAffineMap().getNumDims(), p); + p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{"map"}); +} + +LogicalResult AffineApplyOp::verify() { + // Check that affine map attribute was specified. + auto affineMapAttr = getAttrOfType("map"); + if (!affineMapAttr) + return emitOpError("requires an affine map"); + + // Check input and output dimensions match. + auto map = affineMapAttr.getValue(); + + // Verify that operand count matches affine map dimension and symbol count. + if (getNumOperands() != map.getNumDims() + map.getNumSymbols()) + return emitOpError( + "operand count and affine map dimension and symbol count must match"); + + // Verify that all operands are of `index` type. + for (Type t : getOperandTypes()) { + if (!t.isIndex()) + return emitOpError("operands must be of type 'index'"); + } + + if (!getResult()->getType().isIndex()) + return emitOpError("result must be of type 'index'"); + + // Verify that the operands are valid dimension and symbol identifiers. + if (failed(verifyDimAndSymbolIdentifiers(*this, getOperands(), + map.getNumDims()))) + return failure(); + + // Verify that the map only produces one result. + if (map.getNumResults() != 1) + return emitOpError("mapping must produce one value"); + + return success(); +} + +// The result of the affine apply operation can be used as a dimension id if it +// is a CFG value or if it is an Value, and all the operands are valid +// dimension ids. +bool AffineApplyOp::isValidDim() { + return llvm::all_of(getOperands(), + [](Value *op) { return mlir::isValidDim(op); }); +} + +// The result of the affine apply operation can be used as a symbol if it is +// a CFG value or if it is an Value, and all the operands are symbols. +bool AffineApplyOp::isValidSymbol() { + return llvm::all_of(getOperands(), + [](Value *op) { return mlir::isValidSymbol(op); }); +} + +OpFoldResult AffineApplyOp::fold(ArrayRef operands) { + auto map = getAffineMap(); + + // Fold dims and symbols to existing values. + auto expr = map.getResult(0); + if (auto dim = expr.dyn_cast()) + return getOperand(dim.getPosition()); + if (auto sym = expr.dyn_cast()) + return getOperand(map.getNumDims() + sym.getPosition()); + + // Otherwise, default to folding the map. + SmallVector result; + if (failed(map.constantFold(operands, result))) + return {}; + return result[0]; +} + +namespace { +/// An `AffineApplyNormalizer` is a helper class that is not visible to the user +/// and supports renumbering operands of AffineApplyOp. This acts as a +/// reindexing map of Value* to positional dims or symbols and allows +/// simplifications such as: +/// +/// ```mlir +/// %1 = affine.apply (d0, d1) -> (d0 - d1) (%0, %0) +/// ``` +/// +/// into: +/// +/// ```mlir +/// %1 = affine.apply () -> (0) +/// ``` +struct AffineApplyNormalizer { + AffineApplyNormalizer(AffineMap map, ArrayRef operands); + + /// Returns the AffineMap resulting from normalization. + AffineMap getAffineMap() { return affineMap; } + + SmallVector getOperands() { + SmallVector res(reorderedDims); + res.append(concatenatedSymbols.begin(), concatenatedSymbols.end()); + return res; + } + +private: + /// Helper function to insert `v` into the coordinate system of the current + /// AffineApplyNormalizer. Returns the AffineDimExpr with the corresponding + /// renumbered position. + AffineDimExpr renumberOneDim(Value *v); + + /// Given an `other` normalizer, this rewrites `other.affineMap` in the + /// coordinate system of the current AffineApplyNormalizer. + /// Returns the rewritten AffineMap and updates the dims and symbols of + /// `this`. + AffineMap renumber(const AffineApplyNormalizer &other); + + /// Maps of Value* to position in `affineMap`. + DenseMap dimValueToPosition; + + /// Ordered dims and symbols matching positional dims and symbols in + /// `affineMap`. + SmallVector reorderedDims; + SmallVector concatenatedSymbols; + + AffineMap affineMap; + + /// Used with RAII to control the depth at which AffineApply are composed + /// recursively. Only accepts depth 1 for now to allow a behavior where a + /// newly composed AffineApplyOp does not increase the length of the chain of + /// AffineApplyOps. Full composition is implemented iteratively on top of + /// this behavior. + static unsigned &affineApplyDepth() { + static thread_local unsigned depth = 0; + return depth; + } + static constexpr unsigned kMaxAffineApplyDepth = 1; + + AffineApplyNormalizer() { affineApplyDepth()++; } + +public: + ~AffineApplyNormalizer() { affineApplyDepth()--; } +}; +} // end anonymous namespace. + +AffineDimExpr AffineApplyNormalizer::renumberOneDim(Value *v) { + DenseMap::iterator iterPos; + bool inserted = false; + std::tie(iterPos, inserted) = + dimValueToPosition.insert(std::make_pair(v, dimValueToPosition.size())); + if (inserted) { + reorderedDims.push_back(v); + } + return getAffineDimExpr(iterPos->second, v->getContext()) + .cast(); +} + +AffineMap AffineApplyNormalizer::renumber(const AffineApplyNormalizer &other) { + SmallVector dimRemapping; + for (auto *v : other.reorderedDims) { + auto kvp = other.dimValueToPosition.find(v); + if (dimRemapping.size() <= kvp->second) + dimRemapping.resize(kvp->second + 1); + dimRemapping[kvp->second] = renumberOneDim(kvp->first); + } + unsigned numSymbols = concatenatedSymbols.size(); + unsigned numOtherSymbols = other.concatenatedSymbols.size(); + SmallVector symRemapping(numOtherSymbols); + for (unsigned idx = 0; idx < numOtherSymbols; ++idx) { + symRemapping[idx] = + getAffineSymbolExpr(idx + numSymbols, other.affineMap.getContext()); + } + concatenatedSymbols.insert(concatenatedSymbols.end(), + other.concatenatedSymbols.begin(), + other.concatenatedSymbols.end()); + auto map = other.affineMap; + return map.replaceDimsAndSymbols(dimRemapping, symRemapping, + dimRemapping.size(), symRemapping.size()); +} + +// Gather the positions of the operands that are produced by an AffineApplyOp. +static llvm::SetVector +indicesFromAffineApplyOp(ArrayRef operands) { + llvm::SetVector res; + for (auto en : llvm::enumerate(operands)) + if (isa_and_nonnull(en.value()->getDefiningOp())) + res.insert(en.index()); + return res; +} + +// Support the special case of a symbol coming from an AffineApplyOp that needs +// to be composed into the current AffineApplyOp. +// This case is handled by rewriting all such symbols into dims for the purpose +// of allowing mathematical AffineMap composition. +// Returns an AffineMap where symbols that come from an AffineApplyOp have been +// rewritten as dims and are ordered after the original dims. +// TODO(andydavis,ntv): This promotion makes AffineMap lose track of which +// symbols are represented as dims. This loss is static but can still be +// recovered dynamically (with `isValidSymbol`). Still this is annoying for the +// semi-affine map case. A dynamic canonicalization of all dims that are valid +// symbols (a.k.a `canonicalizePromotedSymbols`) into symbols helps and even +// results in better simplifications and foldings. But we should evaluate +// whether this behavior is what we really want after using more. +static AffineMap promoteComposedSymbolsAsDims(AffineMap map, + ArrayRef symbols) { + if (symbols.empty()) { + return map; + } + + // Sanity check on symbols. + for (auto *sym : symbols) { + assert(isValidSymbol(sym) && "Expected only valid symbols"); + (void)sym; + } + + // Extract the symbol positions that come from an AffineApplyOp and + // needs to be rewritten as dims. + auto symPositions = indicesFromAffineApplyOp(symbols); + if (symPositions.empty()) { + return map; + } + + // Create the new map by replacing each symbol at pos by the next new dim. + unsigned numDims = map.getNumDims(); + unsigned numSymbols = map.getNumSymbols(); + unsigned numNewDims = 0; + unsigned numNewSymbols = 0; + SmallVector symReplacements(numSymbols); + for (unsigned i = 0; i < numSymbols; ++i) { + symReplacements[i] = + symPositions.count(i) > 0 + ? getAffineDimExpr(numDims + numNewDims++, map.getContext()) + : getAffineSymbolExpr(numNewSymbols++, map.getContext()); + } + assert(numSymbols >= numNewDims); + AffineMap newMap = map.replaceDimsAndSymbols( + {}, symReplacements, numDims + numNewDims, numNewSymbols); + + return newMap; +} + +/// The AffineNormalizer composes AffineApplyOp recursively. Its purpose is to +/// keep a correspondence between the mathematical `map` and the `operands` of +/// a given AffineApplyOp. This correspondence is maintained by iterating over +/// the operands and forming an `auxiliaryMap` that can be composed +/// mathematically with `map`. To keep this correspondence in cases where +/// symbols are produced by affine.apply operations, we perform a local rewrite +/// of symbols as dims. +/// +/// Rationale for locally rewriting symbols as dims: +/// ================================================ +/// The mathematical composition of AffineMap must always concatenate symbols +/// because it does not have enough information to do otherwise. For example, +/// composing `(d0)[s0] -> (d0 + s0)` with itself must produce +/// `(d0)[s0, s1] -> (d0 + s0 + s1)`. +/// +/// The result is only equivalent to `(d0)[s0] -> (d0 + 2 * s0)` when +/// applied to the same mlir::Value* for both s0 and s1. +/// As a consequence mathematical composition of AffineMap always concatenates +/// symbols. +/// +/// When AffineMaps are used in AffineApplyOp however, they may specify +/// composition via symbols, which is ambiguous mathematically. This corner case +/// is handled by locally rewriting such symbols that come from AffineApplyOp +/// into dims and composing through dims. +/// TODO(andydavis, ntv): Composition via symbols comes at a significant code +/// complexity. Alternatively we should investigate whether we want to +/// explicitly disallow symbols coming from affine.apply and instead force the +/// user to compose symbols beforehand. The annoyances may be small (i.e. 1 or 2 +/// extra API calls for such uses, which haven't popped up until now) and the +/// benefit potentially big: simpler and more maintainable code for a +/// non-trivial, recursive, procedure. +AffineApplyNormalizer::AffineApplyNormalizer(AffineMap map, + ArrayRef operands) + : AffineApplyNormalizer() { + static_assert(kMaxAffineApplyDepth > 0, "kMaxAffineApplyDepth must be > 0"); + assert(map.getNumInputs() == operands.size() && + "number of operands does not match the number of map inputs"); + + LLVM_DEBUG(map.print(dbgs() << "\nInput map: ")); + + // Promote symbols that come from an AffineApplyOp to dims by rewriting the + // map to always refer to: + // (dims, symbols coming from AffineApplyOp, other symbols). + // The order of operands can remain unchanged. + // This is a simplification that relies on 2 ordering properties: + // 1. rewritten symbols always appear after the original dims in the map; + // 2. operands are traversed in order and either dispatched to: + // a. auxiliaryExprs (dims and symbols rewritten as dims); + // b. concatenatedSymbols (all other symbols) + // This allows operand order to remain unchanged. + unsigned numDimsBeforeRewrite = map.getNumDims(); + map = promoteComposedSymbolsAsDims(map, + operands.take_back(map.getNumSymbols())); + + LLVM_DEBUG(map.print(dbgs() << "\nRewritten map: ")); + + SmallVector auxiliaryExprs; + bool furtherCompose = (affineApplyDepth() <= kMaxAffineApplyDepth); + // We fully spell out the 2 cases below. In this particular instance a little + // code duplication greatly improves readability. + // Note that the first branch would disappear if we only supported full + // composition (i.e. infinite kMaxAffineApplyDepth). + if (!furtherCompose) { + // 1. Only dispatch dims or symbols. + for (auto en : llvm::enumerate(operands)) { + auto *t = en.value(); + assert(t->getType().isIndex()); + bool isDim = (en.index() < map.getNumDims()); + if (isDim) { + // a. The mathematical composition of AffineMap composes dims. + auxiliaryExprs.push_back(renumberOneDim(t)); + } else { + // b. The mathematical composition of AffineMap concatenates symbols. + // We do the same for symbol operands. + concatenatedSymbols.push_back(t); + } + } + } else { + assert(numDimsBeforeRewrite <= operands.size()); + // 2. Compose AffineApplyOps and dispatch dims or symbols. + for (unsigned i = 0, e = operands.size(); i < e; ++i) { + auto *t = operands[i]; + auto affineApply = dyn_cast_or_null(t->getDefiningOp()); + if (affineApply) { + // a. Compose affine.apply operations. + LLVM_DEBUG(affineApply.getOperation()->print( + dbgs() << "\nCompose AffineApplyOp recursively: ")); + AffineMap affineApplyMap = affineApply.getAffineMap(); + SmallVector affineApplyOperands( + affineApply.getOperands().begin(), affineApply.getOperands().end()); + AffineApplyNormalizer normalizer(affineApplyMap, affineApplyOperands); + + LLVM_DEBUG(normalizer.affineMap.print( + dbgs() << "\nRenumber into current normalizer: ")); + + auto renumberedMap = renumber(normalizer); + + LLVM_DEBUG( + renumberedMap.print(dbgs() << "\nRecursive composition yields: ")); + + auxiliaryExprs.push_back(renumberedMap.getResult(0)); + } else { + if (i < numDimsBeforeRewrite) { + // b. The mathematical composition of AffineMap composes dims. + auxiliaryExprs.push_back(renumberOneDim(t)); + } else { + // c. The mathematical composition of AffineMap concatenates symbols. + // We do the same for symbol operands. + concatenatedSymbols.push_back(t); + } + } + } + } + + // Early exit if `map` is already composed. + if (auxiliaryExprs.empty()) { + affineMap = map; + return; + } + + assert(concatenatedSymbols.size() >= map.getNumSymbols() && + "Unexpected number of concatenated symbols"); + auto numDims = dimValueToPosition.size(); + auto numSymbols = concatenatedSymbols.size() - map.getNumSymbols(); + auto auxiliaryMap = AffineMap::get(numDims, numSymbols, auxiliaryExprs); + + LLVM_DEBUG(map.print(dbgs() << "\nCompose map: ")); + LLVM_DEBUG(auxiliaryMap.print(dbgs() << "\nWith map: ")); + LLVM_DEBUG(map.compose(auxiliaryMap).print(dbgs() << "\nResult: ")); + + // TODO(andydavis,ntv): Disabling simplification results in major speed gains. + // Another option is to cache the results as it is expected a lot of redundant + // work is performed in practice. + affineMap = simplifyAffineMap(map.compose(auxiliaryMap)); + + LLVM_DEBUG(affineMap.print(dbgs() << "\nSimplified result: ")); + LLVM_DEBUG(dbgs() << "\n"); +} + +/// Implements `map` and `operands` composition and simplification to support +/// `makeComposedAffineApply`. This can be called to achieve the same effects +/// on `map` and `operands` without creating an AffineApplyOp that needs to be +/// immediately deleted. +static void composeAffineMapAndOperands(AffineMap *map, + SmallVectorImpl *operands) { + AffineApplyNormalizer normalizer(*map, *operands); + auto normalizedMap = normalizer.getAffineMap(); + auto normalizedOperands = normalizer.getOperands(); + canonicalizeMapAndOperands(&normalizedMap, &normalizedOperands); + *map = normalizedMap; + *operands = normalizedOperands; + assert(*map); +} + +void mlir::fullyComposeAffineMapAndOperands( + AffineMap *map, SmallVectorImpl *operands) { + while (llvm::any_of(*operands, [](Value *v) { + return isa_and_nonnull(v->getDefiningOp()); + })) { + composeAffineMapAndOperands(map, operands); + } +} + +AffineApplyOp mlir::makeComposedAffineApply(OpBuilder &b, Location loc, + AffineMap map, + ArrayRef operands) { + AffineMap normalizedMap = map; + SmallVector normalizedOperands(operands.begin(), operands.end()); + composeAffineMapAndOperands(&normalizedMap, &normalizedOperands); + assert(normalizedMap); + return b.create(loc, normalizedMap, normalizedOperands); +} + +// A symbol may appear as a dim in affine.apply operations. This function +// canonicalizes dims that are valid symbols into actual symbols. +static void +canonicalizePromotedSymbols(AffineMap *map, + llvm::SmallVectorImpl *operands) { + if (!map || operands->empty()) + return; + + assert(map->getNumInputs() == operands->size() && + "map inputs must match number of operands"); + + auto *context = map->getContext(); + SmallVector resultOperands; + resultOperands.reserve(operands->size()); + SmallVector remappedSymbols; + remappedSymbols.reserve(operands->size()); + unsigned nextDim = 0; + unsigned nextSym = 0; + unsigned oldNumSyms = map->getNumSymbols(); + SmallVector dimRemapping(map->getNumDims()); + for (unsigned i = 0, e = map->getNumInputs(); i != e; ++i) { + if (i < map->getNumDims()) { + if (isValidSymbol((*operands)[i])) { + // This is a valid symbols that appears as a dim, canonicalize it. + dimRemapping[i] = getAffineSymbolExpr(oldNumSyms + nextSym++, context); + remappedSymbols.push_back((*operands)[i]); + } else { + dimRemapping[i] = getAffineDimExpr(nextDim++, context); + resultOperands.push_back((*operands)[i]); + } + } else { + resultOperands.push_back((*operands)[i]); + } + } + + resultOperands.append(remappedSymbols.begin(), remappedSymbols.end()); + *operands = resultOperands; + *map = map->replaceDimsAndSymbols(dimRemapping, {}, nextDim, + oldNumSyms + nextSym); + + assert(map->getNumInputs() == operands->size() && + "map inputs must match number of operands"); +} + +void mlir::canonicalizeMapAndOperands( + AffineMap *map, llvm::SmallVectorImpl *operands) { + if (!map || operands->empty()) + return; + + assert(map->getNumInputs() == operands->size() && + "map inputs must match number of operands"); + + canonicalizePromotedSymbols(map, operands); + + // Check to see what dims are used. + llvm::SmallBitVector usedDims(map->getNumDims()); + llvm::SmallBitVector usedSyms(map->getNumSymbols()); + map->walkExprs([&](AffineExpr expr) { + if (auto dimExpr = expr.dyn_cast()) + usedDims[dimExpr.getPosition()] = true; + else if (auto symExpr = expr.dyn_cast()) + usedSyms[symExpr.getPosition()] = true; + }); + + auto *context = map->getContext(); + + SmallVector resultOperands; + resultOperands.reserve(operands->size()); + + llvm::SmallDenseMap seenDims; + SmallVector dimRemapping(map->getNumDims()); + unsigned nextDim = 0; + for (unsigned i = 0, e = map->getNumDims(); i != e; ++i) { + if (usedDims[i]) { + auto it = seenDims.find((*operands)[i]); + if (it == seenDims.end()) { + dimRemapping[i] = getAffineDimExpr(nextDim++, context); + resultOperands.push_back((*operands)[i]); + seenDims.insert(std::make_pair((*operands)[i], dimRemapping[i])); + } else { + dimRemapping[i] = it->second; + } + } + } + llvm::SmallDenseMap seenSymbols; + SmallVector symRemapping(map->getNumSymbols()); + unsigned nextSym = 0; + for (unsigned i = 0, e = map->getNumSymbols(); i != e; ++i) { + if (usedSyms[i]) { + auto it = seenSymbols.find((*operands)[i + map->getNumDims()]); + if (it == seenSymbols.end()) { + symRemapping[i] = getAffineSymbolExpr(nextSym++, context); + resultOperands.push_back((*operands)[i + map->getNumDims()]); + seenSymbols.insert(std::make_pair((*operands)[i + map->getNumDims()], + symRemapping[i])); + } else { + symRemapping[i] = it->second; + } + } + } + *map = + map->replaceDimsAndSymbols(dimRemapping, symRemapping, nextDim, nextSym); + *operands = resultOperands; +} + +namespace { +/// Simplify AffineApply operations. +/// +struct SimplifyAffineApply : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + PatternMatchResult matchAndRewrite(AffineApplyOp apply, + PatternRewriter &rewriter) const override { + auto map = apply.getAffineMap(); + + AffineMap oldMap = map; + SmallVector resultOperands(apply.getOperands()); + composeAffineMapAndOperands(&map, &resultOperands); + if (map == oldMap) + return matchFailure(); + + rewriter.replaceOpWithNewOp(apply, map, resultOperands); + return matchSuccess(); + } +}; +} // end anonymous namespace. + +void AffineApplyOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + results.insert(context); +} + +//===----------------------------------------------------------------------===// +// Common canonicalization pattern support logic +//===----------------------------------------------------------------------===// + +namespace { +/// This is a common class used for patterns of the form +/// "someop(memrefcast) -> someop". It folds the source of any memref_cast +/// into the root operation directly. +struct MemRefCastFolder : public RewritePattern { + /// The rootOpName is the name of the root operation to match against. + MemRefCastFolder(StringRef rootOpName, MLIRContext *context) + : RewritePattern(rootOpName, 1, context) {} + + PatternMatchResult match(Operation *op) const override { + for (auto *operand : op->getOperands()) + if (matchPattern(operand, m_Op())) + return matchSuccess(); + + return matchFailure(); + } + + void rewrite(Operation *op, PatternRewriter &rewriter) const override { + for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i) + if (auto *memref = op->getOperand(i)->getDefiningOp()) + if (auto cast = dyn_cast(memref)) + op->setOperand(i, cast.getOperand()); + rewriter.updatedRootInPlace(op); + } +}; + +} // end anonymous namespace. + +//===----------------------------------------------------------------------===// +// AffineDmaStartOp +//===----------------------------------------------------------------------===// + +// TODO(b/133776335) Check that map operands are loop IVs or symbols. +void AffineDmaStartOp::build(Builder *builder, OperationState *result, + Value *srcMemRef, AffineMap srcMap, + ArrayRef srcIndices, Value *destMemRef, + AffineMap dstMap, ArrayRef destIndices, + Value *tagMemRef, AffineMap tagMap, + ArrayRef tagIndices, Value *numElements, + Value *stride, Value *elementsPerStride) { + result->addOperands(srcMemRef); + result->addAttribute(getSrcMapAttrName(), builder->getAffineMapAttr(srcMap)); + result->addOperands(srcIndices); + result->addOperands(destMemRef); + result->addAttribute(getDstMapAttrName(), builder->getAffineMapAttr(dstMap)); + result->addOperands(destIndices); + result->addOperands(tagMemRef); + result->addAttribute(getTagMapAttrName(), builder->getAffineMapAttr(tagMap)); + result->addOperands(tagIndices); + result->addOperands(numElements); + if (stride) { + result->addOperands({stride, elementsPerStride}); + } +} + +void AffineDmaStartOp::print(OpAsmPrinter *p) { + *p << "affine.dma_start " << *getSrcMemRef() << '['; + SmallVector operands(getSrcIndices()); + p->printAffineMapOfSSAIds(getSrcMapAttr(), operands); + *p << "], " << *getDstMemRef() << '['; + operands.assign(getDstIndices().begin(), getDstIndices().end()); + p->printAffineMapOfSSAIds(getDstMapAttr(), operands); + *p << "], " << *getTagMemRef() << '['; + operands.assign(getTagIndices().begin(), getTagIndices().end()); + p->printAffineMapOfSSAIds(getTagMapAttr(), operands); + *p << "], " << *getNumElements(); + if (isStrided()) { + *p << ", " << *getStride(); + *p << ", " << *getNumElementsPerStride(); + } + *p << " : " << getSrcMemRefType() << ", " << getDstMemRefType() << ", " + << getTagMemRefType(); +} + +// 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) { + OpAsmParser::OperandType srcMemRefInfo; + AffineMapAttr srcMapAttr; + SmallVector srcMapOperands; + OpAsmParser::OperandType dstMemRefInfo; + AffineMapAttr dstMapAttr; + SmallVector dstMapOperands; + OpAsmParser::OperandType tagMemRefInfo; + AffineMapAttr tagMapAttr; + SmallVector tagMapOperands; + OpAsmParser::OperandType numElementsInfo; + SmallVector strideInfo; + + SmallVector types; + auto indexType = parser->getBuilder().getIndexType(); + + // Parse and resolve the following list of operands: + // *) dst memref followed by its affine maps operands (in square brackets). + // *) src memref followed by its affine map operands (in square brackets). + // *) tag memref followed by its affine map operands (in square brackets). + // *) number of elements transferred by DMA operation. + if (parser->parseOperand(srcMemRefInfo) || + parser->parseAffineMapOfSSAIds(srcMapOperands, srcMapAttr, + getSrcMapAttrName(), result->attributes) || + parser->parseComma() || parser->parseOperand(dstMemRefInfo) || + parser->parseAffineMapOfSSAIds(dstMapOperands, dstMapAttr, + getDstMapAttrName(), result->attributes) || + parser->parseComma() || parser->parseOperand(tagMemRefInfo) || + parser->parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr, + getTagMapAttrName(), result->attributes) || + parser->parseComma() || parser->parseOperand(numElementsInfo)) + return failure(); + + // Parse optional stride and elements per stride. + if (parser->parseTrailingOperandList(strideInfo)) { + return failure(); + } + if (!strideInfo.empty() && strideInfo.size() != 2) { + return parser->emitError(parser->getNameLoc(), + "expected two stride related operands"); + } + bool isStrided = strideInfo.size() == 2; + + if (parser->parseColonTypeList(types)) + return failure(); + + if (types.size() != 3) + return parser->emitError(parser->getNameLoc(), "expected three types"); + + if (parser->resolveOperand(srcMemRefInfo, types[0], result->operands) || + parser->resolveOperands(srcMapOperands, indexType, result->operands) || + parser->resolveOperand(dstMemRefInfo, types[1], result->operands) || + parser->resolveOperands(dstMapOperands, indexType, result->operands) || + parser->resolveOperand(tagMemRefInfo, types[2], result->operands) || + parser->resolveOperands(tagMapOperands, indexType, result->operands) || + parser->resolveOperand(numElementsInfo, indexType, result->operands)) + return failure(); + + if (isStrided) { + if (parser->resolveOperands(strideInfo, indexType, result->operands)) + return failure(); + } + + // Check that src/dst/tag operand counts match their map.numInputs. + if (srcMapOperands.size() != srcMapAttr.getValue().getNumInputs() || + dstMapOperands.size() != dstMapAttr.getValue().getNumInputs() || + tagMapOperands.size() != tagMapAttr.getValue().getNumInputs()) + return parser->emitError(parser->getNameLoc(), + "memref operand count not equal to map.numInputs"); + return success(); +} + +LogicalResult AffineDmaStartOp::verify() { + if (!getOperand(getSrcMemRefOperandIndex())->getType().isa()) + return emitOpError("expected DMA source to be of memref type"); + if (!getOperand(getDstMemRefOperandIndex())->getType().isa()) + return emitOpError("expected DMA destination to be of memref type"); + if (!getOperand(getTagMemRefOperandIndex())->getType().isa()) + return emitOpError("expected DMA tag to be of memref type"); + + // DMAs from different memory spaces supported. + if (getSrcMemorySpace() == getDstMemorySpace()) { + return emitOpError("DMA should be between different memory spaces"); + } + unsigned numInputsAllMaps = getSrcMap().getNumInputs() + + getDstMap().getNumInputs() + + getTagMap().getNumInputs(); + if (getNumOperands() != numInputsAllMaps + 3 + 1 && + getNumOperands() != numInputsAllMaps + 3 + 1 + 2) { + return emitOpError("incorrect number of operands"); + } + + for (auto *idx : getSrcIndices()) { + if (!idx->getType().isIndex()) + return emitOpError("src index to dma_start must have 'index' type"); + if (!isValidAffineIndexOperand(idx)) + return emitOpError("src index must be a dimension or symbol identifier"); + } + for (auto *idx : getDstIndices()) { + if (!idx->getType().isIndex()) + return emitOpError("dst index to dma_start must have 'index' type"); + if (!isValidAffineIndexOperand(idx)) + return emitOpError("dst index must be a dimension or symbol identifier"); + } + for (auto *idx : getTagIndices()) { + if (!idx->getType().isIndex()) + return emitOpError("tag index to dma_start must have 'index' type"); + if (!isValidAffineIndexOperand(idx)) + return emitOpError("tag index must be a dimension or symbol identifier"); + } + return success(); +} + +void AffineDmaStartOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + /// dma_start(memrefcast) -> dma_start + results.insert(getOperationName(), context); +} + +//===----------------------------------------------------------------------===// +// AffineDmaWaitOp +//===----------------------------------------------------------------------===// + +// TODO(b/133776335) Check that map operands are loop IVs or symbols. +void AffineDmaWaitOp::build(Builder *builder, OperationState *result, + Value *tagMemRef, AffineMap tagMap, + ArrayRef tagIndices, Value *numElements) { + result->addOperands(tagMemRef); + result->addAttribute(getTagMapAttrName(), builder->getAffineMapAttr(tagMap)); + result->addOperands(tagIndices); + result->addOperands(numElements); +} + +void AffineDmaWaitOp::print(OpAsmPrinter *p) { + *p << "affine.dma_wait " << *getTagMemRef() << '['; + SmallVector operands(getTagIndices()); + p->printAffineMapOfSSAIds(getTagMapAttr(), operands); + *p << "], "; + p->printOperand(getNumElements()); + *p << " : " << getTagMemRef()->getType(); +} + +// Parse AffineDmaWaitOp. +// Eg: +// affine.dma_wait %tag[%index], %num_elements +// : memref<1 x i32, (d0) -> (d0), 4> +// +ParseResult AffineDmaWaitOp::parse(OpAsmParser *parser, + OperationState *result) { + OpAsmParser::OperandType tagMemRefInfo; + AffineMapAttr tagMapAttr; + SmallVector tagMapOperands; + Type type; + auto indexType = parser->getBuilder().getIndexType(); + OpAsmParser::OperandType numElementsInfo; + + // Parse tag memref, its map operands, and dma size. + if (parser->parseOperand(tagMemRefInfo) || + parser->parseAffineMapOfSSAIds(tagMapOperands, tagMapAttr, + getTagMapAttrName(), result->attributes) || + parser->parseComma() || parser->parseOperand(numElementsInfo) || + parser->parseColonType(type) || + parser->resolveOperand(tagMemRefInfo, type, result->operands) || + parser->resolveOperands(tagMapOperands, indexType, result->operands) || + parser->resolveOperand(numElementsInfo, indexType, result->operands)) + return failure(); + + if (!type.isa()) + return parser->emitError(parser->getNameLoc(), + "expected tag to be of memref type"); + + if (tagMapOperands.size() != tagMapAttr.getValue().getNumInputs()) + return parser->emitError(parser->getNameLoc(), + "tag memref operand count != to map.numInputs"); + return success(); +} + +LogicalResult AffineDmaWaitOp::verify() { + if (!getOperand(0)->getType().isa()) + return emitOpError("expected DMA tag to be of memref type"); + for (auto *idx : getTagIndices()) { + if (!idx->getType().isIndex()) + return emitOpError("index to dma_wait must have 'index' type"); + if (!isValidAffineIndexOperand(idx)) + return emitOpError("index must be a dimension or symbol identifier"); + } + return success(); +} + +void AffineDmaWaitOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + /// dma_wait(memrefcast) -> dma_wait + results.insert(getOperationName(), context); +} + +//===----------------------------------------------------------------------===// +// AffineForOp +//===----------------------------------------------------------------------===// + +void AffineForOp::build(Builder *builder, OperationState *result, + ArrayRef lbOperands, AffineMap lbMap, + ArrayRef ubOperands, AffineMap ubMap, + int64_t step) { + assert(((!lbMap && lbOperands.empty()) || + lbOperands.size() == lbMap.getNumInputs()) && + "lower bound operand count does not match the affine map"); + assert(((!ubMap && ubOperands.empty()) || + ubOperands.size() == ubMap.getNumInputs()) && + "upper bound operand count does not match the affine map"); + assert(step > 0 && "step has to be a positive integer constant"); + + // Add an attribute for the step. + result->addAttribute(getStepAttrName(), + builder->getIntegerAttr(builder->getIndexType(), step)); + + // Add the lower bound. + result->addAttribute(getLowerBoundAttrName(), + builder->getAffineMapAttr(lbMap)); + result->addOperands(lbOperands); + + // Add the upper bound. + result->addAttribute(getUpperBoundAttrName(), + builder->getAffineMapAttr(ubMap)); + result->addOperands(ubOperands); + + // Create a region and a block for the body. The argument of the region is + // the loop induction variable. + Region *bodyRegion = result->addRegion(); + Block *body = new Block(); + body->addArgument(IndexType::get(builder->getContext())); + bodyRegion->push_back(body); + ensureTerminator(*bodyRegion, *builder, result->location); + + // Set the operands list as resizable so that we can freely modify the bounds. + result->setOperandListToResizable(); +} + +void AffineForOp::build(Builder *builder, OperationState *result, int64_t lb, + int64_t ub, int64_t step) { + auto lbMap = AffineMap::getConstantMap(lb, builder->getContext()); + auto ubMap = AffineMap::getConstantMap(ub, builder->getContext()); + return build(builder, result, {}, lbMap, {}, ubMap, step); +} + +static LogicalResult verify(AffineForOp op) { + // Check that the body defines as single block argument for the induction + // variable. + auto *body = op.getBody(); + if (body->getNumArguments() != 1 || + !body->getArgument(0)->getType().isIndex()) + return op.emitOpError( + "expected body to have a single index argument for the " + "induction variable"); + + // Verify that there are enough operands for the bounds. + AffineMap lowerBoundMap = op.getLowerBoundMap(), + upperBoundMap = op.getUpperBoundMap(); + if (op.getNumOperands() != + (lowerBoundMap.getNumInputs() + upperBoundMap.getNumInputs())) + return op.emitOpError( + "operand count must match with affine map dimension and symbol count"); + + // Verify that the bound operands are valid dimension/symbols. + /// Lower bound. + if (failed(verifyDimAndSymbolIdentifiers(op, op.getLowerBoundOperands(), + op.getLowerBoundMap().getNumDims()))) + return failure(); + /// Upper bound. + if (failed(verifyDimAndSymbolIdentifiers(op, op.getUpperBoundOperands(), + op.getUpperBoundMap().getNumDims()))) + return failure(); + return success(); +} + +/// Parse a for operation loop bounds. +static ParseResult parseBound(bool isLower, OperationState *result, + OpAsmParser *p) { + // 'min' / 'max' prefixes are generally syntactic sugar, but are required if + // the map has multiple results. + bool failedToParsedMinMax = + failed(p->parseOptionalKeyword(isLower ? "max" : "min")); + + auto &builder = p->getBuilder(); + auto boundAttrName = isLower ? AffineForOp::getLowerBoundAttrName() + : AffineForOp::getUpperBoundAttrName(); + + // Parse ssa-id as identity map. + SmallVector boundOpInfos; + if (p->parseOperandList(boundOpInfos)) + return failure(); + + if (!boundOpInfos.empty()) { + // Check that only one operand was parsed. + if (boundOpInfos.size() > 1) + return p->emitError(p->getNameLoc(), + "expected only one loop bound operand"); + + // TODO: improve error message when SSA value is not an affine integer. + // Currently it is 'use of value ... expects different type than prior uses' + if (p->resolveOperand(boundOpInfos.front(), builder.getIndexType(), + result->operands)) + return failure(); + + // Create an identity map using symbol id. This representation is optimized + // for storage. Analysis passes may expand it into a multi-dimensional map + // if desired. + AffineMap map = builder.getSymbolIdentityMap(); + result->addAttribute(boundAttrName, builder.getAffineMapAttr(map)); + return success(); + } + + // Get the attribute location. + llvm::SMLoc attrLoc = p->getCurrentLocation(); + + Attribute boundAttr; + if (p->parseAttribute(boundAttr, builder.getIndexType(), boundAttrName, + result->attributes)) + return failure(); + + // Parse full form - affine map followed by dim and symbol list. + if (auto affineMapAttr = boundAttr.dyn_cast()) { + unsigned currentNumOperands = result->operands.size(); + unsigned numDims; + if (parseDimAndSymbolList(p, result->operands, numDims)) + return failure(); + + auto map = affineMapAttr.getValue(); + if (map.getNumDims() != numDims) + return p->emitError( + p->getNameLoc(), + "dim operand count and integer set dim count must match"); + + unsigned numDimAndSymbolOperands = + result->operands.size() - currentNumOperands; + if (numDims + map.getNumSymbols() != numDimAndSymbolOperands) + return p->emitError( + p->getNameLoc(), + "symbol operand count and integer set symbol count must match"); + + // If the map has multiple results, make sure that we parsed the min/max + // prefix. + if (map.getNumResults() > 1 && failedToParsedMinMax) { + if (isLower) { + return p->emitError(attrLoc, "lower loop bound affine map with " + "multiple results requires 'max' prefix"); + } + return p->emitError(attrLoc, "upper loop bound affine map with multiple " + "results requires 'min' prefix"); + } + return success(); + } + + // Parse custom assembly form. + if (auto integerAttr = boundAttr.dyn_cast()) { + result->attributes.pop_back(); + result->addAttribute( + boundAttrName, builder.getAffineMapAttr( + builder.getConstantAffineMap(integerAttr.getInt()))); + return success(); + } + + return p->emitError( + p->getNameLoc(), + "expected valid affine map representation for loop bounds"); +} + +ParseResult parseAffineForOp(OpAsmParser *parser, OperationState *result) { + auto &builder = parser->getBuilder(); + OpAsmParser::OperandType inductionVariable; + // Parse the induction variable followed by '='. + if (parser->parseRegionArgument(inductionVariable) || parser->parseEqual()) + return failure(); + + // Parse loop bounds. + if (parseBound(/*isLower=*/true, result, parser) || + parser->parseKeyword("to", " between bounds") || + parseBound(/*isLower=*/false, result, parser)) + return failure(); + + // Parse the optional loop step, we default to 1 if one is not present. + if (parser->parseOptionalKeyword("step")) { + result->addAttribute( + AffineForOp::getStepAttrName(), + builder.getIntegerAttr(builder.getIndexType(), /*value=*/1)); + } else { + llvm::SMLoc stepLoc = parser->getCurrentLocation(); + IntegerAttr stepAttr; + if (parser->parseAttribute(stepAttr, builder.getIndexType(), + AffineForOp::getStepAttrName().data(), + result->attributes)) + return failure(); + + if (stepAttr.getValue().getSExtValue() < 0) + return parser->emitError( + stepLoc, + "expected step to be representable as a positive signed integer"); + } + + // Parse the body region. + Region *body = result->addRegion(); + if (parser->parseRegion(*body, inductionVariable, builder.getIndexType())) + return failure(); + + AffineForOp::ensureTerminator(*body, builder, result->location); + + // Parse the optional attribute list. + if (parser->parseOptionalAttributeDict(result->attributes)) + return failure(); + + // Set the operands list as resizable so that we can freely modify the bounds. + result->setOperandListToResizable(); + return success(); +} + +static void printBound(AffineMapAttr boundMap, + Operation::operand_range boundOperands, + const char *prefix, OpAsmPrinter *p) { + AffineMap map = boundMap.getValue(); + + // Check if this bound should be printed using custom assembly form. + // The decision to restrict printing custom assembly form to trivial cases + // comes from the will to roundtrip MLIR binary -> text -> binary in a + // lossless way. + // Therefore, custom assembly form parsing and printing is only supported for + // zero-operand constant maps and single symbol operand identity maps. + if (map.getNumResults() == 1) { + AffineExpr expr = map.getResult(0); + + // Print constant bound. + if (map.getNumDims() == 0 && map.getNumSymbols() == 0) { + if (auto constExpr = expr.dyn_cast()) { + *p << constExpr.getValue(); + return; + } + } + + // Print bound that consists of a single SSA symbol if the map is over a + // single symbol. + if (map.getNumDims() == 0 && map.getNumSymbols() == 1) { + if (auto symExpr = expr.dyn_cast()) { + p->printOperand(*boundOperands.begin()); + return; + } + } + } else { + // Map has multiple results. Print 'min' or 'max' prefix. + *p << prefix << ' '; + } + + // Print the map and its operands. + *p << boundMap; + printDimAndSymbolList(boundOperands.begin(), boundOperands.end(), + map.getNumDims(), p); +} + +void print(OpAsmPrinter *p, AffineForOp op) { + *p << "affine.for "; + p->printOperand(op.getBody()->getArgument(0)); + *p << " = "; + printBound(op.getLowerBoundMapAttr(), op.getLowerBoundOperands(), "max", p); + *p << " to "; + printBound(op.getUpperBoundMapAttr(), op.getUpperBoundOperands(), "min", p); + + if (op.getStep() != 1) + *p << " step " << op.getStep(); + p->printRegion(op.region(), + /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); + p->printOptionalAttrDict(op.getAttrs(), + /*elidedAttrs=*/{op.getLowerBoundAttrName(), + op.getUpperBoundAttrName(), + op.getStepAttrName()}); +} + +namespace { +/// This is a pattern to fold constant loop bounds. +struct AffineForLoopBoundFolder : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + PatternMatchResult matchAndRewrite(AffineForOp forOp, + PatternRewriter &rewriter) const override { + auto foldLowerOrUpperBound = [&forOp](bool lower) { + // Check to see if each of the operands is the result of a constant. If + // so, get the value. If not, ignore it. + SmallVector operandConstants; + auto boundOperands = + lower ? forOp.getLowerBoundOperands() : forOp.getUpperBoundOperands(); + for (auto *operand : boundOperands) { + Attribute operandCst; + matchPattern(operand, m_Constant(&operandCst)); + operandConstants.push_back(operandCst); + } + + AffineMap boundMap = + lower ? forOp.getLowerBoundMap() : forOp.getUpperBoundMap(); + assert(boundMap.getNumResults() >= 1 && + "bound maps should have at least one result"); + SmallVector foldedResults; + if (failed(boundMap.constantFold(operandConstants, foldedResults))) + return failure(); + + // Compute the max or min as applicable over the results. + assert(!foldedResults.empty() && + "bounds should have at least one result"); + auto maxOrMin = foldedResults[0].cast().getValue(); + for (unsigned i = 1, e = foldedResults.size(); i < e; i++) { + auto foldedResult = foldedResults[i].cast().getValue(); + maxOrMin = lower ? llvm::APIntOps::smax(maxOrMin, foldedResult) + : llvm::APIntOps::smin(maxOrMin, foldedResult); + } + lower ? forOp.setConstantLowerBound(maxOrMin.getSExtValue()) + : forOp.setConstantUpperBound(maxOrMin.getSExtValue()); + return success(); + }; + + // Try to fold the lower bound. + bool folded = false; + if (!forOp.hasConstantLowerBound()) + folded |= succeeded(foldLowerOrUpperBound(/*lower=*/true)); + + // Try to fold the upper bound. + if (!forOp.hasConstantUpperBound()) + folded |= succeeded(foldLowerOrUpperBound(/*lower=*/false)); + + // If any of the bounds were folded we return success. + if (!folded) + return matchFailure(); + rewriter.updatedRootInPlace(forOp); + return matchSuccess(); + } +}; +} // end anonymous namespace + +void AffineForOp::getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context) { + results.insert(context); +} + +AffineBound AffineForOp::getLowerBound() { + auto lbMap = getLowerBoundMap(); + return AffineBound(AffineForOp(*this), 0, lbMap.getNumInputs(), lbMap); +} + +AffineBound AffineForOp::getUpperBound() { + auto lbMap = getLowerBoundMap(); + auto ubMap = getUpperBoundMap(); + return AffineBound(AffineForOp(*this), lbMap.getNumInputs(), getNumOperands(), + ubMap); +} + +void AffineForOp::setLowerBound(ArrayRef lbOperands, AffineMap map) { + assert(lbOperands.size() == map.getNumInputs()); + assert(map.getNumResults() >= 1 && "bound map has at least one result"); + + SmallVector newOperands(lbOperands.begin(), lbOperands.end()); + + auto ubOperands = getUpperBoundOperands(); + newOperands.append(ubOperands.begin(), ubOperands.end()); + getOperation()->setOperands(newOperands); + + setAttr(getLowerBoundAttrName(), AffineMapAttr::get(map)); +} + +void AffineForOp::setUpperBound(ArrayRef ubOperands, AffineMap map) { + assert(ubOperands.size() == map.getNumInputs()); + assert(map.getNumResults() >= 1 && "bound map has at least one result"); + + SmallVector newOperands(getLowerBoundOperands()); + newOperands.append(ubOperands.begin(), ubOperands.end()); + getOperation()->setOperands(newOperands); + + setAttr(getUpperBoundAttrName(), AffineMapAttr::get(map)); +} + +void AffineForOp::setLowerBoundMap(AffineMap map) { + auto lbMap = getLowerBoundMap(); + assert(lbMap.getNumDims() == map.getNumDims() && + lbMap.getNumSymbols() == map.getNumSymbols()); + assert(map.getNumResults() >= 1 && "bound map has at least one result"); + (void)lbMap; + setAttr(getLowerBoundAttrName(), AffineMapAttr::get(map)); +} + +void AffineForOp::setUpperBoundMap(AffineMap map) { + auto ubMap = getUpperBoundMap(); + assert(ubMap.getNumDims() == map.getNumDims() && + ubMap.getNumSymbols() == map.getNumSymbols()); + assert(map.getNumResults() >= 1 && "bound map has at least one result"); + (void)ubMap; + setAttr(getUpperBoundAttrName(), AffineMapAttr::get(map)); +} + +bool AffineForOp::hasConstantLowerBound() { + return getLowerBoundMap().isSingleConstant(); +} + +bool AffineForOp::hasConstantUpperBound() { + return getUpperBoundMap().isSingleConstant(); +} + +int64_t AffineForOp::getConstantLowerBound() { + return getLowerBoundMap().getSingleConstantResult(); +} + +int64_t AffineForOp::getConstantUpperBound() { + return getUpperBoundMap().getSingleConstantResult(); +} + +void AffineForOp::setConstantLowerBound(int64_t value) { + setLowerBound({}, AffineMap::getConstantMap(value, getContext())); +} + +void AffineForOp::setConstantUpperBound(int64_t value) { + setUpperBound({}, AffineMap::getConstantMap(value, getContext())); +} + +AffineForOp::operand_range AffineForOp::getLowerBoundOperands() { + return {operand_begin(), operand_begin() + getLowerBoundMap().getNumInputs()}; +} + +AffineForOp::operand_range AffineForOp::getUpperBoundOperands() { + return {operand_begin() + getLowerBoundMap().getNumInputs(), operand_end()}; +} + +bool AffineForOp::matchingBoundOperandList() { + auto lbMap = getLowerBoundMap(); + auto ubMap = getUpperBoundMap(); + if (lbMap.getNumDims() != ubMap.getNumDims() || + lbMap.getNumSymbols() != ubMap.getNumSymbols()) + return false; + + unsigned numOperands = lbMap.getNumInputs(); + for (unsigned i = 0, e = lbMap.getNumInputs(); i < e; i++) { + // Compare Value *'s. + if (getOperand(i) != getOperand(numOperands + i)) + return false; + } + return true; +} + +/// Returns if the provided value is the induction variable of a AffineForOp. +bool mlir::isForInductionVar(Value *val) { + return getForInductionVarOwner(val) != AffineForOp(); +} + +/// Returns the loop parent of an induction variable. If the provided value is +/// not an induction variable, then return nullptr. +AffineForOp mlir::getForInductionVarOwner(Value *val) { + auto *ivArg = dyn_cast(val); + if (!ivArg || !ivArg->getOwner()) + return AffineForOp(); + auto *containingInst = ivArg->getOwner()->getParent()->getParentOp(); + return dyn_cast(containingInst); +} + +/// Extracts the induction variables from a list of AffineForOps and returns +/// them. +void mlir::extractForInductionVars(ArrayRef forInsts, + SmallVectorImpl *ivs) { + ivs->reserve(forInsts.size()); + for (auto forInst : forInsts) + ivs->push_back(forInst.getInductionVar()); +} + +//===----------------------------------------------------------------------===// +// AffineIfOp +//===----------------------------------------------------------------------===// + +static LogicalResult verify(AffineIfOp op) { + // Verify that we have a condition attribute. + auto conditionAttr = + op.getAttrOfType(op.getConditionAttrName()); + if (!conditionAttr) + return op.emitOpError( + "requires an integer set attribute named 'condition'"); + + // Verify that there are enough operands for the condition. + IntegerSet condition = conditionAttr.getValue(); + if (op.getNumOperands() != condition.getNumOperands()) + return op.emitOpError( + "operand count and condition integer set dimension and " + "symbol count must match"); + + // Verify that the operands are valid dimension/symbols. + if (failed(verifyDimAndSymbolIdentifiers( + op, op.getOperation()->getNonSuccessorOperands(), + condition.getNumDims()))) + return failure(); + + // Verify that the entry of each child region does not have arguments. + for (auto ®ion : op.getOperation()->getRegions()) { + for (auto &b : region) + if (b.getNumArguments() != 0) + return op.emitOpError( + "requires that child entry blocks have no arguments"); + } + return success(); +} + +ParseResult parseAffineIfOp(OpAsmParser *parser, OperationState *result) { + // Parse the condition attribute set. + IntegerSetAttr conditionAttr; + unsigned numDims; + if (parser->parseAttribute(conditionAttr, AffineIfOp::getConditionAttrName(), + result->attributes) || + parseDimAndSymbolList(parser, result->operands, numDims)) + return failure(); + + // Verify the condition operands. + auto set = conditionAttr.getValue(); + if (set.getNumDims() != numDims) + return parser->emitError( + parser->getNameLoc(), + "dim operand count and integer set dim count must match"); + if (numDims + set.getNumSymbols() != result->operands.size()) + return parser->emitError( + parser->getNameLoc(), + "symbol operand count and integer set symbol count must match"); + + // Create the regions for 'then' and 'else'. The latter must be created even + // if it remains empty for the validity of the operation. + result->regions.reserve(2); + Region *thenRegion = result->addRegion(); + Region *elseRegion = result->addRegion(); + + // Parse the 'then' region. + if (parser->parseRegion(*thenRegion, {}, {})) + return failure(); + AffineIfOp::ensureTerminator(*thenRegion, parser->getBuilder(), + result->location); + + // If we find an 'else' keyword then parse the 'else' region. + if (!parser->parseOptionalKeyword("else")) { + if (parser->parseRegion(*elseRegion, {}, {})) + return failure(); + AffineIfOp::ensureTerminator(*elseRegion, parser->getBuilder(), + result->location); + } + + // Parse the optional attribute list. + if (parser->parseOptionalAttributeDict(result->attributes)) + return failure(); + + return success(); +} + +void print(OpAsmPrinter *p, AffineIfOp op) { + auto conditionAttr = + op.getAttrOfType(op.getConditionAttrName()); + *p << "affine.if " << conditionAttr; + printDimAndSymbolList(op.operand_begin(), op.operand_end(), + conditionAttr.getValue().getNumDims(), p); + p->printRegion(op.thenRegion(), + /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); + + // Print the 'else' regions if it has any blocks. + auto &elseRegion = op.elseRegion(); + if (!elseRegion.empty()) { + *p << " else"; + p->printRegion(elseRegion, + /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); + } + + // Print the attribute list. + p->printOptionalAttrDict(op.getAttrs(), + /*elidedAttrs=*/op.getConditionAttrName()); +} + +IntegerSet AffineIfOp::getIntegerSet() { + return getAttrOfType(getConditionAttrName()).getValue(); +} +void AffineIfOp::setIntegerSet(IntegerSet newSet) { + setAttr(getConditionAttrName(), IntegerSetAttr::get(newSet)); +} + +//===----------------------------------------------------------------------===// +// AffineLoadOp +//===----------------------------------------------------------------------===// + +void AffineLoadOp::build(Builder *builder, OperationState *result, + AffineMap map, ArrayRef operands) { + result->addOperands(operands); + if (map) + result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); + auto memrefType = operands[0]->getType().cast(); + result->types.push_back(memrefType.getElementType()); +} + +void AffineLoadOp::build(Builder *builder, OperationState *result, + Value *memref, ArrayRef indices) { + result->addOperands(memref); + result->addOperands(indices); + auto memrefType = memref->getType().cast(); + auto rank = memrefType.getRank(); + // Create identity map for memrefs with at least one dimension or () -> () + // for zero-dimensional memrefs. + auto map = rank ? builder->getMultiDimIdentityMap(rank) + : builder->getEmptyAffineMap(); + result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); + result->types.push_back(memrefType.getElementType()); +} + +ParseResult AffineLoadOp::parse(OpAsmParser *parser, OperationState *result) { + auto &builder = parser->getBuilder(); + auto affineIntTy = builder.getIndexType(); + + MemRefType type; + OpAsmParser::OperandType memrefInfo; + AffineMapAttr mapAttr; + SmallVector mapOperands; + return failure( + parser->parseOperand(memrefInfo) || + parser->parseAffineMapOfSSAIds(mapOperands, mapAttr, getMapAttrName(), + result->attributes) || + parser->parseOptionalAttributeDict(result->attributes) || + parser->parseColonType(type) || + parser->resolveOperand(memrefInfo, type, result->operands) || + parser->resolveOperands(mapOperands, affineIntTy, result->operands) || + parser->addTypeToList(type.getElementType(), result->types)); +} + +void AffineLoadOp::print(OpAsmPrinter *p) { + *p << "affine.load " << *getMemRef() << '['; + AffineMapAttr mapAttr = getAttrOfType(getMapAttrName()); + if (mapAttr) { + SmallVector operands(getIndices()); + p->printAffineMapOfSSAIds(mapAttr, operands); + } + *p << ']'; + p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{getMapAttrName()}); + *p << " : " << getMemRefType(); +} + +LogicalResult AffineLoadOp::verify() { + if (getType() != getMemRefType().getElementType()) + return emitOpError("result type must match element type of memref"); + + auto mapAttr = getAttrOfType(getMapAttrName()); + if (mapAttr) { + AffineMap map = getAttrOfType(getMapAttrName()).getValue(); + if (map.getNumResults() != getMemRefType().getRank()) + return emitOpError("affine.load affine map num results must equal" + " memref rank"); + if (map.getNumInputs() != getNumOperands() - 1) + return emitOpError("expects as many subscripts as affine map inputs"); + } else { + if (getMemRefType().getRank() != getNumOperands() - 1) + return emitOpError( + "expects the number of subscripts to be equal to memref rank"); + } + + for (auto *idx : getIndices()) { + if (!idx->getType().isIndex()) + return emitOpError("index to load must have 'index' type"); + if (!isValidAffineIndexOperand(idx)) + return emitOpError("index must be a dimension or symbol identifier"); + } + return success(); +} + +void AffineLoadOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + /// load(memrefcast) -> load + results.insert(getOperationName(), context); +} + +//===----------------------------------------------------------------------===// +// AffineStoreOp +//===----------------------------------------------------------------------===// + +void AffineStoreOp::build(Builder *builder, OperationState *result, + Value *valueToStore, AffineMap map, + ArrayRef operands) { + result->addOperands(valueToStore); + result->addOperands(operands); + if (map) + result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); +} + +void AffineStoreOp::build(Builder *builder, OperationState *result, + Value *valueToStore, Value *memref, + ArrayRef operands) { + result->addOperands(valueToStore); + result->addOperands(memref); + result->addOperands(operands); + auto memrefType = memref->getType().cast(); + auto rank = memrefType.getRank(); + // Create identity map for memrefs with at least one dimension or () -> () + // for zero-dimensional memrefs. + auto map = rank ? builder->getMultiDimIdentityMap(rank) + : builder->getEmptyAffineMap(); + result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map)); +} + +ParseResult AffineStoreOp::parse(OpAsmParser *parser, OperationState *result) { + auto affineIntTy = parser->getBuilder().getIndexType(); + + MemRefType type; + OpAsmParser::OperandType storeValueInfo; + OpAsmParser::OperandType memrefInfo; + AffineMapAttr mapAttr; + SmallVector mapOperands; + return failure( + parser->parseOperand(storeValueInfo) || parser->parseComma() || + parser->parseOperand(memrefInfo) || + parser->parseAffineMapOfSSAIds(mapOperands, mapAttr, getMapAttrName(), + result->attributes) || + parser->parseOptionalAttributeDict(result->attributes) || + parser->parseColonType(type) || + parser->resolveOperand(storeValueInfo, type.getElementType(), + result->operands) || + parser->resolveOperand(memrefInfo, type, result->operands) || + parser->resolveOperands(mapOperands, affineIntTy, result->operands)); +} + +void AffineStoreOp::print(OpAsmPrinter *p) { + *p << "affine.store " << *getValueToStore(); + *p << ", " << *getMemRef() << '['; + AffineMapAttr mapAttr = getAttrOfType(getMapAttrName()); + if (mapAttr) { + SmallVector operands(getIndices()); + p->printAffineMapOfSSAIds(mapAttr, operands); + } + *p << ']'; + p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{getMapAttrName()}); + *p << " : " << getMemRefType(); +} + +LogicalResult AffineStoreOp::verify() { + // First operand must have same type as memref element type. + if (getValueToStore()->getType() != getMemRefType().getElementType()) + return emitOpError("first operand must have same type memref element type"); + + auto mapAttr = getAttrOfType(getMapAttrName()); + if (mapAttr) { + AffineMap map = mapAttr.getValue(); + if (map.getNumResults() != getMemRefType().getRank()) + return emitOpError("affine.store affine map num results must equal" + " memref rank"); + if (map.getNumInputs() != getNumOperands() - 2) + return emitOpError("expects as many subscripts as affine map inputs"); + } else { + if (getMemRefType().getRank() != getNumOperands() - 2) + return emitOpError( + "expects the number of subscripts to be equal to memref rank"); + } + + for (auto *idx : getIndices()) { + if (!idx->getType().isIndex()) + return emitOpError("index to store must have 'index' type"); + if (!isValidAffineIndexOperand(idx)) + return emitOpError("index must be a dimension or symbol identifier"); + } + return success(); +} + +void AffineStoreOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + /// load(memrefcast) -> load + results.insert(getOperationName(), context); +} + +#define GET_OP_CLASSES +#include "mlir/Dialect/AffineOps/AffineOps.cpp.inc" diff --git a/mlir/lib/Dialect/AffineOps/CMakeLists.txt b/mlir/lib/Dialect/AffineOps/CMakeLists.txt new file mode 100644 index 00000000000..dbe469369a3 --- /dev/null +++ b/mlir/lib/Dialect/AffineOps/CMakeLists.txt @@ -0,0 +1,10 @@ +add_llvm_library(MLIRAffineOps + AffineOps.cpp + DialectRegistration.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/AffineOps + ) +add_dependencies(MLIRAffineOps MLIRAffineOpsIncGen MLIRIR MLIRStandardOps) +target_link_libraries(MLIRAffineOps MLIRIR MLIRStandardOps) + diff --git a/mlir/lib/Dialect/AffineOps/DialectRegistration.cpp b/mlir/lib/Dialect/AffineOps/DialectRegistration.cpp new file mode 100644 index 00000000000..9197e3c619f --- /dev/null +++ b/mlir/lib/Dialect/AffineOps/DialectRegistration.cpp @@ -0,0 +1,22 @@ +//===- DialectRegistration.cpp - Register Affine Op dialect ---------------===// +// +// Copyright 2019 The MLIR Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================= + +#include "mlir/Dialect/AffineOps/AffineOps.h" +using namespace mlir; + +// Static initialization for Affine op dialect registration. +static DialectRegistration StandardOps; diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt index 294041df4a5..b0641a9611f 100644 --- a/mlir/lib/Dialect/CMakeLists.txt +++ b/mlir/lib/Dialect/CMakeLists.txt @@ -1,3 +1,4 @@ +add_subdirectory(AffineOps) add_subdirectory(FxpMathOps) add_subdirectory(GPU) add_subdirectory(Linalg) diff --git a/mlir/lib/Dialect/Linalg/Transforms/LowerToLoops.cpp b/mlir/lib/Dialect/Linalg/Transforms/LowerToLoops.cpp index 1c5bb6e70c8..c48437f60db 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/LowerToLoops.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/LowerToLoops.cpp @@ -15,7 +15,12 @@ // limitations under the License. // ============================================================================= -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" +#include "mlir/Dialect/Linalg/IR/LinalgOps.h" +#include "mlir/Dialect/Linalg/IR/LinalgTypes.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/Linalg/Utils/Intrinsics.h" +#include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/LoopOps/LoopOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/EDSC/Helpers.h" @@ -23,11 +28,6 @@ #include "mlir/IR/AffineMap.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/OpImplementation.h" -#include "mlir/Dialect/Linalg/IR/LinalgOps.h" -#include "mlir/Dialect/Linalg/IR/LinalgTypes.h" -#include "mlir/Dialect/Linalg/Passes.h" -#include "mlir/Dialect/Linalg/Utils/Intrinsics.h" -#include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Pass/Pass.h" #include "mlir/Support/LLVM.h" #include "mlir/Support/STLExtras.h" diff --git a/mlir/lib/Transforms/AffineDataCopyGeneration.cpp b/mlir/lib/Transforms/AffineDataCopyGeneration.cpp index 33b73336ff7..71f6c78462d 100644 --- a/mlir/lib/Transforms/AffineDataCopyGeneration.cpp +++ b/mlir/lib/Transforms/AffineDataCopyGeneration.cpp @@ -28,9 +28,9 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp index 98798938077..46713dcff49 100644 --- a/mlir/lib/Transforms/LoopFusion.cpp +++ b/mlir/lib/Transforms/LoopFusion.cpp @@ -19,11 +19,11 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" diff --git a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp index 094f8fc421d..293e565cda7 100644 --- a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp +++ b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp @@ -19,12 +19,12 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" diff --git a/mlir/lib/Transforms/LoopTiling.cpp b/mlir/lib/Transforms/LoopTiling.cpp index c521a8f6f5d..02787b12e3d 100644 --- a/mlir/lib/Transforms/LoopTiling.cpp +++ b/mlir/lib/Transforms/LoopTiling.cpp @@ -19,11 +19,11 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/IR/Builders.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/LoopUtils.h" diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp index fbe1dcc09f9..2acc5a90f5f 100644 --- a/mlir/lib/Transforms/LoopUnroll.cpp +++ b/mlir/lib/Transforms/LoopUnroll.cpp @@ -21,8 +21,8 @@ #include "mlir/Transforms/Passes.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/LoopAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/Builders.h" diff --git a/mlir/lib/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Transforms/LoopUnrollAndJam.cpp index ef92861adf9..3e92ad739e8 100644 --- a/mlir/lib/Transforms/LoopUnrollAndJam.cpp +++ b/mlir/lib/Transforms/LoopUnrollAndJam.cpp @@ -43,8 +43,8 @@ //===----------------------------------------------------------------------===// #include "mlir/Transforms/Passes.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/LoopAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/BlockAndValueMapping.h" diff --git a/mlir/lib/Transforms/LowerAffine.cpp b/mlir/lib/Transforms/LowerAffine.cpp index 5a7d926d4f9..e8a8284d392 100644 --- a/mlir/lib/Transforms/LowerAffine.cpp +++ b/mlir/lib/Transforms/LowerAffine.cpp @@ -21,7 +21,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Transforms/LowerAffine.h" -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/LoopOps/LoopOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExprVisitor.h" diff --git a/mlir/lib/Transforms/MaterializeVectors.cpp b/mlir/lib/Transforms/MaterializeVectors.cpp index 0c6a3567ef3..bfdd5bf05f2 100644 --- a/mlir/lib/Transforms/MaterializeVectors.cpp +++ b/mlir/lib/Transforms/MaterializeVectors.cpp @@ -20,7 +20,6 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/Dominance.h" #include "mlir/Analysis/LoopAnalysis.h" @@ -28,6 +27,7 @@ #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Analysis/Utils.h" #include "mlir/Analysis/VectorAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/Dialect/VectorOps/VectorOps.h" #include "mlir/IR/AffineExpr.h" diff --git a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp index 33433e50d0f..9b71ada100c 100644 --- a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp +++ b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp @@ -22,10 +22,10 @@ // SSA scalars live out of 'affine.for'/'affine.if' statements is available. //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/Dominance.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/Passes.h" diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp index b58b6debc05..0cd979a1c82 100644 --- a/mlir/lib/Transforms/PipelineDataTransfer.cpp +++ b/mlir/lib/Transforms/PipelineDataTransfer.cpp @@ -21,10 +21,10 @@ #include "mlir/Transforms/Passes.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp b/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp index 63150c14742..8b314780c9f 100644 --- a/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopFusionUtils.cpp @@ -21,11 +21,11 @@ #include "mlir/Transforms/LoopFusionUtils.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp index 8b62d007f47..d6a31f92aed 100644 --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -21,11 +21,11 @@ #include "mlir/Transforms/LoopUtils.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/SliceAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/LoopOps/LoopOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExpr.h" diff --git a/mlir/lib/Transforms/Utils/Utils.cpp b/mlir/lib/Transforms/Utils/Utils.cpp index e2253c77f67..8d7b7a8b3a1 100644 --- a/mlir/lib/Transforms/Utils/Utils.cpp +++ b/mlir/lib/Transforms/Utils/Utils.cpp @@ -22,11 +22,11 @@ #include "mlir/Transforms/Utils.h" -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Dominance.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Function.h" diff --git a/mlir/lib/Transforms/Vectorize.cpp b/mlir/lib/Transforms/Vectorize.cpp index 08ee944dc45..cbf616eae10 100644 --- a/mlir/lib/Transforms/Vectorize.cpp +++ b/mlir/lib/Transforms/Vectorize.cpp @@ -20,12 +20,12 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/LoopAnalysis.h" #include "mlir/Analysis/NestedMatcher.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Analysis/Utils.h" #include "mlir/Analysis/VectorAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/Dialect/VectorOps/VectorOps.h" #include "mlir/IR/AffineExpr.h" diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp index f4c04eeb7fb..529aaa052f6 100644 --- a/mlir/test/EDSC/builder-api-test.cpp +++ b/mlir/test/EDSC/builder-api-test.cpp @@ -17,7 +17,7 @@ // RUN: mlir-edsc-builder-api-test | FileCheck %s -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/EDSC/Builders.h" #include "mlir/EDSC/Helpers.h" diff --git a/mlir/test/lib/Transforms/TestConstantFold.cpp b/mlir/test/lib/Transforms/TestConstantFold.cpp index 35a7eba5478..9c541699e99 100644 --- a/mlir/test/lib/Transforms/TestConstantFold.cpp +++ b/mlir/test/lib/Transforms/TestConstantFold.cpp @@ -15,7 +15,7 @@ // limitations under the License. // ============================================================================= -#include "mlir/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Function.h" diff --git a/mlir/test/lib/Transforms/TestLoopFusion.cpp b/mlir/test/lib/Transforms/TestLoopFusion.cpp index 4dd06a58904..604b42817e2 100644 --- a/mlir/test/lib/Transforms/TestLoopFusion.cpp +++ b/mlir/test/lib/Transforms/TestLoopFusion.cpp @@ -19,11 +19,11 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Passes.h" #include "mlir/Analysis/Utils.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Builders.h" #include "mlir/Pass/Pass.h" diff --git a/mlir/test/lib/Transforms/TestVectorizationUtils.cpp b/mlir/test/lib/Transforms/TestVectorizationUtils.cpp index 6fe277dcfcb..3f00eb01e11 100644 --- a/mlir/test/lib/Transforms/TestVectorizationUtils.cpp +++ b/mlir/test/lib/Transforms/TestVectorizationUtils.cpp @@ -19,11 +19,11 @@ // //===----------------------------------------------------------------------===// -#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/NestedMatcher.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Analysis/VectorAnalysis.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Diagnostics.h" #include "mlir/IR/StandardTypes.h" -- cgit v1.2.3