Build changes to support integration with pytorch (#4354)

Summary:
Pull Request resolved: #4354

We plan to start integration of FBGEMM GenAI build into pytorch directly. This change would make the FBGEMM CMake compatible with building FBGEMM GenAI as part of pytorch as a dependency. Later we will have a PR in pytorch to update the pinned fbgemm and update their build.

Changes for torch:
- Allow option to keep `lib` prefix, to keep it more standard with other libs
- Allow passing in torch cmake targets so we can link against them in fbgemm
- Use the library name `libfbgemm_gpu_gen_ai.so` (without experimental) as it felt a bit wierd to keep that in.
- We will export the target so torch can later link with `libfbgemm_gpu_gen_ai.so`

X-link: facebookresearch/FBGEMM#1422

Test Plan:
Built locally and confirmed things work:

```
test_fbgemm_gpu_build_and_install $CONDA_ENV genai/cuda
```

```
test_fbgemm_gpu_build_and_install $CONDA_ENV default/cuda
```

Will let OSS CI run for more comprehensive test.

Reviewed By: q10

Differential Revision: D76734540

Pulled By: cthi

fbshipit-source-id: c138ec93d3e78d93c77ade4d1b8b8ed902489db3

Author: cthi

cmake/modules/CudaSetup.cmake

@@ -4,9 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/modules/Utilities.cmake)
-
-
 ################################################################################
 # CUDA Setup
 ################################################################################

cmake/modules/GpuCppLibrary.cmake

@@ -4,8 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/modules/Utilities.cmake)
-
 function(prepare_target_sources)
     # This function does the following:
     #
@@ -150,6 +148,7 @@ function(gpu_cpp_library)
         PREFIX          # Desired name for the library target (and by extension, the prefix for naming intermediate targets)
         TYPE            # Target type, e.g., MODULE, OBJECT.  See https://cmake.org/cmake/help/latest/command/add_library.html
         DESTINATION     # The install destination directory to place the build target into
+        KEEP_PREFIX     # Whether to keep the prefix for the library target, e.g. libfoo.so vs foo.so
     )
     set(multiValueArgs
         CPU_SRCS            # Sources for CPU-only build
@@ -162,6 +161,7 @@ function(gpu_cpp_library)
         HIPCC_FLAGS         # Compilation flags specific to HIPCC
         INCLUDE_DIRS        # Include directories for compilation
         DEPS                # Target dependencies, i.e. built STATIC targets
+        TORCH_LIBS          # PyTorch libraries to link against. Note that we provide the TORCH_LIBS automatically - this is for PyTorch build.
     )
 
     cmake_parse_arguments(
@@ -258,19 +258,23 @@ function(gpu_cpp_library)
         ${NCCL_INCLUDE_DIRS})
 
     # Set additional target properties
