Fix nvshmem build (#2815)

Signed-off-by: Gaetan Lepage <gaetan@glepage.com>

Author: GaetanLepage

transformer_engine/common/nvshmem_api/CMakeLists.txt

@@ -16,12 +16,13 @@ set(NVSHMEMAPI_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}" PARENT_SCOPE)
 target_link_directories(nvshmemapi PUBLIC ${NVSHMEM_HOME}/lib)
 target_link_libraries(nvshmemapi PUBLIC -static-libstdc++ nvshmem_device nvshmem_host CUDA::nvml CUDA::cublas CUDA::cuda_driver)
 target_include_directories(nvshmemapi PRIVATE
-                           ${NVSHMEM_HOME}/include/)
+                           ${NVSHMEM_HOME}/include/
+                           ${CMAKE_CURRENT_SOURCE_DIR}/../include)
 target_include_directories(nvshmemapi PUBLIC
                            ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
                            "${CMAKE_CURRENT_SOURCE_DIR}")
 
 set_target_properties(nvshmemapi PROPERTIES
                       CUDA_STANDARD 17
                       POSITION_INDEPENDENT_CODE ON
-                      CUDA_SEPARABLE_COMPILATION ON)
+                      CUDA_SEPARABLE_COMPILATION ON)

transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu

@@ -15,6 +15,7 @@
 #include <sstream>
 #include <string>
 
+#include "../util/cuda_driver.h"
 #include "../util/logging.h"
 #include "nvshmem_waitkernel.h"