wip 2

pytorch · Nov 17, 2024 · 0e1c38e · 0e1c38e
1 parent a3cdb79
commit 0e1c38e
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 39 deletions.
diff --git a/cmake/modules/GpuCppLibrary.cmake b/cmake/modules/GpuCppLibrary.cmake
@@ -83,7 +83,10 @@ function(prepare_target_sources)
             ${${args_PREFIX}_sources_cpp})
     endif()
 
-    # Set the output variable
+    ############################################################################
+    # Set the Output Variable(s)
+    ############################################################################
+
     set(${args_PREFIX}_sources ${${args_PREFIX}_sources} PARENT_SCOPE)
 endfunction()
 
@@ -161,7 +164,7 @@ function(gpu_cpp_library)
     endif()
 
     ############################################################################
-    # Additional Library Includes and Linking
+    # Library Includes and Linking
     ############################################################################
 
     # Add PyTorch include/
@@ -189,9 +192,25 @@ function(gpu_cpp_library)
         -Wno-deprecated-anon-enum-enum-conversion
         -Wno-deprecated-declarations)
 
-    # Set the output library target
-    set(${args_PREFIX}_py ${lib_name} PARENT_SCOPE)
+    ############################################################################
+    # Post-Build Steps
+    ############################################################################
+
+    # Add a post-build step to remove errant RPATHs from the .SO
+    add_custom_target(${lib_name}_postbuild ALL
+        DEPENDS
+        WORKING_DIRECTORY ${OUTPUT_DIR}
+        COMMAND bash ${FBGEMM}/.github/scripts/fbgemm_gpu_postbuild.bash)
 
