[cuda] Port over channel implementation via NCCL (#14059)
This commit ports over channel implementation via NCCL.
Modifications over the existing code is mostly around removing
the context wrapper and making status handling better.
Progress towards https://github.com/openxla/iree/issues/13245
diff --git a/experimental/cuda2/CMakeLists.txt b/experimental/cuda2/CMakeLists.txt
index d8f3df0..3f3fdef 100644
--- a/experimental/cuda2/CMakeLists.txt
+++ b/experimental/cuda2/CMakeLists.txt
@@ -26,6 +26,8 @@
"cuda_driver.c"
"memory_pools.c"
"memory_pools.h"
+ "nccl_channel.c"
+ "nccl_channel.h"
"pipeline_layout.c"
"pipeline_layout.h"
DEPS
@@ -61,6 +63,7 @@
iree::base
iree::base::internal::dynamic_library
iree_cuda::headers
+ nccl::headers
PUBLIC
)
diff --git a/experimental/cuda2/api.h b/experimental/cuda2/api.h
index 777d6ef..20951b0 100644
--- a/experimental/cuda2/api.h
+++ b/experimental/cuda2/api.h
@@ -20,6 +20,11 @@
// iree_hal_cuda2_device_t
//===----------------------------------------------------------------------===//
+// ncclUniqueId exposed without exporting the NCCL headers.
+typedef struct {
+ char data[128];
+} iree_hal_cuda2_nccl_id_t;
+
// Parameters defining a CUmemoryPool.
typedef struct iree_hal_cuda2_memory_pool_params_t {
// Minimum number of bytes to keep in the pool when trimming with
diff --git a/experimental/cuda2/cuda_device.c b/experimental/cuda2/cuda_device.c
index 5ed2dc3..7bbfdd5 100644
--- a/experimental/cuda2/cuda_device.c
+++ b/experimental/cuda2/cuda_device.c
@@ -6,7 +6,6 @@
#include "experimental/cuda2/cuda_device.h"
-#include <iree/hal/device.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
@@ -16,6 +15,8 @@
#include "experimental/cuda2/cuda_dynamic_symbols.h"
#include "experimental/cuda2/cuda_status_util.h"
#include "experimental/cuda2/memory_pools.h"
+#include "experimental/cuda2/nccl_channel.h"
+#include "experimental/cuda2/nccl_dynamic_symbols.h"
#include "experimental/cuda2/pipeline_layout.h"
#include "iree/base/internal/arena.h"
#include "iree/base/internal/math.h"
@@ -41,6 +42,7 @@
iree_hal_driver_t* driver;
const iree_hal_cuda2_dynamic_symbols_t* cuda_symbols;
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols;
// Parameters used to control device behavior.
iree_hal_cuda2_device_params_t params;
@@ -56,6 +58,9 @@
bool supports_memory_pools;
iree_hal_cuda2_memory_pools_t memory_pools;
iree_hal_allocator_t* device_allocator;
+
+ // Optional provider used for creating/configuring collective channels.
+ iree_hal_channel_provider_t* channel_provider;
} iree_hal_cuda2_device_t;
static const iree_hal_device_vtable_t iree_hal_cuda2_device_vtable;
@@ -96,7 +101,8 @@
iree_hal_driver_t* driver, iree_string_view_t identifier,
const iree_hal_cuda2_device_params_t* params, CUdevice cu_device,
CUstream stream, CUcontext context,
- const iree_hal_cuda2_dynamic_symbols_t* symbols,
+ const iree_hal_cuda2_dynamic_symbols_t* cuda_symbols,
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols,
iree_allocator_t host_allocator, iree_hal_device_t** out_device) {
iree_hal_cuda2_device_t* device = NULL;
iree_host_size_t total_size = iree_sizeof_struct(*device) + identifier.size;
@@ -113,7 +119,8 @@
&device->block_pool);
device->driver = driver;
iree_hal_driver_retain(device->driver);
- device->cuda_symbols = symbols;
+ device->cuda_symbols = cuda_symbols;
+ device->nccl_symbols = nccl_symbols;
device->params = *params;
device->cu_context = context;
device->cu_device = cu_device;
@@ -126,7 +133,7 @@
if (iree_status_is_ok(status) && params->async_allocations) {
int supports_memory_pools = 0;
status = IREE_CURESULT_TO_STATUS(
- symbols,
+ cuda_symbols,
cuDeviceGetAttribute(&supports_memory_pools,
CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED,
cu_device),
@@ -137,13 +144,13 @@
// Create memory pools first so that we can share them with the allocator.
if (iree_status_is_ok(status) && device->supports_memory_pools) {
status = iree_hal_cuda2_memory_pools_initialize(
- symbols, cu_device, ¶ms->memory_pools, host_allocator,
+ cuda_symbols, cu_device, ¶ms->memory_pools, host_allocator,
&device->memory_pools);
}
if (iree_status_is_ok(status)) {
status = iree_hal_cuda2_allocator_create(
- (iree_hal_device_t*)device, symbols, cu_device, stream,
+ (iree_hal_device_t*)device, cuda_symbols, cu_device, stream,
device->supports_memory_pools ? &device->memory_pools : NULL,
host_allocator, &device->device_allocator);
}
@@ -159,11 +166,12 @@
iree_status_t iree_hal_cuda2_device_create(
iree_hal_driver_t* driver, iree_string_view_t identifier,
const iree_hal_cuda2_device_params_t* params,
- const iree_hal_cuda2_dynamic_symbols_t* symbols, CUdevice device,
+ const iree_hal_cuda2_dynamic_symbols_t* cuda_symbols,
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols, CUdevice device,
iree_allocator_t host_allocator, iree_hal_device_t** out_device) {
IREE_ASSERT_ARGUMENT(driver);
IREE_ASSERT_ARGUMENT(params);
- IREE_ASSERT_ARGUMENT(symbols);
+ IREE_ASSERT_ARGUMENT(cuda_symbols);
IREE_ASSERT_ARGUMENT(out_device);
IREE_TRACE_ZONE_BEGIN(z0);
@@ -173,27 +181,27 @@
CUcontext context = NULL;
if (iree_status_is_ok(status)) {
status = IREE_CURESULT_TO_STATUS(
- symbols, cuDevicePrimaryCtxRetain(&context, device));
+ cuda_symbols, cuDevicePrimaryCtxRetain(&context, device));
}
if (iree_status_is_ok(status)) {
- status = IREE_CURESULT_TO_STATUS(symbols, cuCtxSetCurrent(context));
+ status = IREE_CURESULT_TO_STATUS(cuda_symbols, cuCtxSetCurrent(context));
}
// Create the default stream for the device.
CUstream stream = NULL;
if (iree_status_is_ok(status)) {
status = IREE_CURESULT_TO_STATUS(
- symbols, cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING));
+ cuda_symbols, cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING));
}
if (iree_status_is_ok(status)) {
status = iree_hal_cuda2_device_create_internal(
- driver, identifier, params, device, stream, context, symbols,
- host_allocator, out_device);
+ driver, identifier, params, device, stream, context, cuda_symbols,
+ nccl_symbols, host_allocator, out_device);
}
if (!iree_status_is_ok(status)) {
- if (stream) symbols->cuStreamDestroy(stream);
- if (context) symbols->cuDevicePrimaryCtxRelease(device);
+ if (stream) cuda_symbols->cuStreamDestroy(stream);
+ if (context) cuda_symbols->cuDevicePrimaryCtxRelease(device);
}
IREE_TRACE_ZONE_END(z0);
@@ -221,6 +229,9 @@
// There should be no more buffers live that use the allocator.
iree_hal_allocator_release(device->device_allocator);
+ // Buffers may have been retaining collective resources.
+ iree_hal_channel_provider_release(device->channel_provider);
+
// Destroy memory pools that hold on to reserved memory.
iree_hal_cuda2_memory_pools_deinitialize(&device->memory_pools);
@@ -269,7 +280,10 @@
static void iree_hal_cuda2_replace_channel_provider(
iree_hal_device_t* base_device, iree_hal_channel_provider_t* new_provider) {
- // TODO: implement this together with channel support.
+ iree_hal_cuda2_device_t* device = iree_hal_cuda2_device_cast(base_device);
+ iree_hal_channel_provider_retain(new_provider);
+ iree_hal_channel_provider_release(device->channel_provider);
+ device->channel_provider = new_provider;
}
static iree_status_t iree_hal_cuda2_device_trim(
@@ -327,8 +341,87 @@
static iree_status_t iree_hal_cuda2_device_create_channel(
iree_hal_device_t* base_device, iree_hal_queue_affinity_t queue_affinity,
iree_hal_channel_params_t params, iree_hal_channel_t** out_channel) {
- return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
- "channel not yet implmeneted");
+ iree_hal_cuda2_device_t* device = iree_hal_cuda2_device_cast(base_device);
+ if (!device->nccl_symbols || !device->nccl_symbols->dylib) {
+ return iree_make_status(
+ IREE_STATUS_UNAVAILABLE,
+ "NCCL runtime library (%d.%d.%d) not available; ensure installed and "
+ "the shared library is on your PATH/LD_LIBRARY_PATH "
+ "(nccl.dll/libnccl.so)",
+ NCCL_MAJOR, NCCL_MINOR, NCCL_PATCH);
+ }
+
+ // Today we only allow a single logical device per channel.
+ // We could multiplex channels but it'd be better to surface that to the
+ // compiler so that it can emit the right rank math.
+ int requested_count = iree_math_count_ones_u64(queue_affinity);
+ // TODO(#12206): properly assign affinity in the compiler.
+ if (requested_count != 64 && requested_count != 1) {
+ return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "exactly one participant is allowed in a "
+ "channel but %d were specified",
+ requested_count);
+ }
+
+ // Ask the channel provider (if configured) for the default rank and count
+ // if the user did not set them.
+ if (device->channel_provider &&
+ (params.rank == IREE_HAL_CHANNEL_RANK_DEFAULT ||
+ params.count == IREE_HAL_CHANNEL_COUNT_DEFAULT)) {
+ IREE_RETURN_IF_ERROR(
+ iree_hal_channel_provider_query_default_rank_and_count(
+ device->channel_provider, ¶ms.rank, ¶ms.count),
+ "querying default collective group rank and count");
+ }
+
+ // An ID is required to initialize NCCL. On the root it'll be the local ID and
+ // on all other participants it'll be the root ID.
+ iree_hal_cuda2_nccl_id_t id;
+ memset(&id, 0, sizeof(id));
+ if (iree_const_byte_span_is_empty(params.id)) {
+ // User wants the default ID.
+ if (!device->channel_provider) {
+ return iree_make_status(
+ IREE_STATUS_INVALID_ARGUMENT,
+ "default collective channel ID requested but no channel provider has "
+ "been set on the device to provide it");
+ }
+ if (params.rank == 0) {
+ // Bootstrap NCCL to get the root ID.
+ IREE_RETURN_IF_ERROR(
+ iree_hal_cuda2_nccl_get_unique_id(device->nccl_symbols, &id),
+ "bootstrapping NCCL root");
+ }
+ // Exchange NCCL ID with all participants.
+ IREE_RETURN_IF_ERROR(iree_hal_channel_provider_exchange_default_id(
+ device->channel_provider,
+ iree_make_byte_span((void*)&id, sizeof(id))),
+ "exchanging NCCL ID with other participants");
+ } else if (params.id.data_length != IREE_ARRAYSIZE(id.data)) {
+ // User provided something but it's not what we expect.
+ return iree_make_status(
+ IREE_STATUS_INVALID_ARGUMENT,
+ "NCCL ID must be %" PRIhsz
+ " bytes matching the ncclUniqueId struct but caller provided %" PRIhsz
+ " bytes",
+ IREE_ARRAYSIZE(id.data), sizeof(id));
+ } else {
+ // User provided the ID - we treat it as opaque here and let NCCL validate.
+ memcpy(id.data, params.id.data, IREE_ARRAYSIZE(id.data));
+ }
+
+ if (iree_hal_cuda2_nccl_id_is_empty(&id)) {
+ // TODO: maybe this is ok? a localhost alias or something?
+ return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "no default NCCL ID specified (all zeros)");
+ }
+
+ // TODO: when we support multiple logical devices we'll want to pass in the
+ // context of the device mapped to the queue_affinity. For now since this
+ // implementation only supports one device we pass in the only one we have.
+ return iree_hal_cuda2_nccl_channel_create(
+ device->cuda_symbols, device->nccl_symbols, &id, params.rank,
+ params.count, device->host_allocator, out_channel);
}
static iree_status_t iree_hal_cuda2_device_create_command_buffer(
diff --git a/experimental/cuda2/cuda_device.h b/experimental/cuda2/cuda_device.h
index 70dbd36..5def0de 100644
--- a/experimental/cuda2/cuda_device.h
+++ b/experimental/cuda2/cuda_device.h
@@ -9,6 +9,7 @@
#include "experimental/cuda2/api.h"
#include "experimental/cuda2/cuda_dynamic_symbols.h"
+#include "experimental/cuda2/nccl_dynamic_symbols.h"
#include "iree/base/api.h"
#include "iree/hal/api.h"
@@ -20,7 +21,8 @@
iree_status_t iree_hal_cuda2_device_create(
iree_hal_driver_t* driver, iree_string_view_t identifier,
const iree_hal_cuda2_device_params_t* params,
- const iree_hal_cuda2_dynamic_symbols_t* symbols, CUdevice device,
+ const iree_hal_cuda2_dynamic_symbols_t* symbols,
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols, CUdevice device,
iree_allocator_t host_allocator, iree_hal_device_t** out_device);
// Returns the CUDA context bound to the given |device| if it is a CUDA device
diff --git a/experimental/cuda2/cuda_driver.c b/experimental/cuda2/cuda_driver.c
index 40c09a4..b72dab2 100644
--- a/experimental/cuda2/cuda_driver.c
+++ b/experimental/cuda2/cuda_driver.c
@@ -489,7 +489,7 @@
// Attempt to create the device now.
iree_status_t status = iree_hal_cuda2_device_create(
base_driver, device_name, &driver->device_params, &driver->cuda_symbols,
- device, host_allocator, out_device);
+ &driver->nccl_symbols, device, host_allocator, out_device);
IREE_TRACE_ZONE_END(z0);
return status;
diff --git a/experimental/cuda2/nccl_channel.c b/experimental/cuda2/nccl_channel.c
new file mode 100644
index 0000000..6d3cbb1
--- /dev/null
+++ b/experimental/cuda2/nccl_channel.c
@@ -0,0 +1,568 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "experimental/cuda2/nccl_channel.h"
+
+#include <stddef.h>
+#include <stdlib.h>
+
+#include "experimental/cuda2/cuda_buffer.h"
+#include "experimental/cuda2/cuda_status_util.h"
+#include "experimental/cuda2/nccl_headers.h"
+#include "experimental/cuda2/nccl_status_util.h"
+#include "iree/base/api.h"
+#include "iree/base/tracing.h"
+
+typedef struct iree_hal_cuda2_nccl_channel_t {
+ iree_hal_resource_t resource;
+
+ const iree_hal_cuda2_dynamic_symbols_t* cuda_symbols;
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols;
+
+ iree_allocator_t host_allocator;
+
+ // Parent channel this was split from, if any.
+ // This is only used to keep the parent channel live for as long as there are
+ // any split channels live (including transitive splits).
+ iree_hal_channel_t* parent_channel;
+
+ // This participant's rank in the communicator.
+ // Equivalent to ncclCommUserRank.
+ int rank;
+ // Total number of participants in the communicator.
+ // Equivalent to ncclCommCount.
+ int count;
+
+ // Communicator handle.
+ ncclComm_t comm;
+
+ // Hash of the unique ID used to create the communicator.
+ // This is consistent with the hashes NCCL itself uses for logging but is not
+ // guaranteed to be unique - only use for informational purposes.
+ IREE_TRACE(uint64_t id_hash;)
+} iree_hal_cuda2_nccl_channel_t;
+
+static const iree_hal_channel_vtable_t iree_hal_cuda2_nccl_channel_vtable;
+
+static iree_hal_cuda2_nccl_channel_t* iree_hal_cuda2_nccl_channel_cast(
+ iree_hal_channel_t* base_value) {
+ IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_cuda2_nccl_channel_vtable);
+ return (iree_hal_cuda2_nccl_channel_t*)base_value;
+}
+
+static const iree_hal_cuda2_nccl_channel_t*
+iree_hal_cuda2_nccl_channel_const_cast(const iree_hal_channel_t* base_value) {
+ IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_cuda2_nccl_channel_vtable);
+ return (const iree_hal_cuda2_nccl_channel_t*)base_value;
+}
+
+// Returns the same value as NCCL's init.cc hashUniqueId.
+// These magic constants were chosen by their implementation and unlikely to
+// be stable as it's not part of their public API. So they are only meant to be
+// used for correlating debug logging/traces. We keep it internal here too so
+// that we aren't tempted to use it in other places.
+static uint64_t iree_hal_cuda2_nccl_hash_id(
+ const iree_hal_cuda2_nccl_id_t* id) {
+ uint64_t hash = 0xDEADBEEF;
+ for (iree_host_size_t i = 0; i < sizeof(*id); i++) {
+ hash ^= hash >> 32;
+ hash *= 0x8DB3DB47FA2994ADull;
+ hash += id->data[i];
+ }
+ return hash;
+}
+
+iree_status_t iree_hal_cuda2_nccl_get_unique_id(
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* symbols,
+ iree_hal_cuda2_nccl_id_t* out_id) {
+ static_assert(sizeof(*out_id) == sizeof(ncclUniqueId),
+ "NCCL ID size mismatch");
+
+ IREE_ASSERT_ARGUMENT(symbols);
+ IREE_ASSERT_ARGUMENT(out_id);
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ memset(out_id, 0, sizeof(*out_id));
+ iree_status_t status = IREE_NCCL_RESULT_TO_STATUS(
+ symbols, ncclGetUniqueId((ncclUniqueId*)out_id), "ncclGetUniqueId");
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
+iree_status_t iree_hal_cuda2_nccl_channel_create(
+ const iree_hal_cuda2_dynamic_symbols_t* cuda_symbols,
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols,
+ const iree_hal_cuda2_nccl_id_t* id, int rank, int count,
+ iree_allocator_t host_allocator, iree_hal_channel_t** out_channel) {
+ IREE_ASSERT_ARGUMENT(cuda_symbols);
+ IREE_ASSERT_ARGUMENT(nccl_symbols);
+ IREE_ASSERT_ARGUMENT(id);
+ IREE_ASSERT_ARGUMENT(out_channel);
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ *out_channel = NULL;
+ IREE_TRACE(const uint64_t id_hash = iree_hal_cuda2_nccl_hash_id(id));
+ IREE_TRACE_ZONE_APPEND_VALUE(z0, id_hash);
+ IREE_TRACE_ZONE_APPEND_VALUE(z0, rank);
+ IREE_TRACE_ZONE_APPEND_VALUE(z0, count);
+
+ ncclComm_t comm = NULL;
+ ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+ // TODO: use async to check a timeout.
+ config.blocking = 1;
+ IREE_NCCL_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, nccl_symbols,
+ ncclCommInitRankConfig(&comm, count, *((const ncclUniqueId*)id), rank,
+ &config),
+ "ncclCommInitRankConfig");
+
+ iree_hal_cuda2_nccl_channel_t* channel = NULL;
+ iree_status_t status =
+ iree_allocator_malloc(host_allocator, sizeof(*channel), (void**)&channel);
+
+ if (iree_status_is_ok(status)) {
+ iree_hal_resource_initialize(&iree_hal_cuda2_nccl_channel_vtable,
+ &channel->resource);
+ channel->cuda_symbols = cuda_symbols;
+ channel->nccl_symbols = nccl_symbols;
+ channel->host_allocator = host_allocator;
+ channel->parent_channel = NULL;
+ channel->rank = rank;
+ channel->count = count;
+ channel->comm = comm;
+ IREE_TRACE(channel->id_hash = id_hash);
+ *out_channel = (iree_hal_channel_t*)channel;
+ }
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
+static void iree_hal_cuda2_nccl_channel_destroy(
+ iree_hal_channel_t* base_channel) {
+ iree_hal_cuda2_nccl_channel_t* channel =
+ iree_hal_cuda2_nccl_channel_cast(base_channel);
+ IREE_TRACE_ZONE_BEGIN(z0);
+ IREE_TRACE_ZONE_APPEND_VALUE(z0, channel->id_hash);
+ IREE_TRACE_ZONE_APPEND_VALUE(z0, channel->rank);
+ IREE_TRACE_ZONE_APPEND_VALUE(z0, channel->count);
+
+ iree_allocator_t host_allocator = channel->host_allocator;
+
+ // TODO(#9580): support async tear down
+ // We could be smarter about starting finalization of all channels async and
+ // then waiting for them to complete but we aren't currently optimizing for
+ // lifetime performance. To do that we'd probably want to track each open
+ // channel on the device that created them and manage teardown there.
+ //
+ // Recommended:
+ // ncclCommFinalize(channel->comm); // non-blocking!
+ // while (ncclCommGetAsyncError == ncclInProgress) sleep(1);
+ // ncclCommDestroy(channel->comm)
+ // Should work the same (as we are doing a blocking teardown):
+ // ncclCommDestroy(channel->comm)
+ IREE_NCCL_IGNORE_ERROR(channel->nccl_symbols,
+ ncclCommFinalize(channel->comm));
+
+ IREE_NCCL_IGNORE_ERROR(channel->nccl_symbols, ncclCommDestroy(channel->comm));
+
+ iree_hal_channel_release(channel->parent_channel);
+ iree_allocator_free(host_allocator, channel);
+
+ IREE_TRACE_ZONE_END(z0);
+}
+
+static iree_status_t iree_hal_cuda2_nccl_channel_split(
+ iree_hal_channel_t* base_channel, int32_t color, int32_t key,
+ iree_hal_channel_flags_t flags, iree_hal_channel_t** out_split_channel) {
+ iree_hal_cuda2_nccl_channel_t* channel =
+ iree_hal_cuda2_nccl_channel_cast(base_channel);
+
+ // TODO: see if we need to set the sharing config - we may always want to.
+ ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+ // TODO: use async to check a timeout.
+ config.blocking = 1;
+
+ // Split the communicator.
+ ncclComm_t split_comm = NULL;
+ IREE_NCCL_RETURN_IF_ERROR(
+ channel->nccl_symbols,
+ ncclCommSplit(channel->comm, color, key, &split_comm, &config),
+ "ncclCommSplit");
+
+ // Query the local rank/count from the split communicator.
+ int split_rank = 0;
+ int split_count = 0;
+ iree_status_t status = IREE_NCCL_RESULT_TO_STATUS(
+ channel->nccl_symbols, ncclCommUserRank(split_comm, &split_rank),
+ "ncclCommUserRank");
+ if (iree_status_is_ok(status)) {
+ status = IREE_NCCL_RESULT_TO_STATUS(channel->nccl_symbols,
+ ncclCommCount(split_comm, &split_count),
+ "ncclCommCount");
+ }
+
+ // Wrap the split communicator in a new channel.
+ iree_hal_cuda2_nccl_channel_t* split_channel = NULL;
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_allocator_malloc(channel->host_allocator, sizeof(*split_channel),
+ (void**)&split_channel);
+ }
+ if (iree_status_is_ok(status)) {
+ iree_hal_resource_initialize(&iree_hal_cuda2_nccl_channel_vtable,
+ &split_channel->resource);
+ split_channel->cuda_symbols = channel->cuda_symbols;
+ split_channel->nccl_symbols = channel->nccl_symbols;
+ split_channel->host_allocator = channel->host_allocator;
+ split_channel->parent_channel = base_channel;
+ iree_hal_channel_retain(base_channel);
+ split_channel->rank = split_rank;
+ split_channel->count = split_count;
+ split_channel->comm = split_comm;
+ *out_split_channel = (iree_hal_channel_t*)split_channel;
+ }
+
+ if (!iree_status_is_ok(status)) {
+ IREE_NCCL_IGNORE_ERROR(channel->nccl_symbols, ncclCommDestroy(split_comm));
+ }
+ return status;
+}
+
+static void iree_hal_cuda2_nccl_channel_query_rank_and_count(
+ const iree_hal_channel_t* base_channel, int32_t* out_rank,
+ int32_t* out_count) {
+ IREE_ASSERT_ARGUMENT(base_channel);
+ IREE_ASSERT_ARGUMENT(out_count);
+ const iree_hal_cuda2_nccl_channel_t* channel =
+ iree_hal_cuda2_nccl_channel_const_cast(base_channel);
+ // NOTE: since it's cheap we keep rank/count local - this lets us trace them
+ // out without needing to call into NCCL each time.
+ *out_rank = channel->rank;
+ *out_count = channel->count;
+}
+
+// Returns the NCCL communicator for the given |channel|, if available.
+static ncclComm_t iree_hal_cuda2_nccl_channel_comm(
+ iree_hal_channel_t* base_channel) {
+ IREE_ASSERT_ARGUMENT(base_channel);
+ iree_hal_cuda2_nccl_channel_t* channel =
+ iree_hal_cuda2_nccl_channel_cast(base_channel);
+ return channel->comm;
+}
+
+static iree_status_t iree_hal_cuda2_get_nccl_data_type(
+ iree_hal_collective_element_type_t in, ncclDataType_t* out) {
+ switch (in) {
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_SINT_8:
+ *out = ncclInt8;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_UINT_8:
+ *out = ncclUint8;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_SINT_16:
+ return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "SINT16 is not supported for collective op");
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_UINT_16:
+ return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "UINT16 is not supported for collective op");
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_SINT_32:
+ *out = ncclInt32;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_UINT_32:
+ *out = ncclUint32;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_SINT_64:
+ *out = ncclInt64;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_UINT_64:
+ *out = ncclUint64;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_FLOAT_16:
+ *out = ncclFloat16;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_FLOAT_32:
+ *out = ncclFloat32;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_FLOAT_64:
+ *out = ncclFloat64;
+ break;
+ case IREE_HAL_COLLECTIVE_ELEMENT_TYPE_BFLOAT_16:
+ *out = ncclFloat64;
+ break;
+ default:
+ return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "unhandled element type for collective op");
+ }
+ return iree_ok_status();
+}
+
+static iree_status_t iree_hal_cuda2_get_nccl_reduction_type(
+ iree_hal_collective_reduction_t in, ncclRedOp_t* out) {
+ switch (in) {
+ case IREE_HAL_COLLECTIVE_REDUCTION_SUM:
+ *out = ncclSum;
+ break;
+ case IREE_HAL_COLLECTIVE_REDUCTION_PRODUCT:
+ *out = ncclProd;
+ break;
+ case IREE_HAL_COLLECTIVE_REDUCTION_MINIMUM:
+ *out = ncclMin;
+ break;
+ case IREE_HAL_COLLECTIVE_REDUCTION_MAXIMUM:
+ *out = ncclMax;
+ break;
+ case IREE_HAL_COLLECTIVE_REDUCTION_AVERAGE:
+ *out = ncclAvg;
+ break;
+ default:
+ return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "unhandled reduction type for collective op");
+ }
+
+ return iree_ok_status();
+}
+
+static iree_status_t iree_hal_cuda2_nccl_submit_batch_entry(
+ const iree_hal_collective_batch_entry_t* entry, CUstream stream) {
+ IREE_ASSERT_ARGUMENT(entry);
+ IREE_ASSERT_ARGUMENT(stream);
+
+ iree_hal_cuda2_nccl_channel_t* channel =
+ iree_hal_cuda2_nccl_channel_cast(entry->channel);
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* symbols = channel->nccl_symbols;
+ ncclComm_t comm = iree_hal_cuda2_nccl_channel_comm(entry->channel);
+
+ ncclDataType_t datatype;
+ IREE_RETURN_IF_ERROR(
+ iree_hal_cuda2_get_nccl_data_type(entry->op.element_type, &datatype));
+
+ switch (entry->op.kind) {
+ case IREE_HAL_COLLECTIVE_KIND_ALL_GATHER: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclAllGather((const void*)sendbuff, (void*)recvbuff,
+ entry->element_count, datatype, comm, stream),
+ "ncclAllGather");
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_ALL_REDUCE: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ ncclRedOp_t redop;
+ IREE_RETURN_IF_ERROR(
+ iree_hal_cuda2_get_nccl_reduction_type(entry->op.reduction, &redop));
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclAllReduce((const void*)sendbuff, (void*)recvbuff,
+ entry->element_count, datatype, redop, comm, stream),
+ "ncclAllReduce");
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_ALL_TO_ALL: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ iree_device_size_t send_count = entry->element_count / channel->count;
+ iree_device_size_t element_size_bytes =
+ iree_hal_collective_element_byte_count(entry->op.element_type);
+ iree_device_size_t rank_offset = send_count * element_size_bytes;
+ // These calls are already grouped by iree_hal_cuda2_nccl_submit_batch.
+ for (iree_host_size_t r = 0; r < channel->count; ++r) {
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclSend((const void*)(sendbuff + r * rank_offset), send_count,
+ datatype, r, comm, stream),
+ "ncclSend");
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclRecv((void*)(recvbuff + r * rank_offset), send_count, datatype,
+ r, comm, stream),
+ "ncclRecv");
+ }
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_BROADCAST: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclBroadcast((const void*)sendbuff, (void*)recvbuff,
+ entry->element_count, datatype, entry->param, comm,
+ stream),
+ "ncclBroadcast");
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_REDUCE: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ ncclRedOp_t redop;
+ IREE_RETURN_IF_ERROR(
+ iree_hal_cuda2_get_nccl_reduction_type(entry->op.reduction, &redop));
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclReduce((const void*)sendbuff, (void*)recvbuff,
+ entry->element_count, datatype, redop, entry->param, comm,
+ stream),
+ "ncclReduce");
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_REDUCE_SCATTER: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ ncclRedOp_t redop;
+ IREE_RETURN_IF_ERROR(
+ iree_hal_cuda2_get_nccl_reduction_type(entry->op.reduction, &redop));
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclReduceScatter((const void*)sendbuff, (void*)recvbuff,
+ entry->element_count, datatype, redop, comm,
+ stream),
+ "ncclReduceScatter");
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_SEND: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclSend((const void*)sendbuff, entry->element_count, datatype,
+ entry->param, comm, stream),
+ "ncclSend");
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_RECV: {
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ IREE_NCCL_RETURN_IF_ERROR(symbols,
+ ncclRecv((void*)recvbuff, entry->element_count,
+ datatype, entry->param, comm, stream),
+ "ncclRecv");
+ break;
+ }
+ case IREE_HAL_COLLECTIVE_KIND_SEND_RECV: {
+ CUdeviceptr sendbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->send_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->send_binding.buffer) +
+ entry->send_binding.offset;
+ CUdeviceptr recvbuff =
+ iree_hal_cuda2_buffer_device_pointer(
+ iree_hal_buffer_allocated_buffer(entry->recv_binding.buffer)) +
+ iree_hal_buffer_byte_offset(entry->recv_binding.buffer) +
+ entry->recv_binding.offset;
+ int16_t sendid;
+ int16_t recvid;
+ memcpy(&sendid, &entry->param, 2);
+ memcpy(&recvid, (char*)&entry->param + 2, 2);
+ if (sendid != -1) {
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclSend((const void*)sendbuff, entry->element_count, datatype,
+ sendid, comm, stream),
+ "ncclSend");
+ }
+ if (recvid != -1) {
+ IREE_NCCL_RETURN_IF_ERROR(
+ symbols,
+ ncclRecv((void*)recvbuff, entry->element_count, datatype, recvid,
+ comm, stream),
+ "ncclRecv");
+ } else {
+ // Zero out recvbuff if this rank is not receiving any data.
+ iree_device_size_t num_bytes =
+ entry->element_count *
+ iree_hal_collective_element_byte_count(entry->op.element_type);
+ IREE_CUDA_RETURN_IF_ERROR(
+ channel->cuda_symbols,
+ cuMemsetD8Async(recvbuff, 0, num_bytes, stream), "cuMemsetD8Async");
+ }
+ break;
+ }
+ } // switch
+ return iree_ok_status();
+}
+
+iree_status_t iree_hal_cuda2_nccl_submit_batch(
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* symbols,
+ const iree_hal_collective_batch_t* batch, CUstream stream) {
+ IREE_ASSERT_ARGUMENT(symbols);
+ IREE_ASSERT_ARGUMENT(batch);
+ IREE_ASSERT_ARGUMENT(stream);
+
+ // Issue all collective operations in the batch as part of a group.
+ // NCCL may be able to fuse or reduce overheads by issuing like this.
+ IREE_NCCL_RETURN_IF_ERROR(symbols, ncclGroupStart(), "ncclGroupStart");
+ for (iree_host_size_t i = 0; i < batch->count; ++i) {
+ IREE_RETURN_IF_ERROR(
+ iree_hal_cuda2_nccl_submit_batch_entry(&batch->entries[i], stream));
+ }
+ IREE_NCCL_RETURN_IF_ERROR(symbols, ncclGroupEnd(), "ncclGroupEnd");
+
+ return iree_ok_status();
+}
+
+static const iree_hal_channel_vtable_t iree_hal_cuda2_nccl_channel_vtable = {
+ .destroy = iree_hal_cuda2_nccl_channel_destroy,
+ .split = iree_hal_cuda2_nccl_channel_split,
+ .query_rank_and_count = iree_hal_cuda2_nccl_channel_query_rank_and_count,
+};
diff --git a/experimental/cuda2/nccl_channel.h b/experimental/cuda2/nccl_channel.h
new file mode 100644
index 0000000..a6da61f
--- /dev/null
+++ b/experimental/cuda2/nccl_channel.h
@@ -0,0 +1,57 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef EXPERIMENTAL_CUDA2_NCCL_CHANNEL_H_
+#define EXPERIMENTAL_CUDA2_NCCL_CHANNEL_H_
+
+#include "experimental/cuda2/api.h"
+#include "experimental/cuda2/cuda_dynamic_symbols.h"
+#include "experimental/cuda2/cuda_headers.h"
+#include "experimental/cuda2/nccl_dynamic_symbols.h"
+#include "iree/base/api.h"
+#include "iree/hal/api.h"
+#include "iree/hal/utils/collective_batch.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif // __cplusplus
+
+// Returns true if |id| is all zeros indicating an empty ID.
+static inline bool iree_hal_cuda2_nccl_id_is_empty(
+ const iree_hal_cuda2_nccl_id_t* id) {
+ for (iree_host_size_t i = 0; i < IREE_ARRAYSIZE(id->data); ++i) {
+ if (id->data[i] != 0) return false;
+ }
+ return true;
+}
+
+// Gets a unique ID to bootstrap a new communicator. It calls ncclGetUniqueId
+// under the hood.
+iree_status_t iree_hal_cuda2_nccl_get_unique_id(
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* symbols,
+ iree_hal_cuda2_nccl_id_t* out_id);
+
+// Creates a IREE HAL channel using the given NCCL |id|, |rank|, and |count|.
+// It calls ncclCommInitRankConfig under the hood.
+iree_status_t iree_hal_cuda2_nccl_channel_create(
+ const iree_hal_cuda2_dynamic_symbols_t* cuda_symbols,
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols,
+ const iree_hal_cuda2_nccl_id_t* id, int rank, int count,
+ iree_allocator_t host_allocator, iree_hal_channel_t** out_channel);
+
+// Performs a non-blocking submission of |batch| to |stream|.
+// The backing storage of |batch| is dropped immediately but all resources
+// referenced will be retained by the parent command buffer for its lifetime.
+// Note that operations in the batch may apply to different channels.
+iree_status_t iree_hal_cuda2_nccl_submit_batch(
+ const iree_hal_cuda2_nccl_dynamic_symbols_t* nccl_symbols,
+ const iree_hal_collective_batch_t* batch, CUstream stream);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif // __cplusplus
+
+#endif // EXPERIMENTAL_CUDA2_NCCL_CHANNEL_H_
diff --git a/experimental/cuda2/nccl_status_util.c b/experimental/cuda2/nccl_status_util.c
index c16e9f9..81ca097 100644
--- a/experimental/cuda2/nccl_status_util.c
+++ b/experimental/cuda2/nccl_status_util.c
@@ -11,7 +11,7 @@
#include "experimental/cuda2/nccl_dynamic_symbols.h"
#include "iree/base/status.h"
-iree_status_t iree_hal_nccl2_result_to_status(
+iree_status_t iree_hal_cuda2_nccl_result_to_status(
const iree_hal_cuda2_nccl_dynamic_symbols_t* syms, ncclResult_t result,
const char* file, uint32_t line) {
iree_status_code_t code;
diff --git a/experimental/cuda2/nccl_status_util.h b/experimental/cuda2/nccl_status_util.h
index 2826b1d..65db9a9 100644
--- a/experimental/cuda2/nccl_status_util.h
+++ b/experimental/cuda2/nccl_status_util.h
@@ -21,30 +21,44 @@
// Usage:
// iree_status_t status = IREE_NCCL_RESULT_TO_STATUS(nccl_symbols,
// ncclDoThing(...));
-#define IREE_NCCL_RESULT_TO_STATUS(syms, expr, ...) \
- iree_hal_nccl2_result_to_status((syms), ((syms)->expr), __FILE__, __LINE__)
+#define IREE_NCCL_RESULT_TO_STATUS(syms, expr, ...) \
+ iree_hal_cuda2_nccl_result_to_status((syms), ((syms)->expr), __FILE__, \
+ __LINE__)
// IREE_RETURN_IF_ERROR but implicitly converts the ncclResult_t return value to
// an iree_status_t.
//
// Usage:
// IREE_NCCL_RETURN_IF_ERROR(nccl_symbols, ncclDoThing(...), "message");
-#define IREE_NCCL_RETURN_IF_ERROR(syms, expr, ...) \
- IREE_RETURN_IF_ERROR(iree_hal_nccl2_result_to_status((syms), ((syms)->expr), \
- __FILE__, __LINE__), \
+#define IREE_NCCL_RETURN_IF_ERROR(syms, expr, ...) \
+ IREE_RETURN_IF_ERROR(iree_hal_cuda2_nccl_result_to_status( \
+ (syms), ((syms)->expr), __FILE__, __LINE__), \
__VA_ARGS__)
+// IREE_RETURN_IF_ERROR but ends the current zone and implicitly converts the
+// ncclResult_t return value to an iree_status_t.
+//
+// Usage:
+// IREE_NCCL_RETURN_AND_END_ZONE_IF_ERROR(zone_id, cuda_symbols,
+// ncclDoThing(...), "message");
+#define IREE_NCCL_RETURN_AND_END_ZONE_IF_ERROR(zone_id, syms, expr, ...) \
+ IREE_RETURN_AND_END_ZONE_IF_ERROR( \
+ zone_id, \
+ iree_hal_cuda2_nccl_result_to_status((syms), ((syms)->expr), __FILE__, \
+ __LINE__), \
+ __VA_ARGS__)
+
// IREE_IGNORE_ERROR but implicitly converts the ncclResult_t return value to
// an iree_status_t.
//
// Usage:
// IREE_NCCL_IGNORE_ERROR(nccl_symbols, ncclDoThing(...));
-#define IREE_NCCL_IGNORE_ERROR(syms, expr) \
- IREE_IGNORE_ERROR(iree_hal_nccl2_result_to_status((syms), ((syms)->expr), \
- __FILE__, __LINE__))
+#define IREE_NCCL_IGNORE_ERROR(syms, expr) \
+ IREE_IGNORE_ERROR(iree_hal_cuda2_nccl_result_to_status( \
+ (syms), ((syms)->expr), __FILE__, __LINE__))
// Converts a ncclResult_t to an iree_status_t object.
-iree_status_t iree_hal_nccl2_result_to_status(
+iree_status_t iree_hal_cuda2_nccl_result_to_status(
const iree_hal_cuda2_nccl_dynamic_symbols_t* syms, ncclResult_t result,
const char* file, uint32_t line);
