add support for P2PCommunication in HIRExecutor

NVIDIA · Sep 19, 2024 · f969cfa · f969cfa
1 parent c48ea11
commit f969cfa
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/csrc/host_ir/executor.cpp b/csrc/host_ir/executor.cpp
@@ -183,6 +183,23 @@ void HostIrExecutor::handle(Communication* communication) {
       output_tensor);
 }
 
+void HostIrExecutor::handle(P2PCommunication* communication) {
+  NVF_ERROR(
+      communicator_ != nullptr && communicator_->is_available(),
+      "A valid communicator must be provided");
+
+  at::Tensor buffer =
+      getKnownTensorOrUndefined(communication->buffer(), expr_evaluator_);
+
+  works_[communication] = postSingleCommunication(
+      communication,
+      communicator_->deviceId(),
+      expr_evaluator_.evaluate(communication->peer()).as<int64_t>(),
+      communicator_->getWorld(),
+      buffer,
+      expr_evaluator_.evaluate(communication->tag()).as<int64_t>());
+}
+
 void HostIrExecutor::handle(Wait* wait) {
   Expr* communication = wait->communication();
   NVF_ERROR(works_.find(communication) != works_.end(), "no wait req");

diff --git a/csrc/host_ir/executor.h b/csrc/host_ir/executor.h
@@ -77,6 +77,7 @@ class HostIrExecutor final : public OptInDispatch {
   void handle(SetCurrentStream* set_current_stream) override;
   void handle(PostOnStream* post_ir) override;
   void handle(Communication* communication) override;
+  void handle(P2PCommunication* communication) override;
   void handle(Wait* wait) override;
   void handle(ForLoop* for_loop) override;
   void handle(SliceOp* slice_op) override;
@@ -93,7 +94,7 @@ class HostIrExecutor final : public OptInDispatch {
   std::unordered_map<HostUnit*, FusionExecutorCache> fec_;
   using StreamKey = std::variant<int64_t, Stream*>;
   std::unordered_map<StreamKey, c10::cuda::CUDAStream> streams_;
-  std::unordered_map<Communication*, c10::intrusive_ptr<c10d::Work>> works_;
+  std::unordered_map<Expr*, c10::intrusive_ptr<c10d::Work>> works_;
 };
 
 } // namespace hir