Added Nccl collectives helpers

lib/realm-execution/include/realm-execution/tasks/impl/nccl_task.h

@@ -0,0 +1,57 @@
+#ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_TASKS_IMPL_NCCL_TASK_H
+#define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_TASKS_IMPL_NCCL_TASK_H
+
+#include "realm-execution/device_specific_managed_per_device_ff_handle.h"
+#include "realm-execution/tensor_instance_backing.dtg.h"
+#include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h"
+#include "kernels/device.h"
+#include <nccl.h>
+#include "realm-execution/realm.h"
+#include "realm-execution/realm_context.h"
+#include <string>
+#include <cstddef>
+
+namespace FlexFlow {
+
+ncclResult_t run_nccl_all_reduce(void const *send_buffer,
+                                void *receive_buffer,
+                                size_t count,
+                                ncclDataType_t data_type,
+                                ncclRedOp_t reduction_op,
+                                ncclComm_t communicator,
+                                ffStream_t stream);
+
+ncclResult_t run_nccl_broadcast(void const *send_buffer,
+                                void *receive_buffer,
+                                size_t count,
+                                ncclDataType_t data_type,
+                                int root_rank,
+                                ncclComm_t communicator,
+                                ffStream_t stream);
+
+ncclResult_t run_nccl_reduce(void const *send_buffer,
+                             void *receive_buffer,
+                             size_t count,
+                             ncclDataType_t data_type,
+                             ncclRedOp_t reduction_op,
+                             int root_rank,
+                             ncclComm_t communicator,
+                             ffStream_t stream);
+
+void nccl_task_body(void const *args,
+    size_t arglen,
+    void const *userdata,
+    size_t userdata_len,
+    Realm::Processor proc);
+
+Realm::Event spawn_nccl_task(
+    RealmContext &ctx,
+    Realm::Processor target_proc,
+    DynamicNodeInvocation const &invocation,
+    TensorInstanceBacking const &tensor_backing,
+    DeviceSpecificPtr<ManagedPerDeviceFFHandle> const &device_handle,
+    Realm::Event precondition);
+
+}
+
+#endif

lib/realm-execution/include/realm-execution/tasks/impl/nccl_task_args.dtg.toml

@@ -0,0 +1,23 @@
+namespace = "FlexFlow"
+name = "NCCLTaskArgs"
+type = "struct"
+features = []
+
+includes = [
+  "realm-execution/device_specific_managed_per_device_ff_handle.h",
+  "realm-execution/device_specific_ptr.h",
+  "realm-execution/tensor_instance_backing.dtg.h",
+  "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h",
+]
+
+[[fields]]
+name = "invocation"
+type = "::FlexFlow::DynamicNodeInvocation"
+
+[[fields]]
+name = "tensor_backing"
+type = "::FlexFlow::TensorInstanceBacking"
+
+[[fields]]
+name = "device_handle"
+type = "::FlexFlow::DeviceSpecificPtr<::FlexFlow::ManagedPerDeviceFFHandle>"

lib/realm-execution/include/realm-execution/tasks/impl/serializable_nccl_task_args.dtg.toml

@@ -0,0 +1,24 @@
+namespace = "FlexFlow"
+name = "SerializableNcclTaskArgs"
+type = "struct"
+features = [
+  "json",
+]
+
+includes = [
+  "realm-execution/tasks/serializer/serializable_device_specific_ptr.dtg.h",
+  "realm-execution/tasks/serializer/serializable_tensor_instance_backing.dtg.h",
+  "task-spec/dynamic_graph/serializable_dynamic_node_invocation.dtg.h",
+]
+
+[[fields]]
+name = "invocation"
+type = "::FlexFlow::SerializableDynamicNodeInvocation"
+
+[[fields]]
+name = "tensor_backing"
+type = "::FlexFlow::SerializableTensorInstanceBacking"
+
+[[fields]]
+name = "device_handle"
+type = "::FlexFlow::SerializableDeviceSpecificPtr"

