diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 5493e1e..6b94aa0 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -63,6 +63,7 @@ Please check all the platforms and/or backends this PR affects (i.e., code is to - [ ] OpenMPI - [ ] MPICH +- [ ] NCCL ## Performance Impact @@ -110,6 +111,7 @@ See `CONTRIBUTING.md` ยง Pull Requests for the official testing requirements and - [ ] OpenMPI - [ ] MPICH +- [ ] NCCL --- diff --git a/README.md b/README.md index 204f7b5..7800ae5 100644 --- a/README.md +++ b/README.md @@ -352,6 +352,7 @@ export LD_LIBRARY_PATH=${INFINI_INSTALL}/lib:$LD_LIBRARY_PATH |---------|---------------|----------------------|---------------| | **OpenMPI** | Full | `WITH_OMPI=ON` | The default backend. Requires the OpenMPI development package.| | **MPICH** | Full | `WITH_MPICH=ON` | Requires the MPICH development package.| +| **NCCL** | Partial | `WITH_NCCL=ON` | Requires NVIDIA NCCL. Currently available only when `WITH_NVIDIA=ON`.| diff --git a/examples/ccl/all_reduce.cc b/examples/ccl/all_reduce.cc new file mode 100644 index 0000000..6f8dc35 --- /dev/null +++ b/examples/ccl/all_reduce.cc @@ -0,0 +1,190 @@ +/** + * InfiniCCL Example: Thread-per-GPU Single-Node AllReduce + * + * This example demonstrates spawning one CPU thread per GPU on a single node, + * generating a shared UniqueID, and performing an AllReduce entirely via + * InfiniCCL's native CCL backend without any OpenMPI dependencies. + * + * Note: to properly run this example, you should specify `--launcher none` + * since it requires no MPI process spawning. + */ + +#include + +#include +#include +#include +#include + +// Public API +#include "infiniccl.h" + +// Example-Specific Utilities +#include "utils.h" + +// Internal Headers (Accessible via example-specific include paths, technically +// not public APIs) +#include "backend_manifest.h" + +using namespace infini::ccl; + +// Structure to pass execution data to each GPU worker thread. +struct ThreadArgs { + int rank; + int size; + infinicclUniqueId id; + size_t num_elements; + int warmup_iter; + int profile_iter; +}; + +// Worker function executed by each CPU thread. +void WorkerThread(ThreadArgs args) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + // Bind this specific CPU thread to its designated local GPU device + // In a thread-per-GPU model, `local_rank` is exactly the thread's rank ID. + int local_device_id = args.rank; + CHECK_RT(Rt, Rt::SetDevice(local_device_id)); + + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitRank(&comm, args.size, args.id, args.rank)); + + // Prepare Host Data Structures + std::vector h_send(args.num_elements); + std::vector h_recv(args.num_elements, 0.0f); + + // Each rank provides its own (rank + 1) as data. + for (size_t i = 0; i < args.num_elements; i++) { + h_send[i] = static_cast(args.rank + 1); + } + + // Allocate GPU Memory using InfiniCCL's Runtime abstraction layer. + float *d_send = nullptr; + float *d_recv = nullptr; + size_t total_bytes = args.num_elements * sizeof(float); + + CHECK_RT(Rt, Rt::Malloc((void **)&d_send, total_bytes)); + CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, total_bytes)); + + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), total_bytes, + Rt::MemcpyHostToDevice)); + + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + // Warm-up Iterations + for (int i = 0; i < args.warmup_iter; ++i) { + CHECK_INFINI(infinicclAllReduce(d_send, d_recv, args.num_elements, + infinicclFloat32, infinicclSum, comm, + nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + // Profiling Iterations + Timer timer; + for (int i = 0; i < args.profile_iter; i++) { + CHECK_INFINI(infinicclAllReduce(d_send, d_recv, args.num_elements, + infinicclFloat32, infinicclSum, comm, + nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + double elapsed = timer.ElapsedMs() / static_cast(args.profile_iter); + + // Copy output reduced data back from device to host. + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, total_bytes, + Rt::MemcpyDeviceToHost)); + + // Result Validation + float expected = 0.0f; + for (int r = 0; r < args.size; r++) { + expected += static_cast(r + 1); + } + Validator::ValidateResult(h_recv.data(), args.num_elements, expected, + args.rank, true, "AllReduce"); + + // Metrics Reporting (Only Rank 0) + if (args.rank == 0) { + std::cout << "\n=== Single-Node Threaded AllReduce Results ===" + << std::endl; + std::cout << "Data size: " << args.num_elements << " floats (" + << total_bytes / 1024 / 1024 << " MB)" << std::endl; + Metrics metrics{elapsed, total_bytes, args.size}; + metrics.Print(); + } + + // Cleanup local rank resources. + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + CHECK_INFINI(infinicclCommDestroy(comm)); +} + +int main(int argc, char **argv) { + int num_gpus = 8; + int warmup_iters = 1; + int profile_iters = 20; + size_t num_elements = 1 << 25; + + int opt; + while ((opt = getopt(argc, argv, "g:w:p:n:h")) != -1) { + switch (opt) { + case 'g': + num_gpus = std::stoi(optarg); + break; + case 'w': + warmup_iters = std::stoi(optarg); + break; + case 'p': + profile_iters = std::stoi(optarg); + break; + case 'n': + num_elements = static_cast(std::stoull(optarg)); + break; + case 'h': + std::cout << "Usage: " << argv[0] << " [options]\n" + << "Options:\n" + << " -g Number of GPUs (default: 8)\n" + << " -w Warmup iterations (default: 1)\n" + << " -p Profile iterations (default: 20)\n" + << " -n Number of elements (default: " + << (1 << 25) << ")\n"; + return EXIT_SUCCESS; + default: + std::cerr << "Invalid argument. Use -h for help." << std::endl; + return EXIT_FAILURE; + } + } + + char hostname[256]; + gethostname(hostname, sizeof(hostname)); + std::cout << "[Main Process] Host: " << hostname + << " | Target GPUs: " << num_gpus << std::endl; + + infinicclUniqueId shared_id; + CHECK_INFINI(infinicclGetUniqueId(&shared_id)); + + // Spawn CPU thread pool. + std::vector threads; + threads.reserve(num_gpus); + + for (int rank = 0; rank < num_gpus; ++rank) { + ThreadArgs args{rank, num_gpus, shared_id, + num_elements, warmup_iters, profile_iters}; + threads.emplace_back(WorkerThread, args); + } + + // Await execution completion across all threads. + for (auto &t : threads) { + if (t.joinable()) { + t.join(); + } + } + + std::cout + << "[Main Process] All worker threads joined. InfiniCCL finalized safely." + << std::endl; + return EXIT_SUCCESS; +} diff --git a/examples/ccl_mpi_hybrid/all_reduce.cc b/examples/ccl_mpi_hybrid/all_reduce.cc new file mode 100644 index 0000000..b1d7562 --- /dev/null +++ b/examples/ccl_mpi_hybrid/all_reduce.cc @@ -0,0 +1,161 @@ +/** + * InfiniCCL Example: AllReduce (Ompi + CCL Hybrid) + * * This example demonstrates a program that performs a + * collective sum-reduction across multiple GPUs and nodes. + */ + +#include + +#include +#include + +// Public API +#include "infiniccl.h" + +// Example-Specific Utilities +#include "utils.h" + +// Internal Headers (Accessible via example-specific include paths, technically +// not public APIs) +#include "backend_manifest.h" +#include "device.h" +#include "runtime.h" +#include "traits.h" + +using namespace infini::ccl; + +void RunAllReduceExample(int argc, char **argv, int warmup_iter, + int profile_iter, const size_t kNumElements) { + constexpr Device::Type kDevType = + ListGetBest(EnabledDevices{}); + using Rt = Runtime; + + CHECK_INFINI(infinicclInit(&argc, &argv)); + + int rank, size; + CHECK_INFINI(infinicclGetRank(&rank)); + CHECK_INFINI(infinicclGetSize(&size)); + + char hostname[256]; + gethostname(hostname, sizeof(hostname)); + + // Map local rank to GPU device. + // Note: this is just for info printing. In practice, this part is not needed. + const char *local_rank_str = std::getenv("OMPI_COMM_WORLD_LOCAL_RANK"); + int local_rank = 0; + if (local_rank_str != nullptr) { + local_rank = std::atoi(local_rank_str); + } + + CHECK_RT(Rt, Rt::SetDevice(local_rank)); + + // Setup Communicator + infinicclUniqueId id; + infinicclComm_t comm = nullptr; + CHECK_INFINI(infinicclCommInitAll(&comm, size, nullptr)); + + if (rank == 0) { + CHECK_INFINI(infinicclGetUniqueId(&id)); + } + + CHECK_INFINI(infinicclBroadcast(&id, &id, sizeof(id), infinicclChar, 0, comm, + nullptr)); + + std::cout << "[Rank " << rank << "] Host: " << hostname + << " | GPU: " << Device::StringFromType(kDevType) << " " + << " | Device " << local_rank << std::endl; + + CHECK_INFINI(infinicclCommInitRank(&comm, size, id, rank)); + + // Prepare Data + std::vector h_send(kNumElements); + std::vector h_recv(kNumElements, 0.0f); + + // Initialize: each rank provides its (rank + 1) as data. + for (size_t i = 0; i < kNumElements; i++) { + h_send[i] = static_cast(rank + 1); + } + + float *d_send, *d_recv; + size_t total_bytes = kNumElements * sizeof(*d_send); + CHECK_RT(Rt, Rt::Malloc((void **)&d_send, total_bytes)); + CHECK_RT(Rt, Rt::Malloc((void **)&d_recv, total_bytes)); + CHECK_RT(Rt, Rt::Memcpy(d_send, h_send.data(), total_bytes, + Rt::MemcpyHostToDevice)); + CHECK_RT(Rt, Rt::Memcpy(d_recv, h_recv.data(), total_bytes, + Rt::MemcpyHostToDevice)); + + if (rank == 0) { + std::cout << "\n=== Performing AllReduce on GPU Memory ===" << std::endl; + std::cout << "Data size: " << kNumElements << " floats (" + << total_bytes / 1024 / 1024 << " MB)" << std::endl; + std::cout << "Operation: Sum" << std::endl; + std::cout << "Warm-up iterations: " << warmup_iter << std::endl; + std::cout << "Profile iterations: " << profile_iter << std::endl; + } + + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + // Warm-up and D2H transfer the answer. + CHECK_INFINI(infinicclAllReduce(d_send, d_recv, kNumElements, + infinicclFloat32, infinicclSum, comm, + nullptr)); + CHECK_RT(Rt, Rt::Memcpy(h_recv.data(), d_recv, kNumElements * sizeof(float), + Rt::MemcpyDeviceToHost)); + + for (int i = 1; i < warmup_iter; ++i) { + CHECK_INFINI(infinicclAllReduce(d_send, d_recv, kNumElements, + infinicclFloat32, infinicclSum, comm, + nullptr)); + } + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + + // Profiling + Timer timer; + + for (int i = 0; i < profile_iter; i++) { + CHECK_INFINI(infinicclAllReduce(d_send, d_recv, kNumElements, + infinicclFloat32, infinicclSum, comm, + nullptr)); + } + + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + double elapsed = timer.ElapsedMs() / static_cast(profile_iter); + + // Result Validation + float expected = 0.0f; + for (int r = 0; r < size; r++) { + expected += static_cast(r + 1); + } + + CHECK_RT(Rt, Rt::StreamSynchronize(nullptr)); + Validator::ValidateResult(h_recv.data(), kNumElements, expected, rank, true, + "AllReduce"); + + // Metrics Reporting (Only from rank 0 for cleaner output) + if (rank == 0) { + Metrics metrics{elapsed, total_bytes, size}; + metrics.Print(); + } + + // Cleanup + CHECK_RT(Rt, Rt::Free(d_send)); + CHECK_RT(Rt, Rt::Free(d_recv)); + + CHECK_INFINI(infinicclCommDestroy(comm)); + CHECK_INFINI(infinicclFinalize()); + + if (rank == 0) { + std::cout << "InfiniCCL finalized." << std::endl; + } +} + +int main(int argc, char **argv) { + int warmup_iters = 2; + int profile_iters = 20; + size_t num_elements = 1 << 25; + + RunAllReduceExample(argc, argv, warmup_iters, profile_iters, num_elements); + + return EXIT_SUCCESS; +} diff --git a/include/comm.h b/include/comm.h index 331a8c0..ec16c92 100644 --- a/include/comm.h +++ b/include/comm.h @@ -10,8 +10,14 @@ extern "C" { #endif +#define INFINICCL_UNIQUE_ID_BYTES 128 + typedef void *infinicclComm_t; +typedef struct { + char internal[INFINICCL_UNIQUE_ID_BYTES]; +} infinicclUniqueId; + // Initialization infinicclResult_t infinicclInit(int *argc, char ***argv); infinicclResult_t infinicclFinalize(void); @@ -21,8 +27,11 @@ infinicclResult_t infinicclGetRank(int *rank); infinicclResult_t infinicclGetSize(int *size); // Communicator Management +infinicclResult_t infinicclGetUniqueId(infinicclUniqueId *id); infinicclResult_t infinicclCommInitAll(infinicclComm_t *comm, int ndev, const int *devlist); +infinicclResult_t infinicclCommInitRank(infinicclComm_t *comm, int nranks, + infinicclUniqueId id, int rank); infinicclResult_t infinicclCommDestroy(infinicclComm_t comm); // --- Reduction Operations --- diff --git a/include/return_status.h b/include/return_status.h index 219d5f2..98cb928 100644 --- a/include/return_status.h +++ b/include/return_status.h @@ -14,7 +14,8 @@ typedef enum { infinicclInvalidUsage = 5, infinicclRemoteError = 6, infinicclInProgress = 7, - infinicclNumResults = 8 + infinicclNotSupported = 8, + infinicclNumResults = 9 } infinicclResult_t; #ifdef __cplusplus diff --git a/scripts/gen_bridge.py b/scripts/gen_bridge.py index 942be82..0cebe8f 100644 --- a/scripts/gen_bridge.py +++ b/scripts/gen_bridge.py @@ -27,7 +27,11 @@ } # Map logical backend names (from CMake) to their internal source paths. -BACKEND_PATH_MAP = {"ompi": "ompi/impl", "mpich": "ompi/impl", "nccl": "nvidia/nccl"} +BACKEND_PATH_MAP = { + "ompi": "ompi/impl", + "mpich": "ompi/impl", + "nccl": "nvidia/nccl/impl", +} # ================================================================= # LOGIC @@ -86,6 +90,7 @@ def generate(project_root, output_dir, devices, backends): sigs = parse_signatures(header_path) # 1. Generate `backend_manifest.h`. + implemented_ops = set() manifest_lines = [ AUTOGEN_HEADER, "#ifndef INFINI_CCL_BACKEND_MANIFEST_H_", @@ -126,6 +131,7 @@ def generate(project_root, output_dir, devices, backends): rel_path = f"{impl_subpath}/{op}.h" if os.path.exists(os.path.join(src_dir, rel_path)): manifest_lines.append(f'#include "{rel_path}"') + implemented_ops.add(op) # Add the Type Alias `EnabledDevices` manifest_lines.append("\nnamespace infini::ccl {\n") @@ -154,36 +160,36 @@ def generate(project_root, output_dir, devices, backends): ] for s in sigs: - # We need to transform the raw args to add casts for specific types. - args_with_casts = [] - # Split params to analyze types (simplified approach). - params_list = s["params"].split(",") - - for p in params_list: - p = p.strip() - if not p or p == "void": - continue - - # Get the type and the name. - parts = p.split() - arg_type = parts[0] - arg_name = parts[-1].replace("*", "") - - # Apply `static_cast` for specialized `Infini` types. - if arg_type == "infinicclDataType_t": - args_with_casts.append(f"static_cast({arg_name})") - elif arg_type == "infinicclRedOp_t": - args_with_casts.append(f"static_cast({arg_name})") - else: - args_with_casts.append(arg_name) - - casted_args_str = ", ".join(args_with_casts) - - bridge_lines.append( - f"\n{s['ret']} {s['name']}({s['params']}) {{\n" - f" return static_cast<{s['ret']}>(Operation<{s['key']}>::Call({casted_args_str}));\n" - f"}}" - ) + op_filename = re.sub( + r"^infiniccl_", "", re.sub(r"(?({arg_name})") + elif arg_type == "infinicclRedOp_t": + args_with_casts.append(f"static_cast({arg_name})") + else: + args_with_casts.append(arg_name) + + casted_args_str = ", ".join(args_with_casts) + body = f" return static_cast<{s['ret']}>(Operation<{s['key']}>::Call({casted_args_str}));" + else: + body = f" return static_cast<{s['ret']}>(ReturnStatus::kNotSupported);" + + bridge_lines.append(f"\n{s['ret']} {s['name']}({s['params']}) {{\n{body}\n}}") bridge_lines.append('\n} // extern "C"') bridge_lines.append("} // namespace infini::ccl\n") diff --git a/scripts/run_examples.py b/scripts/run_examples.py index d938caa..58fca09 100755 --- a/scripts/run_examples.py +++ b/scripts/run_examples.py @@ -90,6 +90,7 @@ def run_iccl_example( target_info: dict, config_path: str, launcher_opt: Optional[str], + executable_args: List[str], log_dir: str, trigger_build: bool, verbose: bool, @@ -117,6 +118,8 @@ def run_iccl_example( else: cmd.append(binary_path) + cmd.extend(executable_args) + # Force environment unbuffered stream states for sub-python instances. custom_env = os.environ.copy() custom_env["PYTHONUNBUFFERED"] = "1" @@ -231,8 +234,16 @@ def main(): action="store_true", help="Enable verbose mode: stream execution `stdout`/`stderr` directly to terminal while logging.", ) + parser.add_argument( + "executable_args", + nargs=argparse.REMAINDER, + help="Arguments forwarded to each selected executable. Use `--` before these args, e.g. `-- -g 4 -n 1024`.", + ) args = parser.parse_args() + executable_args = args.executable_args + if executable_args and executable_args[0] == "--": + executable_args = executable_args[1:] script_dir = Path(__file__).parent.resolve() examples_root = script_dir / "../examples" @@ -280,6 +291,7 @@ def main(): target_info=target, config_path=args.config, launcher_opt=args.launcher, + executable_args=executable_args, log_dir=current_run_log_dir, trigger_build=is_first_run, verbose=args.verbose, diff --git a/src/backend_device_map.h b/src/backend_device_map.h new file mode 100644 index 0000000..687e7da --- /dev/null +++ b/src/backend_device_map.h @@ -0,0 +1,22 @@ +#ifndef INFINI_CCL_BACKEND_DEVICE_MAP_H_ +#define INFINI_CCL_BACKEND_DEVICE_MAP_H_ + +#include "backend.h" +#include "device.h" + +namespace infini::ccl { + +template +struct IsSupportedCombination : std::false_type {}; + +// OpenMPI is device-agnostic. +template +struct IsSupportedCombination : std::true_type {}; + +template <> +struct IsSupportedCombination + : std::true_type {}; + +}; // namespace infini::ccl + +#endif // INFINI_CCL_BACKEND_DEVICE_MAP_H_ diff --git a/src/base/comm_destroy.h b/src/base/comm_destroy.h index 38481c2..454b131 100644 --- a/src/base/comm_destroy.h +++ b/src/base/comm_destroy.h @@ -25,7 +25,6 @@ class CommDestroy : public Operation { CommDestroyImpl::Apply(comm_handle); if (status == ReturnStatus::kSuccess) { - // Pair with the `new` in `CommInitAll`. delete static_cast(comm_handle); } diff --git a/src/base/comm_init_all.h b/src/base/comm_init_all.h index 519c3dc..e737306 100644 --- a/src/base/comm_init_all.h +++ b/src/base/comm_init_all.h @@ -17,16 +17,18 @@ class CommInitAll : public Operation { typename... Args> static ReturnStatus Execute(void **comm_handle, Args &&...args) { Communicator *&comm = *reinterpret_cast(comm_handle); - if (comm) { + if (comm && comm->inter_comm()) { // TODO(lzm): change to use `glog`. - LOG("Invalid communicator handle for CommInitAll."); + LOG("Invalid communicator handle for `CommInitAll`."); return ReturnStatus::kInvalidArgument; } constexpr Device::Type kDev = ListGetBest(ActiveDevices{}); - comm = new Communicator(kDev, 0); + if (!comm) { + comm = new Communicator(kDev, 0); + } return CommInitAllImpl::Apply( comm, std::forward(args)...); diff --git a/src/base/comm_init_rank.h b/src/base/comm_init_rank.h new file mode 100644 index 0000000..46f93f2 --- /dev/null +++ b/src/base/comm_init_rank.h @@ -0,0 +1,45 @@ +#ifndef INFINI_CCL_BASE_COMM_INIT_RANK_H_ +#define INFINI_CCL_BASE_COMM_INIT_RANK_H_ + +#include "logging.h" +#include "operation.h" +#include "return_status_impl.h" + +namespace infini::ccl { + +template +struct CommInitRankImpl; + +class CommInitRank : public Operation { + public: + template + static ReturnStatus Execute(void **comm_handle, Args &&...args) { + Communicator *&comm = *reinterpret_cast(comm_handle); + if (comm && comm->intra_comm()) { + // TODO(lzm): change to use `glog`. + LOG("Invalid communicator handle for `CommInitRank`."); + return ReturnStatus::kInvalidArgument; + } + + constexpr Device::Type kDev = + ListGetBest(ActiveDevices{}); + using Rt = Runtime; + + int current_dev = 0; + CHECK_STATUS(Rt, Rt::GetDevice(¤t_dev)); + + if (!comm) { + comm = new Communicator(kDev, current_dev); + } else { + comm->set_device_id(current_dev); + } + + return CommInitRankImpl::Apply( + comm, std::forward(args)...); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BASE_COMM_INIT_RANK_H_ diff --git a/src/base/get_unique_id.h b/src/base/get_unique_id.h new file mode 100644 index 0000000..d726396 --- /dev/null +++ b/src/base/get_unique_id.h @@ -0,0 +1,23 @@ +#ifndef INFINI_CCL_BASE_GET_UNIQUE_ID_H_ +#define INFINI_CCL_BASE_GET_UNIQUE_ID_H_ + +#include "comm.h" +#include "operation.h" +#include "return_status_impl.h" + +namespace infini::ccl { + +template +struct GetUniqueIdImpl; + +class GetUniqueId : public Operation { + public: + template + static ReturnStatus Execute(infinicclUniqueId *id) { + return GetUniqueIdImpl::Apply(id); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BASE_GET_UNIQUE_ID_H_ diff --git a/src/cpu/runtime_.h b/src/cpu/runtime_.h index a704719..6bce0a8 100644 --- a/src/cpu/runtime_.h +++ b/src/cpu/runtime_.h @@ -47,6 +47,11 @@ struct Runtime : RuntimeBase> { static constexpr int MemcpyDeviceToHost = 1; + static constexpr auto GetDevice = [](int *dev) { + *dev = 0; + return ReturnStatus::kSuccess; + }; + static constexpr auto SetDevice = [](int) { return ReturnStatus::kSuccess; }; static constexpr auto DeviceSynchronize = []() { diff --git a/src/metax/runtime_.h b/src/metax/runtime_.h index 9de9036..bfe55ab 100644 --- a/src/metax/runtime_.h +++ b/src/metax/runtime_.h @@ -42,6 +42,8 @@ struct Runtime static constexpr auto Memset = mcMemset; + static constexpr auto GetDevice = mcGetDevice; + static constexpr auto SetDevice = mcSetDevice; static constexpr auto DeviceSynchronize = mcDeviceSynchronize; diff --git a/src/nvidia/nccl/checks.h b/src/nvidia/nccl/checks.h new file mode 100644 index 0000000..09329e1 --- /dev/null +++ b/src/nvidia/nccl/checks.h @@ -0,0 +1,31 @@ +#ifndef INFINI_CCL_NVIDIA_NCCL_CHECKS_H_ +#define INFINI_CCL_NVIDIA_NCCL_CHECKS_H_ + +#include + +#include + +#include "return_status_impl.h" + +#define INFINI_CHECK_NCCL(result) \ + ::infini::ccl::detail::CheckNcclImpl((result), __FILE__, __LINE__) + +namespace infini::ccl { + +namespace detail { + +inline ReturnStatus CheckNcclImpl(ncclResult_t nccl_result, const char *file, + int line) { + if (nccl_result != ncclSuccess) { + std::cerr << "backend(nccl) NCCL error code: " << nccl_result << " at line " + << line << " in " << file << std::endl; + std::abort(); + } + return ReturnStatus::kSuccess; +} + +} // namespace detail + +} // namespace infini::ccl + +#endif // INFINI_CCL_NVIDIA_NCCL_CHECKS_H_ diff --git a/src/nvidia/nccl/comm_instance.h b/src/nvidia/nccl/comm_instance.h index c8d8956..0d9e925 100644 --- a/src/nvidia/nccl/comm_instance.h +++ b/src/nvidia/nccl/comm_instance.h @@ -3,13 +3,23 @@ #include +#include "checks.h" #include "communicator.h" namespace infini::ccl { struct NcclInstance : public BackendCommInstance { - ncclComm_t handle; + ncclComm_t handle = nullptr; + NcclInstance() { type = BackendType::kNccl; } + ~NcclInstance() override { Destroy(); } + + void Destroy() { + if (handle != nullptr) { + INFINI_CHECK_NCCL(ncclCommDestroy(handle)); + handle = nullptr; + } + } }; } // namespace infini::ccl diff --git a/src/nvidia/nccl/impl/all_reduce.h b/src/nvidia/nccl/impl/all_reduce.h new file mode 100644 index 0000000..e2f9a3e --- /dev/null +++ b/src/nvidia/nccl/impl/all_reduce.h @@ -0,0 +1,50 @@ +#ifndef INFINI_CCL_NVIDIA_NCCL_IMPL_ALL_REDUCE_H_ +#define INFINI_CCL_NVIDIA_NCCL_IMPL_ALL_REDUCE_H_ + +#include "base/all_reduce.h" +#include "communicator.h" +#include "nvidia/nccl/checks.h" +#include "nvidia/nccl/comm_instance.h" +#include "nvidia/nccl/type_map.h" +#include "nvidia/runtime_.h" + +namespace infini::ccl { + +template <> +class AllReduceImpl { + public: + static ReturnStatus Apply(const void *send_buff, void *recv_buff, + size_t count, DataType data_type, + ReductionOpType op, Communicator *comm, + void *stream) { + auto *comm_internal = static_cast(comm); + if (!comm_internal) { + return ReturnStatus::kInternalError; + } + + constexpr Device::Type kDev = + ListGetBest(ActiveDevices{}); + using Rt = Runtime; + + auto *intra = static_cast(comm_internal->intra_comm()); + if (!intra || !intra->handle) { + return ReturnStatus::kInternalError; + } + + ncclDataType_t nccl_type = DataTypeToNcclType(data_type); + ncclRedOp_t nccl_op = RedOpToNcclOp(op); + + INFINI_CHECK_NCCL(ncclAllReduce(send_buff, recv_buff, count, nccl_type, + nccl_op, intra->handle, + reinterpret_cast(stream))); + + return ReturnStatus::kSuccess; + } +}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_NVIDIA_NCCL_IMPL_ALL_REDUCE_H_ diff --git a/src/nvidia/nccl/impl/comm_destroy.h b/src/nvidia/nccl/impl/comm_destroy.h new file mode 100644 index 0000000..452a27e --- /dev/null +++ b/src/nvidia/nccl/impl/comm_destroy.h @@ -0,0 +1,36 @@ +#ifndef INFINI_CCL_NVIDIA_NCCL_IMPL_COMM_DESTROY_H_ +#define INFINI_CCL_NVIDIA_NCCL_IMPL_COMM_DESTROY_H_ + +#include "base/comm_destroy.h" +#include "communicator.h" +#include "nvidia/nccl/checks.h" +#include "nvidia/nccl/comm_instance.h" + +namespace infini::ccl { + +template <> +class CommDestroyImpl { + public: + static ReturnStatus Apply(void *comm) { + auto *comm_internal = static_cast(comm); + if (!comm_internal) { + return ReturnStatus::kInternalError; + } + + if (auto *intra = + static_cast(comm_internal->intra_comm())) { + intra->Destroy(); + } + + comm_internal->set_intra_comm(nullptr); + + return ReturnStatus::kSuccess; + } +}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_NVIDIA_NCCL_IMPL_COMM_DESTROY_H_ diff --git a/src/nvidia/nccl/impl/comm_init_rank.h b/src/nvidia/nccl/impl/comm_init_rank.h new file mode 100644 index 0000000..f03a3de --- /dev/null +++ b/src/nvidia/nccl/impl/comm_init_rank.h @@ -0,0 +1,45 @@ +#ifndef INFINI_CCL_NVIDIA_NCCL_IMPL_COMM_INIT_RANK_H_ +#define INFINI_CCL_NVIDIA_NCCL_IMPL_COMM_INIT_RANK_H_ + +#include + +#include "base/comm_init_rank.h" +#include "communicator.h" +#include "logging.h" +#include "nvidia/nccl/checks.h" +#include "nvidia/nccl/comm_instance.h" + +namespace infini::ccl { + +template <> +class CommInitRankImpl { + public: + static ReturnStatus Apply(Communicator *comm, int nranks, + infinicclUniqueId id, int rank) { + constexpr Device::Type kDev = + ListGetBest(ActiveDevices{}); + using Rt = Runtime; + + const ncclUniqueId *nccl_id = + reinterpret_cast(id.internal); + + ncclComm_t nccl_handle; + INFINI_CHECK_NCCL(ncclCommInitRank(&nccl_handle, nranks, *nccl_id, rank)); + + comm->set_world_info(rank, nranks); + + auto inst = std::make_unique(); + inst->handle = nccl_handle; + + comm->set_intra_comm(std::move(inst)); + + return ReturnStatus::kSuccess; + } +}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif diff --git a/src/nvidia/nccl/impl/get_unique_id.h b/src/nvidia/nccl/impl/get_unique_id.h new file mode 100644 index 0000000..f84e37d --- /dev/null +++ b/src/nvidia/nccl/impl/get_unique_id.h @@ -0,0 +1,38 @@ +#ifndef INFINI_CCL_NVIDIA_NCCL_IMPL_GET_UNIQUE_ID_H_ +#define INFINI_CCL_NVIDIA_NCCL_IMPL_GET_UNIQUE_ID_H_ + +#include + +#include + +#include "base/get_unique_id.h" +#include "nvidia/nccl/checks.h" + +namespace infini::ccl { + +template <> +class GetUniqueIdImpl { + public: + static ReturnStatus Apply(infinicclUniqueId *id) { + ncclUniqueId nccl_id; + + INFINI_CHECK_NCCL(ncclGetUniqueId(&nccl_id)); + + // Safety check: ensure our public struct is large enough. + static_assert( + sizeof(ncclUniqueId) <= sizeof(id->internal), + "`infinicclUniqueId::internal` is too small for `ncclUniqueId`"); + + std::memcpy(id->internal, &nccl_id, sizeof(nccl_id)); + + return ReturnStatus::kSuccess; + } +}; + +// Enable this backend for the `GetUniqueId` operation. +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_NVIDIA_NCCL_IMPL_GET_UNIQUE_ID_H_ diff --git a/src/nvidia/nccl/type_map.h b/src/nvidia/nccl/type_map.h new file mode 100644 index 0000000..a16d227 --- /dev/null +++ b/src/nvidia/nccl/type_map.h @@ -0,0 +1,61 @@ +#ifndef INFINI_CCL_NVIDIA_NCCL_IMPL_TYPE_MAP_H_ +#define INFINI_CCL_NVIDIA_NCCL_IMPL_TYPE_MAP_H_ + +#include + +#include "comm_impl.h" +#include "data_type_impl.h" +#include "logging.h" + +namespace infini::ccl { + +#if defined(__CUDA_BF16_TYPES_EXIST__) +constexpr ncclDataType_t kNcclBFloat16Val = ncclBfloat16; +#else +constexpr ncclDataType_t kNcclBFloat16Val = ncclNumTypes; +#endif + +static const ConstexprMap kNcclTypeMap{{{ + {DataType::kInt8, ncclInt8}, + {DataType::kInt16, ncclNumTypes}, + {DataType::kInt32, ncclInt32}, + {DataType::kInt64, ncclInt64}, + {DataType::kUInt8, ncclUint8}, + {DataType::kUInt16, ncclNumTypes}, + {DataType::kUInt32, ncclUint32}, + {DataType::kUInt64, ncclUint64}, + {DataType::kFloat32, ncclFloat32}, + {DataType::kFloat64, ncclFloat64}, + {DataType::kFloat16, ncclFloat16}, + {DataType::kBFloat16, kNcclBFloat16Val}, +}}}; + +static const ConstexprMap kNcclOpMap{{{ + {ReductionOpType::kSum, ncclSum}, + {ReductionOpType::kProd, ncclProd}, + {ReductionOpType::kMax, ncclMax}, + {ReductionOpType::kMin, ncclMin}, + {ReductionOpType::kAvg, ncclAvg}, +}}}; + +inline ncclDataType_t DataTypeToNcclType(DataType dtype) { + auto nccl_dtype = kNcclTypeMap.at(dtype); + + if (nccl_dtype == ncclNumTypes) { + // This means the requested data type is not supported by NCCL. + // TODO(lzm): change to use `glog`. + LOG(("DataType '" + std::string(kDataTypeToDesc.at(dtype)) + + "' is not supported by the NCCL backend") + .c_str()); + } + + return nccl_dtype; +} + +inline ncclRedOp_t RedOpToNcclOp(ReductionOpType red_op) { + return kNcclOpMap.at(red_op); +} + +} // namespace infini::ccl + +#endif // INFINI_CCL_NVIDIA_NCCL_IMPL_TYPE_MAP_H_ diff --git a/src/nvidia/runtime_.h b/src/nvidia/runtime_.h index 53a1dd2..e8f82bc 100644 --- a/src/nvidia/runtime_.h +++ b/src/nvidia/runtime_.h @@ -44,6 +44,8 @@ struct Runtime static constexpr auto Memset = cudaMemset; + static constexpr auto GetDevice = cudaGetDevice; + static constexpr auto SetDevice = cudaSetDevice; static constexpr auto DeviceSynchronize = cudaDeviceSynchronize; diff --git a/src/ompi/comm_instance.h b/src/ompi/comm_instance.h index 1d79db9..87ecdf0 100644 --- a/src/ompi/comm_instance.h +++ b/src/ompi/comm_instance.h @@ -3,6 +3,7 @@ #include +#include "checks.h" #include "communicator.h" namespace infini::ccl { @@ -11,11 +12,12 @@ struct OmpiInstance : public BackendCommInstance { MPI_Comm handle = MPI_COMM_NULL; OmpiInstance() { type = BackendType::kOmpi; } + ~OmpiInstance() override { Destroy(); } void Destroy() { - // Ensure we don't accidentally leak if a backend duplicates a communicator. if (handle != MPI_COMM_WORLD && handle != MPI_COMM_NULL) { - MPI_Comm_free(&handle); + INFINI_CHECK_MPI(MPI_Comm_free(&handle)); + handle = MPI_COMM_NULL; } } }; diff --git a/src/operation.h b/src/operation.h index 9ef3d19..60c1a38 100644 --- a/src/operation.h +++ b/src/operation.h @@ -4,8 +4,10 @@ #include #include "backend.h" +#include "backend_device_map.h" #include "device.h" #include "dispatcher.h" +#include "return_status_impl.h" #include "traits.h" namespace infini::ccl { @@ -34,8 +36,12 @@ class Operation { constexpr Device::Type kDevice = static_cast(ListGet<1>(resolved_list)); - return Key::template Execute( - std::forward(args)...); + if constexpr (IsSupportedCombination::value) { + return Key::template Execute( + std::forward(args)...); + } else { + return ReturnStatus::kNotSupported; + } }, "Operation::Call"); } diff --git a/src/return_status_impl.h b/src/return_status_impl.h index 9a19ecb..253575b 100644 --- a/src/return_status_impl.h +++ b/src/return_status_impl.h @@ -16,6 +16,7 @@ enum class ReturnStatus : int8_t { kInvalidUsage = infinicclInvalidUsage, kRemoteError = infinicclRemoteError, kInProgress = infinicclInProgress, + kNotSupported = infinicclNotSupported, kNumResults = infinicclNumResults, };