Raptor sync wait

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,7 +1,10 @@
 #include "mlir/IR/ValueRange.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
 
 #include "llvm/ADT/STLExtras.h"
 
@@ -28,6 +31,49 @@ FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anc
   return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
 }
 
+LogicalResult materializePipelineHostBuffer(
+    func::FuncOp funcOp, RewriterBase &rewriter) {
+  auto bytes = funcOp->getAttrOfType<IntegerAttr>(
+    kPipelineHostBufferBytesAttrName);
+  if (!bytes)
+    return success();
+  if (bytes.getInt() <= 0)
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer must be positive");
+  ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
+  if (moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName))
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer symbol already exists");
+  auto type = MemRefType::get(
+    {bytes.getInt()}, rewriter.getI8Type());
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(moduleOp.getBody());
+  memref::GlobalOp::create(
+    rewriter, funcOp.getLoc(),
+    rewriter.getStringAttr(kPipelineHostBufferName),
+    rewriter.getStringAttr("private"), TypeAttr::get(type), Attribute(),
+    UnitAttr(), IntegerAttr());
+  return success();
+}
+
+FailureOr<mlir::Value> getPipelineHostBuffer(
+    OpBuilder &builder, Operation *anchor) {
+  auto funcOp = anchor->getParentOfType<func::FuncOp>();
+  auto moduleOp = anchor->getParentOfType<ModuleOp>();
+  auto bytes = funcOp
+    ? funcOp->getAttrOfType<IntegerAttr>(kPipelineHostBufferBytesAttrName)
+    : IntegerAttr();
+  auto global = moduleOp
+    ? moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName)
+    : memref::GlobalOp();
+  if (!bytes || !global)
+    return anchor->emitOpError(
+      "requires the pipeline host transfer buffer"), failure();
+  auto type = MemRefType::get({bytes.getInt()}, builder.getI8Type());
+  return memref::GetGlobalOp::create(
+    builder, anchor->getLoc(), type, kPipelineHostBufferName).getResult();
+}
+
 Operation* getEarliestUserWithinBlock(mlir::Value value) {
   auto users = value.getUsers();