Added Nccl collectives helpers

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

@@ -0,0 +1,51 @@
+#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 "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,
+    std::string const &message,
+    Realm::Event precondition);
+
+}
+
+#endif

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

@@ -0,0 +1,12 @@
+namespace = "FlexFlow"
+name = "NcclTaskArgs"
+type = "struct"
+features = []
+
+includes = [
+  "string",
+]
+
+[[fields]]
+name = "message"
+type = "std::string"

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

@@ -0,0 +1,13 @@
+namespace = "FlexFlow"
+name = "SerializableNcclTaskArgs"
+type = "struct"
+features = [
+  "json",
+]
+includes = [
+  "string",
+]
+
+[[fields]]
+name = "message"
+type = "std::string"

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,106 @@
+#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("%s\n", task_args.message.c_str());
+  std::printf("NCCL version: %d\n", nccl_version);
+}
+
+Realm::Event spawn_nccl_task(RealmContext &ctx,
+                             Realm::Processor target_proc,
+                             std::string const &message,
+                             Realm::Event precondition) {
+  NcclTaskArgs task_args = NcclTaskArgs{
+      /*message=*/message,
+  };
+
+  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,19 @@
+#include "realm-execution/tasks/impl/serializable_nccl_task_args.h"
+
+namespace FlexFlow {
+
+SerializableNcclTaskArgs
+    nccl_task_args_to_serializable(NcclTaskArgs const &args) {
+  return SerializableNcclTaskArgs{
+      args.message,
+  };
+}
+
+NcclTaskArgs
+    nccl_task_args_from_serializable(SerializableNcclTaskArgs const &args) {
+  return NcclTaskArgs{
+      args.message,
+  };
+}
+
+} // 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,11 @@ Realm::Event register_all_tasks() {
         register_task(Realm::Processor::TOC_PROC, task_id, op_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));

lib/realm-execution/test/src/realm-execution/nccl_task.cc

@@ -0,0 +1,124 @@
+#include "internal/realm_test_utils.h"
+#include "realm-execution/realm_manager.h"
+#include "realm-execution/tasks/impl/nccl_task.h"
+
+#include <cuda_runtime.h>
+#include <doctest/doctest.h>
+#include <nccl.h>
+#include <vector>
+
+namespace test {
+
+using namespace ::FlexFlow;
+namespace Realm = ::FlexFlow::Realm;
+
+TEST_SUITE(FF_CUDA_TEST_SUITE) {
+  TEST_CASE("NCCL task prints Hello World") {
+    std::vector<char *> fake_args =
+        make_fake_realm_args(/*num_cpus=*/1_p, /*num_gpus=*/1_n);
+
+    int fake_argc = fake_args.size();
+    char **fake_argv = fake_args.data();
+
+    RealmManager manager(&fake_argc, &fake_argv);
+
+    ControllerTaskResult result =
+        manager.start_controller([](RealmContext &ctx) {
+          Realm::Event event = spawn_nccl_task(
+              ctx,
+              ctx.get_current_processor(),
+              "Hello World from NCCL!",
+              Realm::Event::NO_EVENT);
+
+          event.wait();
+        });
+
+    result.wait();
+  }
+
+  TEST_CASE("NCCL broadcast and reduce helpers") {
+    constexpr size_t count = 8;
+    size_t const buffer_size = count * sizeof(int);
+
+    ncclUniqueId unique_id;
+    REQUIRE(ncclGetUniqueId(&unique_id) == ncclSuccess);
+
+    ncclComm_t communicator;
+    REQUIRE(ncclCommInitRank(
+                &communicator,
+                /*num_ranks=*/1,
+                unique_id,
+                /*rank=*/0) == ncclSuccess);
+
+    ffStream_t stream;
+    REQUIRE(cudaStreamCreate(&stream) == cudaSuccess);
+
+    int *send_buffer = nullptr;
+    int *receive_buffer = nullptr;
+
+    REQUIRE(cudaMalloc(&send_buffer, buffer_size) == cudaSuccess);
+    REQUIRE(cudaMalloc(&receive_buffer, buffer_size) == cudaSuccess);
+
+    std::vector<int> input = {1, 2, 3, 4, 5, 6, 7, 8};
+    std::vector<int> output(count, 0);
+
+    REQUIRE(cudaMemcpy(send_buffer,
+                       input.data(),
+                       buffer_size,
+                       cudaMemcpyHostToDevice) == cudaSuccess);
+
+    SUBCASE("broadcast") {
+      REQUIRE(cudaMemset(receive_buffer, 0, buffer_size) == cudaSuccess);
+
+      REQUIRE(run_nccl_broadcast(send_buffer,
+                                 receive_buffer,
+                                 count,
+                                 ncclInt32,
+                                 /*root_rank=*/0,
+                                 communicator,
+                                 stream) == ncclSuccess);
+
+      REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
+
+      REQUIRE(cudaMemcpy(output.data(),
+                         receive_buffer,
+                         buffer_size,
+                         cudaMemcpyDeviceToHost) == cudaSuccess);
+
+      for (size_t i = 0; i < count; i++) {
+        CHECK(output[i] == input[i]);
+      }
+    }
+
+    SUBCASE("reduce") {
+      REQUIRE(cudaMemset(receive_buffer, 0, buffer_size) == cudaSuccess);
+
+      REQUIRE(run_nccl_reduce(send_buffer,
+                              receive_buffer,
+                              count,
+                              ncclInt32,
+                              ncclSum,
+                              /*root_rank=*/0,
+                              communicator,
+                              stream) == ncclSuccess);
+
+      REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
+
+      REQUIRE(cudaMemcpy(output.data(),
+                         receive_buffer,
+                         buffer_size,
+                         cudaMemcpyDeviceToHost) == cudaSuccess);
+
+      for (size_t i = 0; i < count; i++) {
+        CHECK(output[i] == input[i]);
+      }
+    }
+
+    REQUIRE(cudaFree(send_buffer) == cudaSuccess);
+    REQUIRE(cudaFree(receive_buffer) == cudaSuccess);
+    REQUIRE(cudaStreamDestroy(stream) == cudaSuccess);
+    REQUIRE(ncclCommDestroy(communicator) == ncclSuccess);
+  }
+}
+
+} // namespace test