Use cuGetProcAddress to load CUDA entry points (#14056)
Sometimes the v2 APIs can have different function signatures which will
lead to errors if it doesn't match what we use in the IREE code. V2
lookup was originally added here
https://github.com/openxla/iree/commit/9f42d0448af7f80700cd613d69f8f4cc67583e1b
NVIDIA recommends using `cuGetProcAddress` to avoid version mixing which
can otherwise cause undefined behavior:
https://docs.nvidia.com/cuda/cuda-driver-api/version-mixing-rules.html#version-mixing-rules
diff --git a/experimental/cuda2/cuda_dynamic_symbol_table.h b/experimental/cuda2/cuda_dynamic_symbol_table.h
index fb8ff5a..aa030c9 100644
--- a/experimental/cuda2/cuda_dynamic_symbol_table.h
+++ b/experimental/cuda2/cuda_dynamic_symbol_table.h
@@ -4,6 +4,10 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// cuGetErrorName, cuGetErrorString should be loaded first in case there are
+// errors loading the other symbols.
+IREE_CU_PFN_DECL(cuGetErrorName, CUresult, const char**)
+IREE_CU_PFN_DECL(cuGetErrorString, CUresult, const char**)
IREE_CU_PFN_DECL(cuDriverGetVersion, int*)
IREE_CU_PFN_DECL(cuCtxCreate, CUcontext*, unsigned int, CUdevice)
IREE_CU_PFN_DECL(cuCtxDestroy, CUcontext)
@@ -23,8 +27,7 @@
IREE_CU_PFN_DECL(cuEventQuery, CUevent)
IREE_CU_PFN_DECL(cuEventRecord, CUevent, CUstream)
IREE_CU_PFN_DECL(cuEventSynchronize, CUevent)
-IREE_CU_PFN_DECL(cuGetErrorName, CUresult, const char**)
-IREE_CU_PFN_DECL(cuGetErrorString, CUresult, const char**)
+IREE_CU_PFN_DECL(cuGetProcAddress, const char*, void**, int, cuuint64_t)
IREE_CU_PFN_DECL(cuGraphAddMemcpyNode, CUgraphNode*, CUgraph,
const CUgraphNode*, size_t, const CUDA_MEMCPY3D*, CUcontext)
IREE_CU_PFN_DECL(cuGraphAddMemsetNode, CUgraphNode*, CUgraph,
@@ -76,8 +79,7 @@
IREE_CU_PFN_DECL(cuMemsetD8Async, unsigned long long, unsigned char, size_t,
CUstream)
IREE_CU_PFN_DECL(cuMemcpyAsync, CUdeviceptr, CUdeviceptr, size_t, CUstream)
-IREE_CU_PFN_DECL(cuMemcpyHtoDAsync_v2, CUdeviceptr, const void*, size_t,
- CUstream)
+IREE_CU_PFN_DECL(cuMemcpyHtoDAsync, CUdeviceptr, const void*, size_t, CUstream)
IREE_CU_PFN_DECL(cuFuncSetAttribute, CUfunction, CUfunction_attribute, int)
IREE_CU_PFN_DECL(cuLaunchKernel, CUfunction, unsigned int, unsigned int,
unsigned int, unsigned int, unsigned int, unsigned int,
diff --git a/experimental/cuda2/cuda_dynamic_symbols.c b/experimental/cuda2/cuda_dynamic_symbols.c
index 530c5f8..fa65809 100644
--- a/experimental/cuda2/cuda_dynamic_symbols.c
+++ b/experimental/cuda2/cuda_dynamic_symbols.c
@@ -20,29 +20,34 @@
#endif // IREE_PLATFORM_WINDOWS
};
-#define IREE_CONCAT(A, B) A B
+// CUDA API version for cuGetProcAddress.
+// 1000 * major + 10 * minor
+#define IREE_CUDA_DRIVER_API_VERSION 11030
-// Resolves all CUDA dynamic symbols in `dynamic_symbol_tables.h`, prefer _v2
-// version if it exists.
+// Load CUDA entry points.
static iree_status_t iree_hal_cuda2_dynamic_symbols_resolve_all(
iree_hal_cuda2_dynamic_symbols_t* syms) {
+ // Since cuGetProcAddress is in the symbol table, it will be loaded again
+ // through cuGetProcAddress. cuGetProcAddress_v2 is added in CUDA 12.0 and has
+ // a new function signature. If IREE_CUDA_DRIVER_API_VERSION is increased to
+ // >=12.0, then make sure we are using the correct signature.
+ IREE_RETURN_IF_ERROR(iree_dynamic_library_lookup_symbol(
+ syms->dylib, "cuGetProcAddress", (void**)&syms->cuGetProcAddress));
#define IREE_CU_PFN_DECL(cuda_symbol_name, ...) \
{ \
static const char* name = #cuda_symbol_name; \
- IREE_RETURN_IF_ERROR(iree_dynamic_library_lookup_symbol( \
- syms->dylib, name, (void**)&syms->cuda_symbol_name)); \
- static const char* name_v2 = IREE_CONCAT(#cuda_symbol_name, "_v2"); \
- void* fptr_v2; \
- iree_dynamic_library_lookup_symbol(syms->dylib, name_v2, &fptr_v2); \
- if (fptr_v2) syms->cuda_symbol_name = fptr_v2; \
+ IREE_CUDA_RETURN_IF_ERROR( \
+ syms, \
+ cuGetProcAddress(name, (void**)&syms->cuda_symbol_name, \
+ IREE_CUDA_DRIVER_API_VERSION, \
+ CU_GET_PROC_ADDRESS_DEFAULT), \
+ "when resolving " #cuda_symbol_name " using cuGetProcAddress"); \
}
#include "experimental/cuda2/cuda_dynamic_symbol_table.h" // IWYU pragma: keep
#undef IREE_CU_PFN_DECL
return iree_ok_status();
}
-#undef IREE_CONCAT
-
iree_status_t iree_hal_cuda2_dynamic_symbols_initialize(
iree_allocator_t host_allocator,
iree_hal_cuda2_dynamic_symbols_t* out_syms) {
diff --git a/experimental/cuda2/cuda_status_util.c b/experimental/cuda2/cuda_status_util.c
index 6d392a4..c8d57b5 100644
--- a/experimental/cuda2/cuda_status_util.c
+++ b/experimental/cuda2/cuda_status_util.c
@@ -124,12 +124,14 @@
if (IREE_LIKELY(result == CUDA_SUCCESS)) return iree_ok_status();
const char* error_name = NULL;
- if (syms->cuGetErrorName(result, &error_name) != CUDA_SUCCESS) {
+ if (!syms->cuGetErrorName ||
+ syms->cuGetErrorName(result, &error_name) != CUDA_SUCCESS) {
error_name = "CUDA_ERROR_UNKNOWN";
}
const char* error_string = NULL;
- if (syms->cuGetErrorString(result, &error_string) != CUDA_SUCCESS) {
+ if (!syms->cuGetErrorString ||
+ syms->cuGetErrorString(result, &error_string) != CUDA_SUCCESS) {
error_string = "unknown error";
}
diff --git a/runtime/src/iree/hal/drivers/cuda/dynamic_symbol_tables.h b/runtime/src/iree/hal/drivers/cuda/dynamic_symbol_tables.h
index e2bb69a..2e9cebe 100644
--- a/runtime/src/iree/hal/drivers/cuda/dynamic_symbol_tables.h
+++ b/runtime/src/iree/hal/drivers/cuda/dynamic_symbol_tables.h
@@ -4,6 +4,10 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// cuGetErrorName, cuGetErrorString should be loaded first in case there are
+// errors loading the other symbols.
+CU_PFN_DECL(cuGetErrorName, CUresult, const char**)
+CU_PFN_DECL(cuGetErrorString, CUresult, const char**)
CU_PFN_DECL(cuCtxCreate, CUcontext*, unsigned int, CUdevice)
CU_PFN_DECL(cuCtxDestroy, CUcontext)
CU_PFN_DECL(cuDevicePrimaryCtxRetain, CUcontext*, CUdevice)
@@ -22,8 +26,7 @@
CU_PFN_DECL(cuEventQuery, CUevent)
CU_PFN_DECL(cuEventRecord, CUevent, CUstream)
CU_PFN_DECL(cuEventSynchronize, CUevent)
-CU_PFN_DECL(cuGetErrorName, CUresult, const char**)
-CU_PFN_DECL(cuGetErrorString, CUresult, const char**)
+CU_PFN_DECL(cuGetProcAddress, const char*, void**, int, cuuint64_t)
CU_PFN_DECL(cuGraphAddMemcpyNode, CUgraphNode*, CUgraph, const CUgraphNode*,
size_t, const CUDA_MEMCPY3D*, CUcontext)
CU_PFN_DECL(cuGraphAddMemsetNode, CUgraphNode*, CUgraph, const CUgraphNode*,
@@ -71,7 +74,7 @@
CU_PFN_DECL(cuMemsetD8Async, unsigned long long, unsigned char, size_t,
CUstream)
CU_PFN_DECL(cuMemcpyAsync, CUdeviceptr, CUdeviceptr, size_t, CUstream)
-CU_PFN_DECL(cuMemcpyHtoDAsync_v2, CUdeviceptr, const void*, size_t, CUstream)
+CU_PFN_DECL(cuMemcpyHtoDAsync, CUdeviceptr, const void*, size_t, CUstream)
CU_PFN_DECL(cuFuncSetAttribute, CUfunction, CUfunction_attribute, int)
CU_PFN_DECL(cuLaunchKernel, CUfunction, unsigned int, unsigned int,
unsigned int, unsigned int, unsigned int, unsigned int,
diff --git a/runtime/src/iree/hal/drivers/cuda/dynamic_symbols.c b/runtime/src/iree/hal/drivers/cuda/dynamic_symbols.c
index eb95e82..436ce82 100644
--- a/runtime/src/iree/hal/drivers/cuda/dynamic_symbols.c
+++ b/runtime/src/iree/hal/drivers/cuda/dynamic_symbols.c
@@ -27,20 +27,28 @@
#endif // IREE_PLATFORM_WINDOWS
};
-#define concat(A, B) A B
+// CUDA API version for cuGetProcAddress.
+// 1000 * major + 10 * minor
+#define IREE_CUDA_DRIVER_API_VERSION 11030
-// Load CUDA entry points, prefer _v2 version if it exists.
+// Load CUDA entry points.
static iree_status_t iree_hal_cuda_dynamic_symbols_resolve_all(
iree_hal_cuda_dynamic_symbols_t* syms) {
-#define CU_PFN_DECL(cudaSymbolName, ...) \
- { \
- static const char* kName = #cudaSymbolName; \
- IREE_RETURN_IF_ERROR(iree_dynamic_library_lookup_symbol( \
- syms->cuda_library, kName, (void**)&syms->cudaSymbolName)); \
- static const char* kNameV2 = concat(#cudaSymbolName, "_v2"); \
- void* funV2; \
- iree_dynamic_library_lookup_symbol(syms->cuda_library, kNameV2, &funV2); \
- if (funV2) syms->cudaSymbolName = funV2; \
+ // Since cuGetProcAddress is in the symbol table, it will be loaded again
+ // through cuGetProcAddress. cuGetProcAddress_v2 is added in CUDA 12.0 and has
+ // a new function signature. If IREE_CUDA_DRIVER_API_VERSION is increased to
+ // >=12.0, then make sure we are using the correct signature.
+ IREE_RETURN_IF_ERROR(iree_dynamic_library_lookup_symbol(
+ syms->cuda_library, "cuGetProcAddress", (void**)&syms->cuGetProcAddress));
+#define CU_PFN_DECL(cudaSymbolName, ...) \
+ { \
+ static const char* kName = #cudaSymbolName; \
+ CUDA_RETURN_IF_ERROR( \
+ syms, \
+ cuGetProcAddress(kName, (void**)&syms->cudaSymbolName, \
+ IREE_CUDA_DRIVER_API_VERSION, \
+ CU_GET_PROC_ADDRESS_DEFAULT), \
+ "when resolving " #cudaSymbolName " using cuGetProcAddress"); \
}
#define NCCL_PFN_DECL(ncclSymbolName, ...)
#define NCCL_PFN_DECL_STR_RETURN(ncclSymbolName, ...)
diff --git a/runtime/src/iree/hal/drivers/cuda/status_util.c b/runtime/src/iree/hal/drivers/cuda/status_util.c
index 2e5eb1b..10ea14d 100644
--- a/runtime/src/iree/hal/drivers/cuda/status_util.c
+++ b/runtime/src/iree/hal/drivers/cuda/status_util.c
@@ -18,12 +18,14 @@
}
const char* error_name = NULL;
- if (syms->cuGetErrorName(result, &error_name) != CUDA_SUCCESS) {
+ if (!syms->cuGetErrorName ||
+ syms->cuGetErrorName(result, &error_name) != CUDA_SUCCESS) {
error_name = "UNKNOWN";
}
const char* error_string = NULL;
- if (syms->cuGetErrorString(result, &error_string) != CUDA_SUCCESS) {
+ if (!syms->cuGetErrorString ||
+ syms->cuGetErrorString(result, &error_string) != CUDA_SUCCESS) {
error_string = "Unknown error.";
}
return iree_make_status_with_location(file, line, IREE_STATUS_INTERNAL,
diff --git a/runtime/src/iree/hal/drivers/cuda/stream_command_buffer.c b/runtime/src/iree/hal/drivers/cuda/stream_command_buffer.c
index 713079f..12b582c 100644
--- a/runtime/src/iree/hal/drivers/cuda/stream_command_buffer.c
+++ b/runtime/src/iree/hal/drivers/cuda/stream_command_buffer.c
@@ -364,8 +364,8 @@
iree_hal_buffer_byte_offset(target_buffer) + target_offset;
CUDA_RETURN_IF_ERROR(
command_buffer->context->syms,
- cuMemcpyHtoDAsync_v2(dst, src, length, command_buffer->stream),
- "cuMemcpyHtoDAsync_v2");
+ cuMemcpyHtoDAsync(dst, src, length, command_buffer->stream),
+ "cuMemcpyHtoDAsync");
return iree_ok_status();
}
