[WIP] Triton extensions research

scripts/compile-triton.sh

@@ -6,7 +6,9 @@ export PIP_DISABLE_PIP_VERSION_CHECK=1
 
 # Select what to build.
 BUILD_LLVM=false
+LLVM_SHARED=false
 BUILD_TRITON=false
+TRITON_SHARED=false
 CLEAN=false
 VENV=false
 CCACHE=false
@@ -16,10 +18,20 @@ for arg in "$@"; do
       BUILD_LLVM=true
       shift
       ;;
+    --llvm-shared)
+      BUILD_LLVM=true
+      LLVM_SHARED=true
+      shift
+      ;;
     --triton)
       BUILD_TRITON=true
       shift
       ;;
+    --triton-shared)
+      BUILD_TRITON=true
+      TRITON_SHARED=true
+      shift
+      ;;
     --clean)
       CLEAN=true
       shift
@@ -33,7 +45,7 @@ for arg in "$@"; do
       shift
       ;;
     --help)
-      echo "Example usage: ./compile-triton.sh [--llvm | --triton | --clean | --venv | --ccache]"
+      echo "Example usage: ./compile-triton.sh [--llvm | --llvm-shared | --triton | --triton-shared | --clean | --venv | --ccache]"
       exit 1
       ;;
     *)
@@ -124,6 +136,11 @@ build_llvm() {
     ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS -DCMAKE_CXX_COMPILER_LAUNCHER=ccache"
   fi
 
+  if [ "$LLVM_SHARED" = true ]
+  then
+    ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS -DBUILD_SHARED_LIBS=true"
+  fi
+
   if [ ! -d "$LLVM_PROJ_BUILD" ]
   then
     mkdir $LLVM_PROJ_BUILD
@@ -179,6 +196,10 @@ build_triton() {
   then
     export TRITON_BUILD_WITH_CCACHE=true
   fi
+  if [ "$TRITON_SHARED" = true ]
+  then
+    export BUILD_SHARED_LIBS=true
+  fi
 
   # Install triton and its dependencies.
   pip install -v -e '.[build,tests]'

setup.py

@@ -324,6 +324,9 @@ def build_extension(self, ext):
         else:
             cmake_args += ["-DLLVM_BUILD_SHARED_LIBS=0"]
 
+        if check_env_flag("BUILD_SHARED_LIBS"):
+            cmake_args += ["-DBUILD_SHARED_LIBS=1"]
+
         # Note that asan doesn't work with binaries that use the GPU, so this is
         # only useful for tools like triton-opt that don't run code on the GPU.
         #

third_party/intel/cmake/FindSPIRVToLLVMTranslator.cmake

@@ -24,7 +24,50 @@ if (NOT SPIRVToLLVMTranslator_FOUND)
             set(LLVM_DIR "${LLVM_LIBRARY_DIR}/cmake/llvm" CACHE PATH "Path to LLVM build dir " FORCE)
             set(LLVM_SPIRV_BUILD_EXTERNAL YES CACHE BOOL "Build SPIRV-LLVM Translator as external" FORCE)
 
-            FetchContent_MakeAvailable(spirv-llvm-translator)
+            # Populate and add SPIRV-LLVM-Translator in two steps rather than
+            # using FetchContent_MakeAvailable(). MakeAvailable() performs
+            # populate + add_subdirectory() in one call, which immediately
+            # configures the subproject. We need a hook to apply patches after
+            # the repo is fetched but before CMake processes the project's
+            # CMakeLists.txt (otherwise patches affecting the configuration
+            # would be applied too late).
+            FetchContent_Populate(spirv-llvm-translator)
+
+            # Apply a patch to the SPIRV-LLVM-Translator project to ensure that
+            # its configuration finds the LLVM config package (LLVMConfig.cmake)
+            # used by the main project. This avoids falling back to the find
+            # module (FindLLVM.cmake). We need this patch to support builds with
+            # LLVM shared libraries because FindLLVM.cmake always invokes
+            # `llvm-config --link-static` and requires that static LLVM
+            # libraries are available.
+            message(STATUS "Applying spirv-llvm-translator-llvm-config.patch to SPIRV-LLVM-Translator")
+            execute_process(
+                COMMAND git apply --check ${CMAKE_CURRENT_LIST_DIR}/spirv-llvm-translator-llvm-config.patch
+                WORKING_DIRECTORY ${spirv-llvm-translator_SOURCE_DIR}
+                ERROR_QUIET
+                RESULT_VARIABLE PATCH_RESULT
+            )
+            if(PATCH_RESULT EQUAL 0)
+                execute_process(
+                        COMMAND git apply ${CMAKE_CURRENT_LIST_DIR}/spirv-llvm-translator-llvm-config.patch
+                        WORKING_DIRECTORY ${spirv-llvm-translator_SOURCE_DIR}
+                        RESULT_VARIABLE PATCH_RESULT
+                )
+            else()
+                execute_process( # Check if the patch is already applied
+                        COMMAND git apply --reverse --check ${CMAKE_CURRENT_LIST_DIR}/spirv-llvm-translator-llvm-config.patch
+                        WORKING_DIRECTORY ${spirv-llvm-translator_SOURCE_DIR}
+                        RESULT_VARIABLE PATCH_RESULT
+                )
+            endif()
+            if(NOT PATCH_RESULT EQUAL 0)
+                message(FATAL_ERROR "Failed to apply spirv-llvm-translator-llvm-config.patch to SPIRV-LLVM-Translator")
+            endif()
+
+            add_subdirectory(
+                ${spirv-llvm-translator_SOURCE_DIR}
+                ${spirv-llvm-translator_BINARY_DIR}
+            )
 
             # FIXME: Don't apply patch when LTS driver is updated.
             execute_process(

third_party/intel/cmake/spirv-llvm-translator-llvm-config.patch

@@ -0,0 +1,13 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index ce492a7a..b1c35b38 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -79,7 +79,7 @@ if(LLVM_SPIRV_BUILD_EXTERNAL)
+     )
+   endif(LLVM_SPIRV_INCLUDE_TESTS)
+
+-  find_package(LLVM ${BASE_LLVM_VERSION} REQUIRED
++  find_package(LLVM ${BASE_LLVM_VERSION} REQUIRED CONFIG
+     COMPONENTS
+       Analysis
+       BitReader