lib/realm-execution/include/realm-execution/tasks/impl/serializable_nccl_task_args.h

@@ -0,0 +1,17 @@
+#ifndef _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_TASKS_IMPL_SERIALIZABLE_NCCL_TASK_ARGS_H
+#define _FLEXFLOW_LIB_REALM_EXECUTION_INCLUDE_REALM_EXECUTION_TASKS_IMPL_SERIALIZABLE_NCCL_TASK_ARGS_H
+
+#include "realm-execution/tasks/impl/nccl_task_args.dtg.h"
+#include "realm-execution/tasks/impl/serializable_nccl_task_args.dtg.h"
+
+namespace FlexFlow {
+
+SerializableNcclTaskArgs
+    nccl_task_args_to_serializable(NCCLTaskArgs const &);
+
+NCCLTaskArgs
+    nccl_task_args_from_serializable(SerializableNcclTaskArgs const &);
+
+} // namespace FlexFlow
+
+#endif

lib/realm-execution/include/realm-execution/tasks/task_id_t.dtg.toml

@@ -270,6 +270,9 @@ name = "NCCL_GETUNIQUEID_TASK_ID"
 [[values]]
 name = "NCCL_INIT_COMMS_TASK_ID"
 
+[[values]]
+name = "NCCL_HELLO_WORLD_TASK_ID"
+
 [[values]]
 name = "STRATEGY_SEARCH_TASK_ID"
 

lib/realm-execution/src/realm-execution/tasks/impl/nccl_task.cc

@@ -0,0 +1,110 @@
+#include "realm-execution/tasks/impl/nccl_task.h"
+#include "realm-execution/tasks/impl/nccl_task_args.dtg.h"
+#include "realm-execution/tasks/impl/serializable_nccl_task_args.h"
+#include "realm-execution/tasks/serializer/task_arg_serializer.h"
+#include "realm-execution/tasks/task_id_t.h"
+
+#include <cstdio>
+#include <nccl.h>
+
+namespace FlexFlow {
+
+ncclResult_t run_nccl_all_reduce(void const *send_buffer,
+                                void *receive_buffer,
+                                size_t count,
+                                ncclDataType_t data_type,
+                                ncclRedOp_t reduction_op,
+                                ncclComm_t communicator,
+                                ffStream_t stream) {
+  return ncclAllReduce(send_buffer,
+                       receive_buffer,
+                       count,
+                       data_type,
+                       reduction_op,
+                       communicator,
+                       stream);
+}
+
+ncclResult_t run_nccl_broadcast(void const *send_buffer,
+                                void *receive_buffer,
+                                size_t count,
+                                ncclDataType_t data_type,
+                                int root_rank,
+                                ncclComm_t communicator,
+                                ffStream_t stream) {
+  return ncclBroadcast(send_buffer,
+                       receive_buffer,
+                       count,
+                       data_type,
+                       root_rank,
+                       communicator,
+                       stream);
+}
+
+ncclResult_t run_nccl_reduce(void const *send_buffer,
+                             void *receive_buffer,
+                             size_t count,
+                             ncclDataType_t data_type,
+                             ncclRedOp_t reduction_op,
+                             int root_rank,
+                             ncclComm_t communicator,
+                             ffStream_t stream) {
+  return ncclReduce(send_buffer,
+                    receive_buffer,
+                    count,
+                    data_type,
+                    reduction_op,
+                    root_rank,
+                    communicator,
+                    stream);
+}
+
+void nccl_task_body(void const *args,
+                    size_t arglen,
+                    void const *userdata,
+                    size_t userdata_len,
+                    Realm::Processor proc) {
+  (void)userdata;
+  (void)userdata_len;
+  (void)proc;
+
+  NCCLTaskArgs task_args = nccl_task_args_from_serializable(
+      deserialize_task_args<SerializableNcclTaskArgs>(args, arglen));
+
+  int nccl_version = 0;
+  ncclResult_t result = ncclGetVersion(&nccl_version);
+
+  if (result != ncclSuccess) {
+    std::printf("NCCL error: %s\n", ncclGetErrorString(result));
+    return;
+  }
+
+  std::printf("NCCL version: %d\n", nccl_version);
+}
+
+Realm::Event spawn_nccl_task(
+    RealmContext &ctx,
+    Realm::Processor target_proc,
+    DynamicNodeInvocation const &invocation,
+    TensorInstanceBacking const &tensor_backing,
+    DeviceSpecificPtr<ManagedPerDeviceFFHandle> const &device_handle,
+    Realm::Event precondition) {
+  NCCLTaskArgs task_args = NCCLTaskArgs{
+      /*invocation=*/invocation,
+      /*tensor_backing=*/tensor_backing,
+      /*device_handle=*/device_handle,
+  };
+
+  std::string serialized_args =
+      serialize_task_args(nccl_task_args_to_serializable(task_args));
+
+  return ctx.spawn_task(
+      target_proc,
+      task_id_t::NCCL_HELLO_WORLD_TASK_ID,
+      serialized_args.data(),
+      serialized_args.size(),
+      Realm::ProfilingRequestSet{},
+      precondition);
+}
+
+} // namespace FlexFlow

