add pipeline stages synchronization

full ops throughput validation now passes

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -368,11 +368,14 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
           return failure();
         PimWaitOp::create(
           rewriter, receiveOp->getLoc(), hostWaitLoad.getEventRegister(),
-          rewriter.getI32IntegerAttr(1));
+          hostWaitLoad.getWaitValue());
         received = PimMemCopyHostToDevOp::create(
           rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
           hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
           .getOutput();
+        PimSyncOp::create(
+          rewriter, receiveOp->getLoc(), hostWaitLoad.getSourceCoreId(),
+          hostWaitLoad.getAcknowledgementEventRegister());
       } else {
         received = PimReceiveOp::create(
           rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,

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

@@ -147,15 +147,42 @@ struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
         return failure();
       auto wait = pim::PimWaitOp::create(
         rewriter, op.getLoc(), op.getEventRegister(),
-        rewriter.getI32IntegerAttr(1));
+        op.getWaitValue());
       copyRaptorDebugAttrs(op.getOperation(), wait.getOperation());
-      return pim::PimMemCopyHostToDevOp::create(
+      Value output = pim::PimMemCopyHostToDevOp::create(
         rewriter, op.getLoc(), outputBuffer.getType(), zero,
         op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput();
+      auto sync = pim::PimSyncOp::create(
+        rewriter, op.getLoc(), op.getSourceCoreId(),
+        op.getAcknowledgementEventRegister());
+      copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
+      return output;
     });
   }
 };
 
+struct SyncLowering : OpRewritePattern<spatial::SpatSyncOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSyncOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<pim::PimSyncOp>(
+      op, op.getTargetCoreId(), op.getEventRegister());
+    return success();
+  }
+};
+
+struct WaitLowering : OpRewritePattern<spatial::SpatWaitOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatWaitOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<pim::PimWaitOp>(
+      op, op.getEventRegister(), op.getWaitValue());
+    return success();
+  }
+};
+
 struct ExtractRowsLowering : OpRewritePattern<spatial::SpatExtractRowsOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -200,7 +227,8 @@ struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
 void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
   patterns.add<ChannelSendLowering, ChannelReceiveLowering,
                HostStoreSyncLowering, HostWaitLoadLowering,
-               ExtractRowsLowering, ConcatLowering>(patterns.getContext());
+               SyncLowering, WaitLowering, ExtractRowsLowering,
+               ConcatLowering>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -128,6 +128,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                     spatial::SpatChannelSendOp,
                     spatial::SpatHostStoreSyncOp,
                     spatial::SpatHostWaitLoadOp,
+                    spatial::SpatSyncOp,
+                    spatial::SpatWaitOp,
                     spatial::SpatExtractRowsOp>();
 
   RewritePatternSet initialPatterns(ctx);
@@ -223,6 +225,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                             spatial::SpatChannelSendOp,
                             spatial::SpatHostStoreSyncOp,
                             spatial::SpatHostWaitLoadOp,
+                            spatial::SpatSyncOp,
+                            spatial::SpatWaitOp,
                             spatial::SpatExtractRowsOp>();
 
   SmallVector<pim::PimCoreOp> coreOps;
@@ -274,6 +278,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                                    spatial::SpatChannelSendOp,
                                    spatial::SpatHostStoreSyncOp,
                                    spatial::SpatHostWaitLoadOp,
+                                   spatial::SpatSyncOp,
+                                   spatial::SpatWaitOp,
                                    spatial::SpatExtractRowsOp>();
 
   RewritePatternSet communicationPatterns(ctx);