dynamically link NVSHMEM

Author: abcdabcd987

csrc/CMakeLists.txt

@@ -65,7 +65,8 @@ target_link_libraries(pplx_kernels PUBLIC
     Python::Module
     CUDA::cuda_driver
     CUDA::cudart
-    nvshmem::nvshmem
+    nvshmem::nvshmem_host
+    nvshmem::nvshmem_device
 )
 set_target_properties(pplx_kernels PROPERTIES
     LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/../src/pplx_kernels

csrc/all_to_all/CMakeLists.txt

@@ -7,7 +7,8 @@ add_library(all_to_all_lib STATIC
     internode.cpp
 )
 target_link_libraries(all_to_all_lib PUBLIC
-    nvshmem::nvshmem
+    nvshmem::nvshmem_host
+    nvshmem::nvshmem_device
     CUDA::cudart
 )
 set_cuda_compile_options(all_to_all_lib)
@@ -22,7 +23,8 @@ if(WITH_TESTS)
         CUDA::cudart
         CUDA::cuda_driver
         MPI::MPI_CXX
-        nvshmem::nvshmem
+        nvshmem::nvshmem_host
+        nvshmem::nvshmem_device
     )
     set_cuda_compile_options(test_all_to_all)
     add_test(NAME AllToAllTest
@@ -39,6 +41,7 @@ if (WITH_BENCHMARKS)
         CUDA::cudart
         CUDA::cuda_driver
         MPI::MPI_CXX
-        nvshmem::nvshmem
+        nvshmem::nvshmem_host
+        nvshmem::nvshmem_device
     )
 endif()

csrc/bindings/nvshmem.cpp

@@ -28,7 +28,6 @@ int64_t init(at::Tensor uid, int64_t rank, int64_t world_size) {
   TORCH_CHECK(
       uid.numel() == sizeof(nvshmemx_uniqueid_t),
       "Invalid unique id size (expected ",
-
       sizeof(nvshmemx_uniqueid_t),
       ", got ",
       uid.numel(),