[Mlir-commits] [mlir] a78a223 - [MLIR] Migrate AffineDma ops to ODS (NFC) (#182497)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Feb 23 05:28:43 PST 2026
Author: Mehdi Amini
Date: 2026-02-23T14:28:38+01:00
New Revision: a78a22356b5acc89ab5f301e47f6bdd94568adb9
URL: https://github.com/llvm/llvm-project/commit/a78a22356b5acc89ab5f301e47f6bdd94568adb9
DIFF: https://github.com/llvm/llvm-project/commit/a78a22356b5acc89ab5f301e47f6bdd94568adb9.diff
LOG: [MLIR] Migrate AffineDma ops to ODS (NFC) (#182497)
For some historical reasons, seems like we never converted these.
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index 333de6bbd8a05..20d7f6436a896 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -58,339 +58,6 @@ Region *getAffineAnalysisScope(Operation *op);
OpFoldResult computeProduct(Location loc, OpBuilder &builder,
ArrayRef<OpFoldResult> terms);
-/// AffineDmaStartOp starts a non-blocking DMA operation that transfers data
-/// from a source memref to a destination memref. The source and destination
-/// memref need not be of the same dimensionality, but need to have the same
-/// elemental type. The operands include the source and destination memref's
-/// each followed by its indices, size of the data transfer in terms of the
-/// number of elements (of the elemental type of the memref), a tag memref with
-/// its indices, and optionally at the end, a stride and a
-/// number_of_elements_per_stride arguments. The tag location is used by an
-/// AffineDmaWaitOp to check for completion. The indices of the source memref,
-/// destination memref, and the tag memref have the same restrictions as any
-/// affine.load/store. In particular, index for each memref dimension must be an
-/// affine expression of loop induction variables and symbols.
-/// The optional stride arguments should be of 'index' type, and specify a
-/// stride for the slower memory space (memory space with a lower memory space
-/// id), transferring chunks of number_of_elements_per_stride every stride until
-/// %num_elements are transferred. Either both or no stride arguments should be
-/// specified. The value of 'num_elements' must be a multiple of
-/// 'number_of_elements_per_stride'. If the source and destination locations
-/// overlap the behavior of this operation is not defined.
-//
-// For example, an AffineDmaStartOp operation that transfers 256 elements of a
-// memref '%src' in memory space 0 at indices [%i + 3, %j] to memref '%dst' in
-// memory space 1 at indices [%k + 7, %l], would be specified as follows:
-//
-// %num_elements = arith.constant 256
-// %idx = arith.constant 0 : index
-// %tag = memref.alloc() : memref<1xi32, 4>
-// affine.dma_start %src[%i + 3, %j], %dst[%k + 7, %l], %tag[%idx],
-// %num_elements :
-// memref<40x128xf32, 0>, memref<2x1024xf32, 1>, memref<1xi32, 2>
-//
-// If %stride and %num_elt_per_stride are specified, the DMA is expected to
-// transfer %num_elt_per_stride elements every %stride elements apart from
-// memory space 0 until %num_elements are transferred.
-//
-// affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%idx], %num_elements,
-// %stride, %num_elt_per_stride : ...
-//
-// TODO: add additional operands to allow source and destination striding, and
-// multiple stride levels (possibly using AffineMaps to specify multiple levels
-// of striding).
-class AffineDmaStartOp
- : public Op<AffineDmaStartOp, OpTrait::MemRefsNormalizable,
- OpTrait::VariadicOperands, OpTrait::ZeroResults,
- OpTrait::OpInvariants, AffineMapAccessInterface::Trait,
- MemoryEffectOpInterface::Trait> {
-public:
- using Op::Op;
- static ArrayRef<StringRef> getAttributeNames() { return {}; }
-
- static void build(OpBuilder &builder, OperationState &result, Value srcMemRef,
- AffineMap srcMap, ValueRange srcIndices, Value destMemRef,
- AffineMap dstMap, ValueRange destIndices, Value tagMemRef,
- AffineMap tagMap, ValueRange tagIndices, Value numElements,
- Value stride = nullptr, Value elementsPerStride = nullptr);
-
- static AffineDmaStartOp
- create(OpBuilder &builder, Location location, Value srcMemRef,
- AffineMap srcMap, ValueRange srcIndices, Value destMemRef,
- AffineMap dstMap, ValueRange destIndices, Value tagMemRef,
- AffineMap tagMap, ValueRange tagIndices, Value numElements,
- Value stride = nullptr, Value elementsPerStride = nullptr);
-
- static AffineDmaStartOp create(ImplicitLocOpBuilder &builder, Value srcMemRef,
- AffineMap srcMap, ValueRange srcIndices,
- Value destMemRef, AffineMap dstMap,
- ValueRange destIndices, Value tagMemRef,
- AffineMap tagMap, ValueRange tagIndices,
- Value numElements, Value stride = nullptr,
- Value elementsPerStride = nullptr);
-
- /// Returns the operand index of the source memref.
- unsigned getSrcMemRefOperandIndex() { return 0; }
-
- /// Returns the source MemRefType for this DMA operation.
- Value getSrcMemRef() { return getOperand(getSrcMemRefOperandIndex()); }
- OpOperand &getSrcMemRefMutable() {
- return getOperation()->getOpOperand(getSrcMemRefOperandIndex());
- }
- MemRefType getSrcMemRefType() {
- return cast<MemRefType>(getSrcMemRef().getType());
- }
-
- /// Returns the rank (number of indices) of the source MemRefType.
- unsigned getSrcMemRefRank() { return getSrcMemRefType().getRank(); }
-
- /// Returns the affine map used to access the source memref.
- AffineMap getSrcMap() { return getSrcMapAttr().getValue(); }
- AffineMapAttr getSrcMapAttr() {
- return cast<AffineMapAttr>(
- *(*this)->getInherentAttr(getSrcMapAttrStrName()));
- }
-
- /// Returns the source memref affine map indices for this DMA operation.
- operand_range getSrcIndices() {
- return {operand_begin() + getSrcMemRefOperandIndex() + 1,
- operand_begin() + getSrcMemRefOperandIndex() + 1 +
- getSrcMap().getNumInputs()};
- }
-
- /// Returns the memory space of the source memref.
- unsigned getSrcMemorySpace() {
- return cast<MemRefType>(getSrcMemRef().getType()).getMemorySpaceAsInt();
- }
-
- /// Returns the operand index of the destination memref.
- unsigned getDstMemRefOperandIndex() {
- return getSrcMemRefOperandIndex() + 1 + getSrcMap().getNumInputs();
- }
-
- /// Returns the destination MemRefType for this DMA operation.
- Value getDstMemRef() { return getOperand(getDstMemRefOperandIndex()); }
- OpOperand &getDstMemRefMutable() {
- return getOperation()->getOpOperand(getDstMemRefOperandIndex());
- }
- MemRefType getDstMemRefType() {
- return cast<MemRefType>(getDstMemRef().getType());
- }
-
- /// Returns the rank (number of indices) of the destination MemRefType.
- unsigned getDstMemRefRank() {
- return cast<MemRefType>(getDstMemRef().getType()).getRank();
- }
-
- /// Returns the memory space of the source memref.
- unsigned getDstMemorySpace() {
- return cast<MemRefType>(getDstMemRef().getType()).getMemorySpaceAsInt();
- }
-
- /// Returns the affine map used to access the destination memref.
- AffineMap getDstMap() { return getDstMapAttr().getValue(); }
- AffineMapAttr getDstMapAttr() {
- return cast<AffineMapAttr>(
- *(*this)->getInherentAttr(getDstMapAttrStrName()));
- }
-
- /// 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()); }
- OpOperand &getTagMemRefMutable() {
- return getOperation()->getOpOperand(getTagMemRefOperandIndex());
- }
- MemRefType getTagMemRefType() {
- return cast<MemRefType>(getTagMemRef().getType());
- }
-
- /// Returns the rank (number of indices) of the tag MemRefType.
- unsigned getTagMemRefRank() {
- return cast<MemRefType>(getTagMemRef().getType()).getRank();
- }
-
- /// Returns the affine map used to access the tag memref.
- AffineMap getTagMap() { return getTagMapAttr().getValue(); }
- AffineMapAttr getTagMapAttr() {
- return cast<AffineMapAttr>(
- *(*this)->getInherentAttr(getTagMapAttrStrName()));
- }
-
- /// Returns the tag memref indices for this DMA operation.
- operand_range getTagIndices() {
- return {operand_begin() + getTagMemRefOperandIndex() + 1,
- operand_begin() + getTagMemRefOperandIndex() + 1 +
- getTagMap().getNumInputs()};
- }
-
- /// Returns the number of elements being transferred by this DMA operation.
- Value getNumElements() {
- return getOperand(getTagMemRefOperandIndex() + 1 +
- getTagMap().getNumInputs());
- }
-
- /// Impelements the AffineMapAccessInterface.
- /// Returns the AffineMapAttr associated with 'memref'.
- NamedAttribute getAffineMapAttrForMemRef(Value memref) {
- if (memref == getSrcMemRef())
- return {StringAttr::get(getContext(), getSrcMapAttrStrName()),
- getSrcMapAttr()};
- if (memref == getDstMemRef())
- return {StringAttr::get(getContext(), getDstMapAttrStrName()),
- getDstMapAttr()};
- assert(memref == getTagMemRef() &&
- "DmaStartOp expected source, destination or tag memref");
- return {StringAttr::get(getContext(), getTagMapAttrStrName()),
- 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();
- }
-
- void
- getEffects(SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
- &effects);
-
- static StringRef getSrcMapAttrStrName() { return "src_map"; }
- static StringRef getDstMapAttrStrName() { return "dst_map"; }
- static StringRef getTagMapAttrStrName() { return "tag_map"; }
-
- static StringRef getOperationName() { return "affine.dma_start"; }
- static ParseResult parse(OpAsmParser &parser, OperationState &result);
- void print(OpAsmPrinter &p);
- LogicalResult verifyInvariantsImpl();
- LogicalResult verifyInvariants() { return verifyInvariantsImpl(); }
- LogicalResult fold(ArrayRef<Attribute> cstOperands,
- SmallVectorImpl<OpFoldResult> &results);
-
- /// Returns true if this DMA operation is strided, returns false otherwise.
- bool isStrided() {
- return getNumOperands() !=
- getTagMemRefOperandIndex() + 1 + getTagMap().getNumInputs() + 1;
- }
-
- /// Returns the stride value for this DMA operation.
- Value getStride() {
- if (!isStrided())
- return nullptr;
- return getOperand(getNumOperands() - 1 - 1);
- }
-
- /// Returns the number of elements to transfer per stride for this DMA op.
- Value getNumElementsPerStride() {
- if (!isStrided())
- return nullptr;
- return getOperand(getNumOperands() - 1);
- }
-};
-
-/// AffineDmaWaitOp blocks until the completion of a DMA operation associated
-/// with the tag element '%tag[%index]'. %tag is a memref, and %index has to be
-/// an index with the same restrictions as any load/store index. In particular,
-/// index for each memref dimension must be an affine expression of loop
-/// induction variables and symbols. %num_elements is the number of elements
-/// associated with the DMA operation. For example:
-//
-// affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %num_elements :
-// memref<2048xf32, 0>, memref<256xf32, 1>, memref<1xi32, 2>
-// ...
-// ...
-// affine.dma_wait %tag[%index], %num_elements : memref<1xi32, 2>
-//
-class AffineDmaWaitOp
- : public Op<AffineDmaWaitOp, OpTrait::MemRefsNormalizable,
- OpTrait::VariadicOperands, OpTrait::ZeroResults,
- OpTrait::OpInvariants, AffineMapAccessInterface::Trait> {
-public:
- using Op::Op;
- static ArrayRef<StringRef> getAttributeNames() { return {}; }
-
- static void build(OpBuilder &builder, OperationState &result, Value tagMemRef,
- AffineMap tagMap, ValueRange tagIndices, Value numElements);
- static AffineDmaWaitOp create(OpBuilder &builder, Location location,
- Value tagMemRef, AffineMap tagMap,
- ValueRange tagIndices, Value numElements);
- static AffineDmaWaitOp create(ImplicitLocOpBuilder &builder, Value tagMemRef,
- AffineMap tagMap, ValueRange tagIndices,
- Value numElements);
-
- static StringRef getOperationName() { return "affine.dma_wait"; }
-
- /// Returns the Tag MemRef associated with the DMA operation being waited on.
- Value getTagMemRef() { return getOperand(0); }
- OpOperand &getTagMemRefMutable() { return getOperation()->getOpOperand(0); }
- MemRefType getTagMemRefType() {
- return cast<MemRefType>(getTagMemRef().getType());
- }
-
- /// Returns the affine map used to access the tag memref.
- AffineMap getTagMap() { return getTagMapAttr().getValue(); }
- AffineMapAttr getTagMapAttr() {
- return cast<AffineMapAttr>(
- *(*this)->getInherentAttr(getTagMapAttrStrName()));
- }
-
- /// 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 cast<MemRefType>(getTagMemRef().getType()).getRank();
- }
-
- /// Impelements the AffineMapAccessInterface. Returns the AffineMapAttr
- /// associated with 'memref'.
- NamedAttribute getAffineMapAttrForMemRef(Value memref) {
- assert(memref == getTagMemRef());
- return {StringAttr::get(getContext(), getTagMapAttrStrName()),
- getTagMapAttr()};
- }
-
- /// Returns the number of elements transferred by the associated DMA op.
- Value getNumElements() { return getOperand(1 + getTagMap().getNumInputs()); }
-
- static StringRef getTagMapAttrStrName() { return "tag_map"; }
- static ParseResult parse(OpAsmParser &parser, OperationState &result);
- void print(OpAsmPrinter &p);
- LogicalResult verifyInvariantsImpl();
- LogicalResult verifyInvariants() { return verifyInvariantsImpl(); }
- LogicalResult fold(ArrayRef<Attribute> cstOperands,
- SmallVectorImpl<OpFoldResult> &results);
- void
- getEffects(SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
- &effects);
-};
-
/// Returns true if the given Value can be used as a dimension id in the region
/// of the closest surrounding op that has the trait `AffineScope`.
bool isValidDim(Value value);
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 482987ebab27d..9cb0f3242db17 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -1267,4 +1267,309 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
let hasCanonicalizer = 1;
}
+//===----------------------------------------------------------------------===//
+// AffineDmaStartOp
+//===----------------------------------------------------------------------===//
+
+def AffineDmaStartOp : Affine_Op<"dma_start", [
+ MemRefsNormalizable,
+ DeclareOpInterfaceMethods<AffineMapAccessInterface>,
+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+ let summary = "affine dma start operation";
+ let description = [{
+ The `affine.dma_start` op 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
+ `affine.dma_wait` to check for completion. The indices of the source memref,
+ destination memref, and the tag memref have the same restrictions as any
+ affine.load/store. In particular, index for each memref dimension must be an
+ affine expression of loop induction variables and symbols.
+
+ The optional stride arguments should be of 'index' type, and specify a
+ stride for the slower memory space (memory space with a lower memory space
+ id), transferring chunks of number_of_elements_per_stride every stride until
+ %num_elements are transferred. Either both or no stride arguments should be
+ specified. The value of 'num_elements' must be a multiple of
+ 'number_of_elements_per_stride'. If the source and destination locations
+ overlap the behavior of this operation is not defined.
+
+ Example:
+
+ ```mlir
+ %num_elements = arith.constant 256
+ %idx = arith.constant 0 : index
+ %tag = memref.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 : ...
+ ```
+ }];
+
+ let arguments = (ins
+ Variadic<AnyType>,
+ AffineMapAttr:$src_map,
+ AffineMapAttr:$dst_map,
+ AffineMapAttr:$tag_map);
+
+ let results = (outs);
+
+ let skipDefaultBuilders = 1;
+ let builders = [
+ OpBuilder<(ins "Value":$srcMemRef, "AffineMap":$srcMap,
+ "ValueRange":$srcIndices, "Value":$destMemRef, "AffineMap":$dstMap,
+ "ValueRange":$destIndices, "Value":$tagMemRef, "AffineMap":$tagMap,
+ "ValueRange":$tagIndices, "Value":$numElements,
+ CArg<"Value", "nullptr">:$stride,
+ CArg<"Value", "nullptr">:$elementsPerStride)>
+ ];
+
+ let extraClassDeclaration = [{
+ /// Returns the operand index of the source memref.
+ unsigned getSrcMemRefOperandIndex() { return 0; }
+
+ /// Returns the source MemRefType for this DMA operation.
+ Value getSrcMemRef() { return getOperand(getSrcMemRefOperandIndex()); }
+ OpOperand &getSrcMemRefMutable() {
+ return getOperation()->getOpOperand(getSrcMemRefOperandIndex());
+ }
+ MemRefType getSrcMemRefType() {
+ return ::llvm::cast<MemRefType>(getSrcMemRef().getType());
+ }
+
+ /// Returns the rank (number of indices) of the source MemRefType.
+ unsigned getSrcMemRefRank() { return getSrcMemRefType().getRank(); }
+
+ /// Returns the source memref affine map indices for this DMA operation.
+ operand_range getSrcIndices() {
+ return {operand_begin() + getSrcMemRefOperandIndex() + 1,
+ operand_begin() + getSrcMemRefOperandIndex() + 1 +
+ getSrcMap().getNumInputs()};
+ }
+
+ /// Returns the memory space of the source memref.
+ unsigned getSrcMemorySpace() {
+ return ::llvm::cast<MemRefType>(getSrcMemRef().getType())
+ .getMemorySpaceAsInt();
+ }
+
+ /// Returns the operand index of the destination memref.
+ unsigned getDstMemRefOperandIndex() {
+ return getSrcMemRefOperandIndex() + 1 + getSrcMap().getNumInputs();
+ }
+
+ /// Returns the destination MemRefType for this DMA operation.
+ Value getDstMemRef() { return getOperand(getDstMemRefOperandIndex()); }
+ OpOperand &getDstMemRefMutable() {
+ return getOperation()->getOpOperand(getDstMemRefOperandIndex());
+ }
+ MemRefType getDstMemRefType() {
+ return ::llvm::cast<MemRefType>(getDstMemRef().getType());
+ }
+
+ /// Returns the rank (number of indices) of the destination MemRefType.
+ unsigned getDstMemRefRank() {
+ return ::llvm::cast<MemRefType>(getDstMemRef().getType()).getRank();
+ }
+
+ /// Returns the memory space of the destination memref.
+ unsigned getDstMemorySpace() {
+ return ::llvm::cast<MemRefType>(getDstMemRef().getType())
+ .getMemorySpaceAsInt();
+ }
+
+ /// 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()); }
+ OpOperand &getTagMemRefMutable() {
+ return getOperation()->getOpOperand(getTagMemRefOperandIndex());
+ }
+ MemRefType getTagMemRefType() {
+ return ::llvm::cast<MemRefType>(getTagMemRef().getType());
+ }
+
+ /// Returns the rank (number of indices) of the tag MemRefType.
+ unsigned getTagMemRefRank() {
+ return ::llvm::cast<MemRefType>(getTagMemRef().getType()).getRank();
+ }
+
+ /// 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());
+ }
+
+ /// Implements the AffineMapAccessInterface.
+ /// Returns the AffineMapAttr associated with 'memref'.
+ NamedAttribute getAffineMapAttrForMemRef(Value memref) {
+ if (memref == getSrcMemRef())
+ return {StringAttr::get(getContext(), getSrcMapAttrStrName()),
+ getSrcMapAttr()};
+ if (memref == getDstMemRef())
+ return {StringAttr::get(getContext(), getDstMapAttrStrName()),
+ getDstMapAttr()};
+ assert(memref == getTagMemRef() &&
+ "DmaStartOp expected source, destination or tag memref");
+ return {StringAttr::get(getContext(), getTagMapAttrStrName()),
+ 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() {
+ return (getDstMemorySpace() < getSrcMemorySpace());
+ }
+
+ /// Returns the operand position of either the source or destination memref
+ /// depending on which is at the higher level of the memory hierarchy.
+ unsigned getFasterMemPos() {
+ assert(isSrcMemorySpaceFaster() || isDestMemorySpaceFaster());
+ return isSrcMemorySpaceFaster() ? 0 : getDstMemRefOperandIndex();
+ }
+
+ static StringRef getSrcMapAttrStrName() { return "src_map"; }
+ static StringRef getDstMapAttrStrName() { return "dst_map"; }
+ static StringRef getTagMapAttrStrName() { return "tag_map"; }
+
+ /// 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);
+ }
+
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+ let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// AffineDmaWaitOp
+//===----------------------------------------------------------------------===//
+
+def AffineDmaWaitOp : Affine_Op<"dma_wait", [
+ MemRefsNormalizable,
+ DeclareOpInterfaceMethods<AffineMapAccessInterface>,
+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+ let summary = "affine dma wait operation";
+ let description = [{
+ The `affine.dma_wait` op 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.
+
+ Example:
+
+ ```mlir
+ 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>
+ ```
+ }];
+
+ let arguments = (ins
+ Variadic<AnyType>,
+ AffineMapAttr:$tag_map);
+
+ let results = (outs);
+
+ let skipDefaultBuilders = 1;
+ let builders = [
+ OpBuilder<(ins "Value":$tagMemRef, "AffineMap":$tagMap,
+ "ValueRange":$tagIndices, "Value":$numElements)>
+ ];
+
+ let extraClassDeclaration = [{
+ /// Returns the tag MemRef associated with the DMA operation being waited on.
+ Value getTagMemRef() { return getOperand(0); }
+ OpOperand &getTagMemRefMutable() {
+ return getOperation()->getOpOperand(0);
+ }
+ MemRefType getTagMemRefType() {
+ return ::llvm::cast<MemRefType>(getTagMemRef().getType());
+ }
+
+ /// 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 ::llvm::cast<MemRefType>(getTagMemRef().getType()).getRank();
+ }
+
+ /// Implements the AffineMapAccessInterface. Returns the AffineMapAttr
+ /// associated with 'memref'.
+ NamedAttribute getAffineMapAttrForMemRef(Value memref) {
+ assert(memref == getTagMemRef());
+ return {StringAttr::get(getContext(), getTagMapAttrStrName()),
+ getTagMapAttr()};
+ }
+
+ /// Returns the number of elements transferred by the associated DMA op.
+ Value getNumElements() {
+ return getOperand(1 + getTagMap().getNumInputs());
+ }
+
+ static StringRef getTagMapAttrStrName() { return "tag_map"; }
+
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+ let hasFolder = 1;
+}
+
#endif // AFFINE_OPS
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 871c7df2f71e5..a5e4b1c4d0a6f 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -224,10 +224,10 @@ struct AffineInlinerInterface : public DialectInlinerInterface {
//===----------------------------------------------------------------------===//
void AffineDialect::initialize() {
- addOperations<AffineDmaStartOp, AffineDmaWaitOp,
+ addOperations<
#define GET_OP_LIST
#include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"
- >();
+ >();
addInterfaces<AffineInlinerInterface>();
declarePromisedInterfaces<ValueBoundsOpInterface, AffineApplyOp, AffineMaxOp,
AffineMinOp>();
@@ -1885,32 +1885,6 @@ void AffineDmaStartOp::build(OpBuilder &builder, OperationState &result,
}
}
-AffineDmaStartOp AffineDmaStartOp::create(
- OpBuilder &builder, Location location, Value srcMemRef, AffineMap srcMap,
- ValueRange srcIndices, Value destMemRef, AffineMap dstMap,
- ValueRange destIndices, Value tagMemRef, AffineMap tagMap,
- ValueRange tagIndices, Value numElements, Value stride,
- Value elementsPerStride) {
- mlir::OperationState state(location, getOperationName());
- build(builder, state, srcMemRef, srcMap, srcIndices, destMemRef, dstMap,
- destIndices, tagMemRef, tagMap, tagIndices, numElements, stride,
- elementsPerStride);
- auto result = dyn_cast<AffineDmaStartOp>(builder.create(state));
- assert(result && "builder didn't return the right type");
- return result;
-}
-
-AffineDmaStartOp AffineDmaStartOp::create(
- ImplicitLocOpBuilder &builder, Value srcMemRef, AffineMap srcMap,
- ValueRange srcIndices, Value destMemRef, AffineMap dstMap,
- ValueRange destIndices, Value tagMemRef, AffineMap tagMap,
- ValueRange tagIndices, Value numElements, Value stride,
- Value elementsPerStride) {
- return create(builder, builder.getLoc(), srcMemRef, srcMap, srcIndices,
- destMemRef, dstMap, destIndices, tagMemRef, tagMap, tagIndices,
- numElements, stride, elementsPerStride);
-}
-
void AffineDmaStartOp::print(OpAsmPrinter &p) {
p << " " << getSrcMemRef() << '[';
p.printAffineMapOfSSAIds(getSrcMapAttr(), getSrcIndices());
@@ -2009,7 +1983,7 @@ ParseResult AffineDmaStartOp::parse(OpAsmParser &parser,
return success();
}
-LogicalResult AffineDmaStartOp::verifyInvariantsImpl() {
+LogicalResult AffineDmaStartOp::verify() {
if (!llvm::isa<MemRefType>(getOperand(getSrcMemRefOperandIndex()).getType()))
return emitOpError("expected DMA source to be of memref type");
if (!llvm::isa<MemRefType>(getOperand(getDstMemRefOperandIndex()).getType()))
@@ -2050,7 +2024,7 @@ LogicalResult AffineDmaStartOp::verifyInvariantsImpl() {
return success();
}
-LogicalResult AffineDmaStartOp::fold(ArrayRef<Attribute> cstOperands,
+LogicalResult AffineDmaStartOp::fold(FoldAdaptor adaptor,
SmallVectorImpl<OpFoldResult> &results) {
/// dma_start(memrefcast) -> dma_start
return memref::foldMemRefCast(*this);
@@ -2081,25 +2055,6 @@ void AffineDmaWaitOp::build(OpBuilder &builder, OperationState &result,
result.addOperands(numElements);
}
-AffineDmaWaitOp AffineDmaWaitOp::create(OpBuilder &builder, Location location,
- Value tagMemRef, AffineMap tagMap,
- ValueRange tagIndices,
- Value numElements) {
- mlir::OperationState state(location, getOperationName());
- build(builder, state, tagMemRef, tagMap, tagIndices, numElements);
- auto result = dyn_cast<AffineDmaWaitOp>(builder.create(state));
- assert(result && "builder didn't return the right type");
- return result;
-}
-
-AffineDmaWaitOp AffineDmaWaitOp::create(ImplicitLocOpBuilder &builder,
- Value tagMemRef, AffineMap tagMap,
- ValueRange tagIndices,
- Value numElements) {
- return create(builder, builder.getLoc(), tagMemRef, tagMap, tagIndices,
- numElements);
-}
-
void AffineDmaWaitOp::print(OpAsmPrinter &p) {
p << " " << getTagMemRef() << '[';
SmallVector<Value, 2> operands(getTagIndices());
@@ -2145,7 +2100,7 @@ ParseResult AffineDmaWaitOp::parse(OpAsmParser &parser,
return success();
}
-LogicalResult AffineDmaWaitOp::verifyInvariantsImpl() {
+LogicalResult AffineDmaWaitOp::verify() {
if (!llvm::isa<MemRefType>(getOperand(0).getType()))
return emitOpError("expected DMA tag to be of memref type");
Region *scope = getAffineScope(*this);
@@ -2159,7 +2114,7 @@ LogicalResult AffineDmaWaitOp::verifyInvariantsImpl() {
return success();
}
-LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
+LogicalResult AffineDmaWaitOp::fold(FoldAdaptor adaptor,
SmallVectorImpl<OpFoldResult> &results) {
/// dma_wait(memrefcast) -> dma_wait
return memref::foldMemRefCast(*this);
More information about the Mlir-commits
mailing list