diff --git a/include/internal/cudecomp_kernels.cuh b/include/internal/cudecomp_kernels.cuh index 66138fd..e38abe3 100644 --- a/include/internal/cudecomp_kernels.cuh +++ b/include/internal/cudecomp_kernels.cuh @@ -22,12 +22,23 @@ #include #include +// Keep NVSHMEM device headers out of non-RDC CUDA translation units. +#if defined(ENABLE_NVSHMEM) && defined(CUDECOMP_NVSHMEM_HOST_ONLY) && !defined(NVSHMEM_HOSTLIB_ONLY) +#define CUDECOMP_UNDEF_NVSHMEM_HOSTLIB_ONLY +#define NVSHMEM_HOSTLIB_ONLY +#endif + #ifdef ENABLE_NVSHMEM #include #endif #include "internal/cudecomp_kernels.h" +#ifdef CUDECOMP_UNDEF_NVSHMEM_HOSTLIB_ONLY +#undef NVSHMEM_HOSTLIB_ONLY +#undef CUDECOMP_UNDEF_NVSHMEM_HOSTLIB_ONLY +#endif + #define CUDECOMP_CUDA_NTHREADS (128) #define CUDECOMP_UNROLL_FACTOR (4) #define CUDECOMP_MIN_BLOCKS_PER_SM (16) @@ -36,7 +47,7 @@ namespace cudecomp { -#ifdef ENABLE_NVSHMEM +#if defined(ENABLE_NVSHMEM) && !defined(CUDECOMP_NVSHMEM_HOST_ONLY) template __launch_bounds__(CUDECOMP_CUDA_NTHREADS) __global__ void cudecomp_nvshmem_alltoallv_k(cudecompNvshmemA2AParams params) { diff --git a/src/cudecomp_kernels.cu b/src/cudecomp_kernels.cu index 311ebb9..16a77b5 100644 --- a/src/cudecomp_kernels.cu +++ b/src/cudecomp_kernels.cu @@ -18,7 +18,11 @@ #include #include "internal/checks.h" + +// This non-RDC translation unit only instantiates generic CUDA kernels. +#define CUDECOMP_NVSHMEM_HOST_ONLY #include "internal/cudecomp_kernels.cuh" +#undef CUDECOMP_NVSHMEM_HOST_ONLY namespace cudecomp {