lib/realm-execution/src/realm-execution/tasks/impl/serializable_nccl_task_args.cc

@@ -0,0 +1,33 @@
+#include "realm-execution/tasks/impl/serializable_nccl_task_args.h"
+#include "realm-execution/tasks/serializer/serializable_device_specific_ptr.h"
+#include "realm-execution/tasks/serializer/serializable_tensor_instance_backing.h"
+#include "task-spec/dynamic_graph/serializable_dynamic_node_invocation.h"
+
+namespace FlexFlow {
+
+SerializableNcclTaskArgs
+nccl_task_args_to_serializable(NCCLTaskArgs const &args) {
+  return SerializableNcclTaskArgs{
+      /*invocation=*/
+      dynamic_node_invocation_to_serializable(args.invocation),
+      /*tensor_backing=*/
+      tensor_instance_backing_to_serializable(args.tensor_backing),
+      /*device_handle=*/
+      device_specific_ptr_to_serializable(args.device_handle),
+  };
+}
+
+NCCLTaskArgs
+nccl_task_args_from_serializable(SerializableNcclTaskArgs const &args) {
+  return NCCLTaskArgs{
+      /*invocation=*/
+      dynamic_node_invocation_from_serializable(args.invocation),
+      /*tensor_backing=*/
+      tensor_instance_backing_from_serializable(args.tensor_backing),
+      /*device_handle=*/
+      device_specific_ptr_from_serializable<ManagedPerDeviceFFHandle>(
+          args.device_handle),
+  };
+}
+
+} // namespace FlexFlow

lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc

@@ -7,6 +7,7 @@
 #include "realm-execution/tasks/impl/per_device_op_state_init_task.h"
 #include "realm-execution/tasks/task_id_t.h"
 #include "utils/exception.h"
+#include "realm-execution/tasks/impl/nccl_task.h"
 
 namespace FlexFlow {
 
@@ -133,6 +134,16 @@ Realm::Event register_all_tasks() {
         register_task(Realm::Processor::TOC_PROC, task_id, op_task_body));
   }
 
+  pending_registrations.push_back(
+      register_task(Realm::Processor::LOC_PROC,
+                    task_id_t::NCCL_HELLO_WORLD_TASK_ID,
+                    nccl_task_body));
+
+  pending_registrations.push_back(
+      register_task(Realm::Processor::TOC_PROC,
+                    task_id_t::NCCL_HELLO_WORLD_TASK_ID,
+                    nccl_task_body));
+
   pending_registrations.push_back(register_task(Realm::Processor::LOC_PROC,
                                                 task_id_t::CONTROLLER_TASK_ID,
                                                 controller_task_body));