wip

cmake/modules/GpuCppLibrary.cmake

@@ -96,71 +96,53 @@ function(prepare_hipified_target_sources)
     set(${args_PREFIX}_sources_hipified ${args_SRCS} PARENT_SCOPE)
 endfunction()
 
-
 function(gpu_cpp_library)
     set(flags)
     set(singleValueArgs PREFIX)
-    set(multiValueArgs CPU_SRCS GPU_CXX_SRCS GPU_CU_SRCS)
+    set(multiValueArgs CPU_SRCS GPU_CXX_SRCS GPU_CU_SRCS OTHER_SRCS)
     cmake_parse_arguments(
         args
         "${flags}" "${singleValueArgs}" "${multiValueArgs}"
         ${ARGN})
 
-    BLOCK_PRINT(
-        "ARGS"
-        "${args_PREFIX}"
-        "===="
-        "${args_CPU_SRCS}"
-        "===="
-        "${args_GPU_CXX_SRCS}"
-        "===="
-        "${args_GPU_CU_SRCS}"
-        "===="
-    )
-
     prepare_target_sources(
         PREFIX ${args_PREFIX}
         CPU_SRCS ${args_CPU_SRCS}
         GPU_CXX_SRCS ${args_GPU_CXX_SRCS}
         GPU_CU_SRCS ${args_GPU_CU_SRCS})
     set(lib_sources ${${args_PREFIX}_sources_static})
 
-    BLOCK_PRINT(
-        "LOCAL SOURCES"
-        "===="
-        "${fbgemm_sources_include_directories}"
-        "===="
-        "${lib_sources}"
-        "===="
-    )
 
     set(lib_name ${args_PREFIX}_py)
 
     if(USE_ROCM)
+        # Fetch the HIPified sources
         prepare_hipified_target_sources(
             PREFIX ${args_PREFIX}
             SRCS ${lib_sources})
         set(lib_sources_hipified ${${args_PREFIX}_sources_hipified})
 
+        # Create the HIP library
         hip_add_library(${lib_name} SHARED
             ${lib_sources_hipified}
+            ${args_OTHER_SRCS}
             ${FBGEMM_HIP_HCC_LIBRARIES}
             HIPCC_OPTIONS
             ${HIP_HCC_FLAGS})
 
+        # Append ROCM includes
         target_include_directories(${lib_name} PUBLIC
             ${FBGEMM_HIP_INCLUDE}
             ${ROCRAND_INCLUDE}
             ${ROCM_SMI_INCLUDE})
 
-        list(GET TORCH_INCLUDE_DIRS 0 TORCH_PATH)
-
     else()
+        # Create the C++/CUDA library
         add_library(${lib_name} MODULE
-            ${lib_sources_hipified})
+            ${lib_sources}
+            ${args_OTHER_SRCS})
     endif()
 
-
     # Add PyTorch include/
     target_include_directories(${lib_name} PRIVATE
         ${TORCH_INCLUDE_DIRS}
@@ -176,15 +158,42 @@ function(gpu_cpp_library)
         ${NCCL_LIBRARIES}
         ${CUDA_DRIVER_LIBRARIES})
 
-    # Link to NVML
+    # Link to NVML if available
     if(NVML_LIB_PATH)
         target_link_libraries(${lib_name} ${NVML_LIB_PATH})
     endif()
 
-    # Silence warnings in asmjit
+    # Silence warnings (in asmjit)
     target_compile_options(${lib_name} PRIVATE
         -Wno-deprecated-anon-enum-enum-conversion
         -Wno-deprecated-declarations)
 
+    # Set the output library target
     set(${args_PREFIX}_py ${lib_name} PARENT_SCOPE)
+
+
+    BLOCK_PRINT(
+        "#### GPU CPP Library Target ####"
+        " "
+        "PREFIX:"
+        "${args_PREFIX}"
+        " "
+        "CPU_SRCS:"
+        "${args_CPU_SRCS}"
+        " "
+        "GPU_CXX_SRCS:"
+        "${args_GPU_CXX_SRCS}"
+        " "
+        "GPU_CU_SRCS:"
+        "${args_GPU_CU_SRCS}"
+        " "
+        "Selected Source Files:"
+        "${lib_sources}"
+        " "
+        "HIPified Source Files:"
+        "${lib_sources_hipified}"
+        " "
+        "Output Library:"
+        "${lib_name}"
+    )
 endfunction()