[fbgemm_gpu] Fix CMakeLists.txt for experimental/gemm

cmake/modules/Utilities.cmake

@@ -74,8 +74,13 @@ endmacro()
 
 function(add_to_package)
   set(flags)
-  set(singleValueArgs DESTINATION)
-  set(multiValueArgs FILES TARGETS)
+  set(singleValueArgs
+    DESTINATION       # The destination directory, RELATIVE to the root of the installation package directory
+  )
+  set(multiValueArgs
+    FILES             # The list of files to place into the DESTINATION directory
+    TARGETS           # THe list of CMake targets whose build artifacts to place into the DESTINATION directory
+  )
 
   cmake_parse_arguments(
     args

fbgemm_gpu/CMakeLists.txt

@@ -236,11 +236,11 @@ endif()
 ################################################################################
 
 if(NOT FBGEMM_CPU_ONLY AND NOT USE_ROCM)
-  # TODO: Figure out NCCL/RCCL integration with ROCm
   add_subdirectory(experimental/example)
 endif()
 
 if(NOT FBGEMM_CPU_ONLY)
+  # Package Triton GEMM (GenAI) kernels
   add_subdirectory(experimental/gemm)
 endif()
 

fbgemm_gpu/experimental/example/CMakeLists.txt

@@ -8,35 +8,33 @@
 # Target Sources
 ################################################################################
 
-set(experimental_example_cpp_source_files
-    src/cutlass_sgemm_nn.cu
-    src/example_ops.cpp
-    src/nccl_example.cpp)
+file(GLOB_RECURSE experimental_example_cpp_source_files
+  src/*.cu
+  src/*.cpp)
 
-set(experimental_example_python_source_files
-    example/__init__.py
-    example/utils.py)
+file(GLOB_RECURSE experimental_example_python_source_files
+  example/*.py)
 
 
 ################################################################################
 # Build Shared Library
 ################################################################################
 
 gpu_cpp_library(
-    PREFIX
-        fbgemm_gpu_experimental_example_py
-    TYPE
-        SHARED
-    INCLUDE_DIRS
-        ${fbgemm_sources_include_directories}
-    GPU_SRCS
-        ${experimental_example_cpp_source_files})
+  PREFIX
+    fbgemm_gpu_experimental_example_py
+  TYPE
+    SHARED
+  INCLUDE_DIRS
+    ${fbgemm_sources_include_directories}
+  GPU_SRCS
+    ${experimental_example_cpp_source_files})
 
 
 ################################################################################
 # Install Shared Library and Python Files
 ################################################################################
 
 add_to_package(DESTINATION fbgemm_gpu/experimental/example
-    TARGETS fbgemm_gpu_experimental_example_py
-    FILES ${experimental_example_python_source_files})
+  TARGETS fbgemm_gpu_experimental_example_py
+  FILES ${experimental_example_python_source_files})

fbgemm_gpu/experimental/gemm/CMakeLists.txt

@@ -8,15 +8,15 @@
 # Target Sources
 ################################################################################
 
-set(experimental_triton_python_source_files
-    triton_gemm/__init__.py
-    triton_gemm/fp8_gemm.py
-    triton_gemm/matmul_perf_model.py)
+# Python sources
+file(GLOB_RECURSE experimental_triton_gemm_python_source_files
+  triton_gemm/*.py)
 
 
 ################################################################################
 # Install Python Files
 ################################################################################
 
-add_to_package(DESTINATION fbgemm_gpu/experimental/gemm/triton_gemm
-    FILES ${experimental_triton_python_source_files})
+add_to_package(
+  DESTINATION fbgemm_gpu/experimental/gemm/triton_gemm
+  FILES ${experimental_triton_gemm_python_source_files})

fbgemm_gpu/experimental/gen_ai/CMakeLists.txt

@@ -42,8 +42,7 @@ file(GLOB_RECURSE experimental_gen_ai_cpp_source_files_hip
 
 # Python sources
 file(GLOB_RECURSE experimental_gen_ai_python_source_files
-  RELATIVE gen_ai
-  *.py)
+  gen_ai/*.py)
 
 
 ################################################################################
@@ -52,7 +51,7 @@ file(GLOB_RECURSE experimental_gen_ai_python_source_files
 
 gpu_cpp_library(
   PREFIX
-    fbgemm_gpu_experimental_gen_ai_py
+    fbgemm_gpu_experimental_gen_ai
   TYPE
     SHARED
   INCLUDE_DIRS
@@ -72,5 +71,5 @@ gpu_cpp_library(
 
 add_to_package(
   DESTINATION fbgemm_gpu/experimental/gen_ai
-  TARGETS fbgemm_gpu_experimental_gen_ai_py
+  TARGETS fbgemm_gpu_experimental_gen_ai
   FILES ${experimental_gen_ai_python_source_files})

fbgemm_gpu/experimental/gen_ai/gen_ai/__init__.py

@@ -25,10 +25,10 @@
 # pyre-ignore[16]
 if open_source:
     torch.ops.load_library(
-        os.path.join(os.path.dirname(__file__), "fbgemm_gpu_experimental_gen_ai_py.so")
+        os.path.join(os.path.dirname(__file__), "fbgemm_gpu_experimental_gen_ai.so")
     )
     torch.classes.load_library(
-        os.path.join(os.path.dirname(__file__), "fbgemm_gpu_experimental_gen_ai_py.so")
+        os.path.join(os.path.dirname(__file__), "fbgemm_gpu_experimental_gen_ai.so")
     )
 else:
     torch.ops.load_library(

fbgemm_gpu/fbgemm_gpu/__init__.py

@@ -46,15 +46,24 @@ def _load_library(filename: str) -> None:
     "fbgemm_gpu_py",
 ]
 
+fbgemm_gpu_genai_libraries = [
+    "experimental/gen_ai/fbgemm_gpu_experimental_gen_ai",
+]
+
+# NOTE: While FBGEMM_GPU GenAI is not available for ROCm yet, we would like to
+# be able to install the existing CUDA variant of the package onto ROCm systems,
+# so that we can at least use the Triton GEMM libraries from experimental/gemm.
+# But loading fbgemm_gpu package will trigger load-checking the .SO file for the
+# GenAI libraries, which will fail.  This workaround ignores check-loading the
+# .SO file for the ROCm case, so that clients can import
+# fbgemm_gpu.experimental.gemm without triggering an error.
+if torch.cuda.is_available() and torch.version.hip:
+    fbgemm_gpu_genai_libraries = []
+
 libraries_to_load = {
     "cpu": fbgemm_gpu_libraries,
-    "cuda": fbgemm_gpu_libraries
-    + [
-        "experimental/gen_ai/fbgemm_gpu_experimental_gen_ai_py",
-    ],
-    "genai": [
-        "experimental/gen_ai/fbgemm_gpu_experimental_gen_ai_py",
-    ],
+    "cuda": fbgemm_gpu_libraries + fbgemm_gpu_genai_libraries,
+    "genai": fbgemm_gpu_genai_libraries,
     "rocm": fbgemm_gpu_libraries,
 }