+    # Run the post-build steps AFTER the build itself
+    add_dependencies(${lib_name}_postbuild ${lib_name})
+
+    ############################################################################
+    # Set the Output Variable(s)
+    ############################################################################
+
+    # PREFIX = `foo` --> Target Library = `foo_py`
+    set(${args_PREFIX}_py ${lib_name} PARENT_SCOPE)
 
     BLOCK_PRINT(
         "#### GPU CPP Library Target ####"
@@ -205,6 +224,9 @@ function(gpu_cpp_library)
         "GPU_SRCS:"
         "${args_GPU_SRCS}"
         " "
+        "INCLUDE_DIRS:"
+        "${args_INCLUDE_DIRS}"
+        " "
         "Selected Source Files:"
         "${lib_sources}"
         " "

diff --git a/fbgemm_gpu/experimental/example/CMakeLists.txt b/fbgemm_gpu/experimental/example/CMakeLists.txt
@@ -10,31 +10,34 @@ include(${CMAKEMODULES}/Utilities.cmake)
 # Target Sources
 ################################################################################
 
-set(fbgemm_sources_include_directories
-  # FBGEMM
-  ${FBGEMM}/include
-  # FBGEMM_GPU
-  ${CMAKE_CURRENT_SOURCE_DIR}../..
-  ${CMAKE_CURRENT_SOURCE_DIR}../../include
-  ${CMAKE_CURRENT_SOURCE_DIR}../../../include
-  # PyTorch
-  ${TORCH_INCLUDE_DIRS}
-  # Third-party
-  ${THIRDPARTY}/asmjit/src
-  ${THIRDPARTY}/cpuinfo/include
-  ${THIRDPARTY}/cutlass/include
-  ${THIRDPARTY}/cutlass/tools/util/include
-  ${THIRDPARTY}/json/include
-  ${NCCL_INCLUDE_DIRS})
+# fbgemm_sources_include_directories
+# gpu_cpp_library
+
+# set(fbgemm_sources_include_directories
+#   # FBGEMM
+#   ${FBGEMM}/include
+#   # FBGEMM_GPU
+#   ${CMAKE_CURRENT_SOURCE_DIR}../..
+#   ${CMAKE_CURRENT_SOURCE_DIR}../../include
+#   ${CMAKE_CURRENT_SOURCE_DIR}../../../include
+#   # PyTorch
+#   ${TORCH_INCLUDE_DIRS}
+#   # Third-party
+#   ${THIRDPARTY}/asmjit/src
+#   ${THIRDPARTY}/cpuinfo/include
+#   ${THIRDPARTY}/cutlass/include
+#   ${THIRDPARTY}/cutlass/tools/util/include
+#   ${THIRDPARTY}/json/include
+#   ${NCCL_INCLUDE_DIRS})
 
 set(experimental_example_cpp_source_files
     src/cutlass_sgemm_nn.cu
     src/example_ops.cpp
     src/nccl_example.cpp)
 
-set_source_files_properties(${experimental_example_cpp_source_files}
-    PROPERTIES INCLUDE_DIRECTORIES
-    "${fbgemm_sources_include_directories}")
+# set_source_files_properties(${experimental_example_cpp_source_files}
+#     PROPERTIES INCLUDE_DIRECTORIES
+#     "${fbgemm_sources_include_directories}")
 
 set(experimental_example_python_source_files
     example/__init__.py
@@ -45,20 +48,29 @@ set(experimental_example_python_source_files
 # Build Shared Library
 ################################################################################
 
-add_library(fbgemm_gpu_experimental_example_py MODULE
-    ${experimental_example_cpp_source_files})
+gpu_cpp_library(
+    PREFIX
+        fbgemm_gpu_experimental_example
+    INCLUDE_DIRS
+        ${fbgemm_sources_include_directories}
+    GPU_SRCS
+        ${experimental_example_cpp_source_files})
 
-target_include_directories(fbgemm_gpu_experimental_example_py PRIVATE
-    ${TORCH_INCLUDE_DIRS}
-    ${NCCL_INCLUDE_DIRS})
 
-target_link_libraries(fbgemm_gpu_experimental_example_py
-    ${TORCH_LIBRARIES}
-    ${NCCL_LIBRARIES}
-    ${CUDA_DRIVER_LIBRARIES})
+# add_library(fbgemm_gpu_experimental_example_py MODULE
+#     ${experimental_example_cpp_source_files})
 
-# Remove `lib` from the output artifact name
-set_target_properties(fbgemm_gpu_experimental_example_py PROPERTIES PREFIX "")
+# target_include_directories(fbgemm_gpu_experimental_example_py PRIVATE
+#     ${TORCH_INCLUDE_DIRS}
+#     ${NCCL_INCLUDE_DIRS})
+
+# target_link_libraries(fbgemm_gpu_experimental_example_py
+#     ${TORCH_LIBRARIES}
+#     ${NCCL_LIBRARIES}
+#     ${CUDA_DRIVER_LIBRARIES})
+
+# # Remove `lib` from the output artifact name
+# set_target_properties(fbgemm_gpu_experimental_example_py PROPERTIES PREFIX "")
 
 
 ################################################################################
@@ -71,9 +83,9 @@ install(TARGETS fbgemm_gpu_experimental_example_py
 install(FILES ${experimental_example_python_source_files}
         DESTINATION fbgemm_gpu/experimental/example)
 
-add_custom_target(fbgemm_gpu_experimental_example_py_clean_rpath ALL
-    WORKING_DIRECTORY ${OUTPUT_DIR}
-    COMMAND bash ${FBGEMM}/.github/scripts/fbgemm_gpu_postbuild.bash)
+# add_custom_target(fbgemm_gpu_experimental_example_py_clean_rpath ALL
+#     WORKING_DIRECTORY ${OUTPUT_DIR}
+#     COMMAND bash ${FBGEMM}/.github/scripts/fbgemm_gpu_postbuild.bash)
 
-add_dependencies(fbgemm_gpu_experimental_example_py_clean_rpath
-    fbgemm_gpu_experimental_example_py)
+# add_dependencies(fbgemm_gpu_experimental_example_py_clean_rpath
+#     fbgemm_gpu_experimental_example_py)