Add objectFifo bufferization pass (nod-ai#781)

Introduces new buffer related ops and a pass to bufferize logical objectFifos. This is part of the effort to move `StatefulTransform` logic into `AMDAIE` logic: https://github.com/nod-ai/iree-amd-aie/compare/main...jtuyls:iree-amd-aie:delete-stateful-transform?expand=1. The newly introduced ops are: - `amdaie.buffer`: buffer op which will get translated into `aie.buffer` - `amdaie.lock`: lock op which will get translated into `aie.lock` - `amdaie.logicalobjectfifo.from_buffers`: logical objectFifo op created from a set of buffers. Will help with lowering to synchronized DMA patterns (`aie.mem`, `aie.memtile_dma`), without the need for AIE objectFifos (`aie.objectfifo`, `aie.objectfifo.link`).
newling · Sep 17, 2024 · e31b56a · e31b56a
1 parent 2cf234c
commit e31b56a
Show file tree

Hide file tree

Showing 18 changed files with 894 additions and 1 deletion.
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.cpp
@@ -8,3 +8,37 @@
 
 /// Include the definitions of the logical-objFifo-like interfaces.
 #include "iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.cpp.inc"
+
+namespace mlir::iree_compiler::AMDAIE {
+
+namespace detail {
+
+SmallVector<mlir::CopyOpInterface> getCopyLikeConsumers(
+    LogicalObjFifoOpInterface op) {
+  SmallVector<mlir::CopyOpInterface> copyLikOps;
+  for (Operation *userOp : op->getUsers()) {
+    if (auto copyOp = dyn_cast<CopyOpInterface>(userOp);
+        dyn_cast_if_present<LogicalObjFifoOpInterface>(
+            copyOp.getSource().getDefiningOp()) == op) {
+      copyLikOps.push_back(copyOp);
+    }
+  }
+  return copyLikOps;
+}
+
+SmallVector<mlir::CopyOpInterface> getCopyLikeProducers(
+    LogicalObjFifoOpInterface op) {
+  SmallVector<mlir::CopyOpInterface> copyLikOps;
+  for (Operation *userOp : op->getUsers()) {
+    if (auto copyOp = dyn_cast<CopyOpInterface>(userOp);
+        dyn_cast_if_present<LogicalObjFifoOpInterface>(
+            copyOp.getTarget().getDefiningOp()) == op) {
+      copyLikOps.push_back(copyOp);
+    }
+  }
+  return copyLikOps;
+}
+
+}  // namespace detail
+
+}  // namespace mlir::iree_compiler::AMDAIE
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.h b/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.h
@@ -8,6 +8,25 @@
 #define IREE_COMPILER_AMDAIE_LOGICALOBJFIFOOPINTERFACE_H_
 
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/CopyOpInterface.h"
+
+namespace mlir::iree_compiler::AMDAIE {
+
+class LogicalObjFifoOpInterface;
+
+namespace detail {
+
+/// Return the consumer copy-like operations of the logical objFifo.
+SmallVector<mlir::CopyOpInterface> getCopyLikeConsumers(
+    LogicalObjFifoOpInterface op);
+
+/// Return the producer copy-like operations of the logical objFifo.
+SmallVector<mlir::CopyOpInterface> getCopyLikeProducers(
+    LogicalObjFifoOpInterface op);
+
+}  // namespace detail
+
+}  // namespace mlir::iree_compiler::AMDAIE
 
 // clang-format off
 #include "iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.h.inc"

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.td b/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIELogicalObjFifoOpInterface.td
@@ -21,6 +21,32 @@ def LogicalObjFifoOpInterface : OpInterface<"LogicalObjFifoOpInterface"> {
   let cppNamespace = "mlir::iree_compiler::AMDAIE";
 
   let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the consumer copy-like operations of the logical objFifo.
+      }],
+      /*retTy=*/"::llvm::SmallVector<::mlir::CopyOpInterface>",
+      /*methodName=*/"getCopyLikeConsumers",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return detail::getCopyLikeConsumers(
+          ::mlir::cast<LogicalObjFifoOpInterface>($_op.getOperation()));
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the producer copy-like operations of the logical objFifo.
+      }],
+      /*retTy=*/"::llvm::SmallVector<::mlir::CopyOpInterface>",
+      /*methodName=*/"getCopyLikeProducers",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return detail::getCopyLikeProducers(
+          ::mlir::cast<LogicalObjFifoOpInterface>($_op.getOperation()));
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/"Return the buffer depth of the logical objectFifo. (E.g. 1 == "
                "single buffer, 2 == double buffer)",
