[Mlir-commits] [mlir] cd67bbd - [MLIR][Affine] add memory effect traits for dmaOp

[Uday Bondhugula]

Author: lipracer
Date: 2023-03-26T06:40:10+05:30
New Revision: cd67bbdc2496fbba68c818c6ff82007d89d2bb40

URL: https://github.com/llvm/llvm-project/commit/cd67bbdc2496fbba68c818c6ff82007d89d2bb40
DIFF: https://github.com/llvm/llvm-project/commit/cd67bbdc2496fbba68c818c6ff82007d89d2bb40.diff

LOG: [MLIR][Affine] add memory effect traits for dmaOp

DmaOp will read the source buffer and write the destination buffer so need to add some traits for it.

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D144712

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
    mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.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 cda38a5580645..3d46bfd016a8c 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -89,7 +89,8 @@ Region *getAffineScope(Operation *op);
 class AffineDmaStartOp
     : public Op<AffineDmaStartOp, OpTrait::MemRefsNormalizable,
                 OpTrait::VariadicOperands, OpTrait::ZeroResults,
-                OpTrait::OpInvariants, AffineMapAccessInterface::Trait> {
+                OpTrait::OpInvariants, AffineMapAccessInterface::Trait,
+                MemoryEffectOpInterface::Trait> {
 public:
   using Op::Op;
   static ArrayRef<StringRef> getAttributeNames() { return {}; }
@@ -233,6 +234,10 @@ class AffineDmaStartOp
     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"; }
@@ -333,6 +338,9 @@ class AffineDmaWaitOp
   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

diff  --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 45b5c9f21b072..b37054b1c5e91 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -761,6 +761,17 @@ def MemRef_DmaStartOp : MemRef_Op<"dma_start"> {
         return nullptr;
       return getOperand(getNumOperands() - 1);
     }
+
+    void getEffects(
+        SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &
+        effects) {
+      effects.emplace_back(MemoryEffects::Read::get(), getSrcMemRef(),
+                           SideEffects::DefaultResource::get());
+      effects.emplace_back(MemoryEffects::Write::get(), getDstMemRef(),
+                           SideEffects::DefaultResource::get());
+      effects.emplace_back(MemoryEffects::Read::get(), getTagMemRef(),
+                           SideEffects::DefaultResource::get());
+    }
   }];
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;
@@ -803,6 +814,12 @@ def MemRef_DmaWaitOp : MemRef_Op<"dma_wait"> {
     unsigned getTagMemRefRank() {
       return getTagMemRef().getType().cast<MemRefType>().getRank();
     }
+    void getEffects(
+        SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &
+        effects) {
+      effects.emplace_back(MemoryEffects::Read::get(), getTagMemRef(),
+                           SideEffects::DefaultResource::get());
+    }
   }];
   let hasFolder = 1;
   let hasVerifier = 1;

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 5ed434fa1e809..e5c88caefdf5b 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1708,6 +1708,17 @@ LogicalResult AffineDmaStartOp::fold(ArrayRef<Attribute> cstOperands,
   return memref::foldMemRefCast(*this);
 }
 
+void AffineDmaStartOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  effects.emplace_back(MemoryEffects::Read::get(), getSrcMemRef(),
+                       SideEffects::DefaultResource::get());
+  effects.emplace_back(MemoryEffects::Write::get(), getDstMemRef(),
+                       SideEffects::DefaultResource::get());
+  effects.emplace_back(MemoryEffects::Read::get(), getTagMemRef(),
+                       SideEffects::DefaultResource::get());
+}
+
 //===----------------------------------------------------------------------===//
 // AffineDmaWaitOp
 //===----------------------------------------------------------------------===//
@@ -1786,6 +1797,13 @@ LogicalResult AffineDmaWaitOp::fold(ArrayRef<Attribute> cstOperands,
   return memref::foldMemRefCast(*this);
 }
 
+void AffineDmaWaitOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  effects.emplace_back(MemoryEffects::Read::get(), getTagMemRef(),
+                       SideEffects::DefaultResource::get());
+}
+
 //===----------------------------------------------------------------------===//
 // AffineForOp
 //===----------------------------------------------------------------------===//