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"