@@ -32,9 +58,39 @@ def LogicalObjFifoOpInterface : OpInterface<"LogicalObjFifoOpInterface"> {
         return $_op.getDepth();
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Return the memory space attribute.",
+      /*retTy=*/"mlir::Attribute",
+      /*methodName=*/"getMemorySpace",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getMemorySpace();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the memory space attribute as an integer.",
+      /*retTy=*/"uint8_t",
+      /*methodName=*/"getMemorySpaceAsUInt",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getMemorySpaceAsUInt();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/"Return the memref type.",
+      /*retTy=*/"mlir::MemRefType",
+      /*methodName=*/"getMemrefType",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getMemrefType();
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/"Return the assigned tiles.",
-      /*retTy=*/"::mlir::ValueRange",
+      /*retTy=*/"::llvm::SmallVector<mlir::Value>",
       /*methodName=*/"getTiles",
       /*args=*/(ins),
       /*methodBody=*/"",

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.cpp
@@ -32,6 +32,15 @@ void BdIdOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "bd_id");
 }
 
+//===----------------------------------------------------------------------===//
+// AMDAIE_BufferOp
+//===----------------------------------------------------------------------===//
+
+void BufferOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+  setNameFn(getResult(), "buffer");
+}
+
 //===----------------------------------------------------------------------===//
 // AMDAIE_ChannelOp
 //===----------------------------------------------------------------------===//
@@ -446,6 +455,14 @@ ConnectionOp::getNpuCircularDmaCpyNdUser() {
   return npuDmaUsers[0];
 }
 
+//===----------------------------------------------------------------------===//
+// AMDAIE_LockOp
+//===----------------------------------------------------------------------===//
+
+void LockOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
+  setNameFn(getResult(), "lock");
+}
+
 //===----------------------------------------------------------------------===//
 // AMDAIE_LogicalObjectFifoAccessOp
 //===----------------------------------------------------------------------===//
@@ -473,6 +490,67 @@ void LogicalObjectFifoAcquire::build(OpBuilder &b, mlir::OperationState &result,
   build(b, result, resultTypes, dma, port, b.getI32IntegerAttr(1));
 }
 
+//===----------------------------------------------------------------------===//
+// AMDAIE_LogicalObjectFifoFromBuffersOp
+//===----------------------------------------------------------------------===//
+
+SmallVector<AMDAIE::BufferOp> LogicalObjectFifoFromBuffersOp::getBuffersOnTile(
+    TileOp tileOp) {
+  SmallVector<AMDAIE::BufferOp> buffers;
+  for (Value res : getBuffers()) {
+    if (auto bufferOp =
+            dyn_cast_if_present<AMDAIE::BufferOp>(res.getDefiningOp());
+        bufferOp.getTile() == tileOp.getResult()) {
+      buffers.push_back(bufferOp);
+    }
+  }
+  return buffers;
+}
+
+SmallVector<AMDAIE::LockOp>
+LogicalObjectFifoFromBuffersOp::getConsumerLocksOnTile(AMDAIE::TileOp tileOp) {
+  SmallVector<AMDAIE::LockOp> locks;
+  for (Value lockRes : getConsumerLocks()) {
+    if (auto lockOp =
+            dyn_cast_if_present<AMDAIE::LockOp>(lockRes.getDefiningOp());
+        lockOp.getTile() == tileOp.getResult()) {
+      locks.push_back(lockOp);
+    }
+  }
+  return locks;
+}
+
+SmallVector<AMDAIE::LockOp>
+LogicalObjectFifoFromBuffersOp::getProducerLocksOnTile(AMDAIE::TileOp tileOp) {
+  SmallVector<AMDAIE::LockOp> locks;
+  for (Value lockRes : getProducerLocks()) {
+    if (auto lockOp =
+            dyn_cast_if_present<AMDAIE::LockOp>(lockRes.getDefiningOp());
+        lockOp.getTile() == tileOp.getResult()) {
+      locks.push_back(lockOp);
+    }
+  }
+  return locks;
+}
+
+SmallVector<Value> LogicalObjectFifoFromBuffersOp::getTiles() {
+  llvm::SmallVector<mlir::Value> tiles;
+  for (Value result : getBuffers()) {
+    Value tile = cast<AMDAIE::BufferOp>(result.getDefiningOp()).getTile();
+    tiles.push_back(tile);
+  }
+  return tiles;
+}
+
+LogicalResult LogicalObjectFifoFromBuffersOp::verify() {
+  if (llvm::any_of(getBuffers(), [](Value result) {
+        return !isa_and_present<AMDAIE::BufferOp>(result.getDefiningOp());
+      })) {
+    return failure();
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // AMDAIE_LogicalObjectFifoFromMemrefOp
 //===----------------------------------------------------------------------===//