diff --git a/lib/realm-execution/include/realm-execution/tasks/impl/nccl_task.h b/lib/realm-execution/include/realm-execution/tasks/impl/nccl_task.h new file mode 100644 index 0000000000..9e3601520d --- /dev/null +++ b/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 +#include "realm-execution/realm.h" +#include "realm-execution/realm_context.h" +#include +#include + +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 diff --git a/lib/realm-execution/include/realm-execution/tasks/impl/nccl_task_args.dtg.toml b/lib/realm-execution/include/realm-execution/tasks/impl/nccl_task_args.dtg.toml new file mode 100644 index 0000000000..fa1d7a3c08 --- /dev/null +++ b/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" diff --git a/lib/realm-execution/include/realm-execution/tasks/impl/serializable_nccl_task_args.dtg.toml b/lib/realm-execution/include/realm-execution/tasks/impl/serializable_nccl_task_args.dtg.toml new file mode 100644 index 0000000000..b77585ee7c --- /dev/null +++ b/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" diff --git a/lib/realm-execution/include/realm-execution/tasks/impl/serializable_nccl_task_args.h b/lib/realm-execution/include/realm-execution/tasks/impl/serializable_nccl_task_args.h new file mode 100644 index 0000000000..9b90be11c2 --- /dev/null +++ b/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 diff --git a/lib/realm-execution/include/realm-execution/tasks/task_id_t.dtg.toml b/lib/realm-execution/include/realm-execution/tasks/task_id_t.dtg.toml index a7c2d54180..10961065da 100644 --- a/lib/realm-execution/include/realm-execution/tasks/task_id_t.dtg.toml +++ b/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" diff --git a/lib/realm-execution/src/realm-execution/tasks/impl/nccl_task.cc b/lib/realm-execution/src/realm-execution/tasks/impl/nccl_task.cc new file mode 100644 index 0000000000..070ac3c937 --- /dev/null +++ b/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 +#include + +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(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 diff --git a/lib/realm-execution/src/realm-execution/tasks/impl/serializable_nccl_task_args.cc b/lib/realm-execution/src/realm-execution/tasks/impl/serializable_nccl_task_args.cc new file mode 100644 index 0000000000..57f6cd1c1b --- /dev/null +++ b/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 diff --git a/lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc b/lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc index 406a2912a3..0fa6769eca 100644 --- a/lib/realm-execution/src/realm-execution/tasks/realm_task_registry.cc +++ b/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)); diff --git a/lib/realm-execution/test/src/realm-execution/nccl_task.cc b/lib/realm-execution/test/src/realm-execution/nccl_task.cc new file mode 100644 index 0000000000..79a34fdac1 --- /dev/null +++ b/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 + +namespace test { + +using namespace ::FlexFlow; +namespace Realm = ::FlexFlow::Realm; + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("NCCL task prints Hello World") { + std::vector 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 diff --git a/lib/utils/test/src/utils/containers/concat_vectors.cc b/lib/utils/test/src/utils/containers/concat_vectors.cc new file mode 100644 index 0000000000..e22b8e8c6c --- /dev/null +++ b/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 +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("concat_vectors") { + SUBCASE("two vectors") { + std::vector prefix = {1, 2}; + std::vector postfix = {3, 4}; + + std::vector result = concat_vectors(prefix, postfix); + std::vector correct = {1, 2, 3, 4}; + + CHECK(result == correct); + } + + SUBCASE("vector of vectors") { + std::vector> vecs = {{1, 2}, {3}, {4, 5}}; + + std::vector result = concat_vectors(vecs); + std::vector correct = {1, 2, 3, 4, 5}; + + CHECK(result == correct); + } + + SUBCASE("empty vector of vectors") { + std::vector> vecs = {}; + + std::vector result = concat_vectors(vecs); + std::vector correct = {}; + + CHECK(result == correct); + } + } +} diff --git a/lib/utils/test/src/utils/containers/extend_vector.cc b/lib/utils/test/src/utils/containers/extend_vector.cc new file mode 100644 index 0000000000..a394cdf22f --- /dev/null +++ b/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 +#include + +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 lhs = {1, 2}; + std::vector rhs = {3, 4}; + + extend_vector(lhs, rhs); + std::vector correct = {1, 2, 3, 4}; + + CHECK(lhs == correct); + } + // checks that lhs stays the same when the rhs is empty + SUBCASE("empty rhs") { + std::vector lhs = {1, 2}; + std::vector rhs = {}; + + extend_vector(lhs, rhs); + std::vector correct = {1, 2}; + + CHECK(lhs == correct); + } + } +} diff --git a/lib/utils/test/src/utils/containers/generate_vector.cc b/lib/utils/test/src/utils/containers/generate_vector.cc new file mode 100644 index 0000000000..e560ba80bd --- /dev/null +++ b/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 +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("generate_vector") { + SUBCASE("empty vector") { + std::vector result = generate_vector(0_n, [](nonnegative_int) -> int { + PANIC("lambda should not be called"); + }); + std::vector correct = {}; + + CHECK(result == correct); + } + SUBCASE("non-empty vector") { + std::vector result = generate_vector(5_n, [](nonnegative_int idx) { + int i = idx.unwrap_nonnegative(); + return i * i; + + }); + std::vector correct = {0, 1, 4, 9, 16}; + + CHECK(result == correct); + } + } +}