Added Nccl collectives helpers

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

@@ -0,0 +1,50 @@
+#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,
+                                ncclRedOp_t reduction_op,
+                                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,
+                                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,36 @@
+#include "internal/realm_test_utils.h"
+#include "realm-execution/realm_manager.h"
+#include "realm-execution/tasks/impl/nccl_task.h"
+#include <doctest/doctest.h>
+
+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();
+  }
+}
+
+} // namespace test

lib/utils/test/src/utils/containers/concat_vectors.cc

@@ -0,0 +1,38 @@
+#include "utils/containers/concat_vectors.h"
+#include "test/utils/doctest/fmt/vector.h"
+#include <doctest/doctest.h>
+#include <vector>
+
+using namespace ::FlexFlow;
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("concat_vectors") {
+    SUBCASE("two vectors") {
+      std::vector<int> prefix = {1, 2};
+      std::vector<int> postfix = {3, 4};
+
+      std::vector<int> result = concat_vectors(prefix, postfix);
+      std::vector<int> correct = {1, 2, 3, 4};
+
+      CHECK(result == correct);
+    }
+
+    SUBCASE("vector of vectors") {
+      std::vector<std::vector<int>> vecs = {{1, 2}, {3}, {4, 5}};
+
+      std::vector<int> result = concat_vectors(vecs);
+      std::vector<int> correct = {1, 2, 3, 4, 5};
+
+      CHECK(result == correct);
+    }
+
+    SUBCASE("empty vector of vectors") {
+      std::vector<std::vector<int>> vecs = {};
+
+      std::vector<int> result = concat_vectors(vecs);
+      std::vector<int> correct = {};
+
+      CHECK(result == correct);
+    }
+  }
+}

lib/utils/test/src/utils/containers/extend_vector.cc

@@ -0,0 +1,30 @@
+#include "utils/containers/extend_vector.h"
+#include "test/utils/doctest/fmt/vector.h"
+#include <doctest/doctest.h>
+#include <vector>
+
+using namespace ::FlexFlow;
+// checks rhs gets added to the end of lhs
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("extend_vector") {
+    SUBCASE("non-empty vectors") {
+      std::vector<int> lhs = {1, 2};
+      std::vector<int> rhs = {3, 4};
+
+      extend_vector(lhs, rhs);
+      std::vector<int> correct = {1, 2, 3, 4};
+
+      CHECK(lhs == correct);
+    }
+    // checks that lhs stays the same when the rhs is empty
+    SUBCASE("empty rhs") {
+      std::vector<int> lhs = {1, 2};
+      std::vector<int> rhs = {};
+
+      extend_vector(lhs, rhs);
+      std::vector<int> correct = {1, 2};
+
+      CHECK(lhs == correct);
+    }
+  }
+}

lib/utils/test/src/utils/containers/generate_vector.cc

@@ -0,0 +1,31 @@
+#include "utils/containers/generate_vector.h"
+#include "test/utils/doctest/fmt/vector.h"
+#include "utils/exception.h"
+#include "utils/nonnegative_int/nonnegative_int.h"
+#include <doctest/doctest.h>
+#include <vector>
+
+using namespace ::FlexFlow;
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("generate_vector") {
+    SUBCASE("empty vector") {
+      std::vector<int> result = generate_vector(0_n, [](nonnegative_int) -> int {
+        PANIC("lambda should not be called");
+      });
+      std::vector<int> correct = {};
+
+      CHECK(result == correct);
+    }
+    SUBCASE("non-empty vector") {
+      std::vector<int> result = generate_vector(5_n, [](nonnegative_int idx) {
+        int i = idx.unwrap_nonnegative();
+        return i * i;
+
+      });
+      std::vector<int> correct = {0, 1, 4, 9, 16};
+
+      CHECK(result == correct);
+    }
+  }
+}