-    set_target_properties(${lib_name} PROPERTIES
+    if(NOT args_KEEP_PREFIX)
         # Remove `lib` prefix from the output artifact name, e.g.
         # `libfoo.so` -> `foo.so`
-        PREFIX ""
+        set_target_properties(${lib_name} PROPERTIES PREFIX "")
+    endif()
+
+    set_target_properties(${lib_name} PROPERTIES
         # Enforce -fPIC for STATIC library option, since they are to be
         # integrated into other libraries down the line
         # https://stackoverflow.com/questions/3961446/why-does-gcc-not-implicitly-supply-the-fpic-flag-when-compiling-static-librarie
         POSITION_INDEPENDENT_CODE ON)
 
-    if (args_DEPS)
+    if (args_DEPS OR CMAKE_INSTALL_RPATH)
         # Only set this if the library has dependencies that we also build,
         # otherwise we will hit the following error:
         #   `No valid ELF RPATH or RUNPATH entry exists in the file`
+        # However, if CMAKE_INSTALL_RPATH is set, respect that logic. Such as when we build with PyTorch.
         set_target_properties(${lib_name} PROPERTIES
             BUILD_WITH_INSTALL_RPATH ON
             # Set the RPATH for the library to include $ORIGIN, so it can look
@@ -288,6 +292,7 @@ function(gpu_cpp_library)
     # Collect external libraries for linking
     set(library_dependencies
         ${TORCH_LIBRARIES}
+        ${args_TORCH_LIBS}
         ${NCCL_LIBRARIES}
         ${CUDA_DRIVER_LIBRARIES}
         ${args_DEPS})
@@ -315,7 +320,7 @@ function(gpu_cpp_library)
     # Post-Build Steps
     ############################################################################
 
-    if (args_DEPS)
+    if (args_DEPS OR CMAKE_INSTALL_RPATH)
         # Only set this if the library has dependencies that we also build,
         # otherwise we will hit the following error:
         #   `No valid ELF RPATH or RUNPATH entry exists in the file`
@@ -342,7 +347,10 @@ function(gpu_cpp_library)
     ############################################################################
 
     if(args_DESTINATION)
-        install(TARGETS ${args_PREFIX}
+        install(
+            TARGETS ${args_PREFIX}
+            # Allows args_PREFIX to be exported as a target, used for PyTorch build integration
+            EXPORT fbgemmGenAILibraryConfig
             DESTINATION ${args_DESTINATION})
     endif()
 

cmake/modules/PyTorchSetup.cmake

@@ -4,9 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/modules/Utilities.cmake)
-
-
 ################################################################################
 # PyTorch Dependencies Setup
 ################################################################################

cmake/modules/RocmSetup.cmake

@@ -4,9 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/modules/Utilities.cmake)
-
-
 ################################################################################
 # ROCm and HIPify Setup
 ################################################################################

fbgemm_gpu/experimental/gen_ai/CMakeLists.txt

@@ -4,6 +4,24 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+################################################################################
+# CMake Prelude
+################################################################################
+
+cmake_minimum_required(VERSION 3.21 FATAL_ERROR)
+
+set(CMAKEMODULES ${CMAKE_CURRENT_SOURCE_DIR}/../../../cmake/modules)
+include(${CMAKEMODULES}/Utilities.cmake)
+include(${CMAKEMODULES}/GpuCppLibrary.cmake)
+
+# CUDA Setup
+include(${CMAKEMODULES}/CudaSetup.cmake)
+
+# ROCm and HIPify Setup
+include(${CMAKEMODULES}/RocmSetup.cmake)
+
+set(CMAKE_VERBOSE_MAKEFILE ON)
+
 ################################################################################
 # Target Sources
 ################################################################################
@@ -91,24 +109,48 @@ list(FILTER experimental_gen_ai_cpp_source_files_hip
 # Build Shared Library
 ################################################################################
 
+# Customizations for PyTorch build integration
+IF(DEFINED CAFFE2_THIRD_PARTY_ROOT)
+  # Remove "experimental" from the library name
+  set(FBGEMM_GENAI_LIB_NAME fbgemm_gpu_gen_ai)
+  # Keep lib prefix
+  set(KEEP_PREFIX ON)
+  # Install the library in the build directory
+  set(DESTINATION ${CMAKE_INSTALL_LIBDIR})
+else()
+  set(FBGEMM_GENAI_LIB_NAME fbgemm_gpu_experimental_gen_ai)
+  set(KEEP_PREFIX OFF)
+  set(DESTINATION "")
+endif()
+
 gpu_cpp_library(
   PREFIX
-    fbgemm_gpu_experimental_gen_ai
+    ${FBGEMM_GENAI_LIB_NAME}
   TYPE
     SHARED
+  DESTINATION
+    ${DESTINATION}
+  KEEP_PREFIX
+    ${KEEP_PREFIX}
   INCLUDE_DIRS
     ${fbgemm_sources_include_directories}
     ${CMAKE_CURRENT_SOURCE_DIR}/src/quantize
     ${CMAKE_CURRENT_SOURCE_DIR}/src/quantize/common/include
     ${CMAKE_CURRENT_SOURCE_DIR}/src/kv_cache
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../include
+    ${FBGEMM_GENAI_INCLUDE_DIRS}
   CPU_SRCS
     ${experimental_gen_ai_cpp_source_files_cpu}
   GPU_SRCS
     ${experimental_gen_ai_cpp_source_files_gpu}
   CUDA_SPECIFIC_SRCS
     ${experimental_gen_ai_cpp_source_files_cuda}
   HIP_SPECIFIC_SRCS
-    ${experimental_gen_ai_cpp_source_files_hip})
+    ${experimental_gen_ai_cpp_source_files_hip}
+  TORCH_LIBS
+    # Used when building as part of PyTorch
+    ${FBGEMM_GENAI_TORCH_LIBS})
+
 
 
 ################################################################################
@@ -117,7 +159,7 @@ gpu_cpp_library(
 
 add_to_package(
   DESTINATION fbgemm_gpu/experimental/gen_ai
-  TARGETS fbgemm_gpu_experimental_gen_ai)
+  TARGETS ${FBGEMM_GENAI_LIB_NAME})
 
 install(
   DIRECTORY bench