Raptor sync wait

src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

@@ -57,10 +57,29 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   }
 };
 
-struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+struct HostStoreSyncLowering : OpRewritePattern<spatial::SpatHostStoreSyncOp> {
   using OpRewritePattern::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatHostStoreSyncOp op, PatternRewriter& rewriter) const override {
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
+    auto hostBuffer = getPipelineHostBuffer(rewriter, op);
+    if (failed(sizeAttr) || failed(hostBuffer))
+      return failure();
+    Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
+    pim::PimMemCopyDevToHostOp::create(
+      rewriter, op.getLoc(), hostBuffer->getType(), op.getHostOffset(), zero,
+      *hostBuffer, op.getInput(), *sizeAttr);
+    auto sync = pim::PimSyncOp::create(
+      rewriter, op.getLoc(), op.getTargetCoreId(), op.getEventRegister());
+    copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+template <typename ReceiveOp, typename CreateReceive>
+static LogicalResult lowerReceive(
+    ReceiveOp op, PatternRewriter& rewriter, CreateReceive createReceive) {
     if (op->use_empty()) {
       rewriter.eraseOp(op);
       return success();
@@ -86,12 +105,11 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     if (failed(sizeAttr))
       return failure();
     Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
-    auto receive = pim::PimReceiveOp::create(
-      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId());
-    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
-    Value received = receive.getOutput();
+    auto received = createReceive(outputBuffer, zero, *sizeAttr);
+    if (failed(received))
+      return failure();
     if (!destinationInsert) {
-      rewriter.replaceOp(op, received);
+      rewriter.replaceOp(op, *received);
       return success();
     }
 
@@ -99,10 +117,42 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
     auto copy = pim::PimMemCopyOp::create(
       rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
-      destinationInsert.getDest(), received, *sizeAttr);
+      destinationInsert.getDest(), *received, *sizeAttr);
     rewriter.replaceOp(destinationInsert, copy.getOutput());
     rewriter.eraseOp(op);
     return success();
+}
+
+struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+    return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
+      auto receive = pim::PimReceiveOp::create(
+        rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero,
+        sizeAttr, op.getSourceCoreId());
+      copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
+      return receive.getOutput();
+    });
+  }
+};
+
+struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatHostWaitLoadOp op, PatternRewriter& rewriter) const override {
+    return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
+      auto hostBuffer = getPipelineHostBuffer(rewriter, op);
+      if (failed(hostBuffer))
+        return failure();
+      auto wait = pim::PimWaitOp::create(
+        rewriter, op.getLoc(), op.getEventRegister(),
+        rewriter.getI32IntegerAttr(1));
+      copyRaptorDebugAttrs(op.getOperation(), wait.getOperation());
+      return pim::PimMemCopyHostToDevOp::create(
+        rewriter, op.getLoc(), outputBuffer.getType(), zero,
+        op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput();
+    });
   }
 };
 
@@ -148,7 +198,9 @@ struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
 } // namespace
 
 void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
-  patterns.add<ChannelSendLowering, ChannelReceiveLowering, ExtractRowsLowering, ConcatLowering>(patterns.getContext());
+  patterns.add<ChannelSendLowering, ChannelReceiveLowering,
+               HostStoreSyncLowering, HostWaitLoadLowering,
+               ExtractRowsLowering, ConcatLowering>(patterns.getContext());
 }
 
 } // namespace onnx_mlir