Remove triton submodule and add triton-hash.txt to track the triton commit ID

.github/workflows/test-plugin.yml

@@ -39,12 +39,18 @@ jobs:
       with:
         ref: ${{ inputs.triton-shared-ref }}
         path: triton_shared
-        # Also checkout triton submodule
-        submodules: recursive
+
+    - name: Checkout triton based on triton-hash.txt
+      run: |
+        git clone https://github.com/triton-lang/triton.git triton
+        cd triton
+        HASH="$(cat ../triton_shared/triton-hash.txt)"
+        echo "Checking out Triton commit: ${HASH}"
+        git checkout "${HASH}"
 
     - name: Checkout triton at ${{ inputs.triton-ref }}
       if: ${{ inputs.triton-ref }}
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         git checkout ${{ inputs.triton-ref }}
         git log -1
@@ -58,13 +64,13 @@ jobs:
         echo "PATH=${HOME}/.local/bin:${PATH}" >> "${GITHUB_ENV}"
 
     - name: Check pre-commit
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         python3 -m pip install --upgrade pre-commit
         python3 -m pre_commit run --all-files --verbose
 
     - name: Build/Install Triton
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install cmake==3.24 ninja pytest-xdist pybind11 setuptools
@@ -74,7 +80,7 @@ jobs:
         TRITON_BUILD_WITH_CLANG_LLD=true TRITON_BUILD_WITH_CCACHE=true python3 -m pip install --no-build-isolation -vvv '.[tests]'
 
     - name: Run shared middle-layer lit tests
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         python3 -m pip install lit
         LIT_TEST_DIR="build/$(ls build | grep -i cmake)/third_party/triton_shared/test"
@@ -89,7 +95,7 @@ jobs:
         python3 -m pip install pytest
 
     - name: Prepare CPU backend environment
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         CMAKE_BUILD_DIR=$(ls $(pwd)/build | grep -i cmake)
         LLVM_BINARY_DIR=$(find  ${HOME}/.triton/llvm/ -name bin | head -1)
@@ -117,7 +123,14 @@ jobs:
       with:
         ref: ${{ inputs.triton-shared-ref }}
         path: triton_shared
-        submodules: recursive
+
+    - name: Checkout triton based on triton-hash.txt
+      run: |
+        git clone https://github.com/triton-lang/triton.git triton
+        cd triton
+        HASH="$(cat ../triton_shared/triton-hash.txt)"
+        echo "Checking out Triton commit: ${HASH}"
+        git checkout "${HASH}"
 
     - name: Clear Triton Cache
       run: |
@@ -132,13 +145,13 @@ jobs:
         echo "PATH=${HOME}/.local/bin:${PATH}" >> "${GITHUB_ENV}"
 
     - name: Check pre-commit
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         python3 -m pip install --upgrade pre-commit
         python3 -m pre_commit run --all-files --verbose
 
     - name: Build/Install Triton
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install cmake==3.24 ninja pytest-xdist pybind11 setuptools
@@ -148,7 +161,7 @@ jobs:
         TRITON_BUILD_WITH_CLANG_LLD=true TRITON_BUILD_WITH_CCACHE=true python3 -m pip install --no-build-isolation -vvv '.[tests]'
 
     - name: Run shared middle-layer lit tests
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         python3 -m pip install lit
         LIT_TEST_DIR="build/$(ls build | grep -i cmake)/third_party/triton_shared/test"
@@ -163,7 +176,7 @@ jobs:
         python3 -m pip install pytest
 
     - name: Prepare CPU backend environment
-      working-directory: triton_shared/triton
+      working-directory: triton
       run: |
         CMAKE_BUILD_DIR=$(ls $(pwd)/build | grep -i cmake)
         LLVM_BINARY_DIR=$(find  ${HOME}/.triton/llvm/ -name bin | head -1)

.gitignore

@@ -2,4 +2,7 @@
 .cache
 compile_commands.json
 build/*
-.vscode/*
+.vscode/*
+/.clang-format
+test_core.py
+test_annotations.py

CMakeLists.txt

@@ -9,6 +9,14 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) # Tablegen'd files
 include_directories(${Python3_INCLUDE_DIR})
 include_directories(${pybind11_INCLUDE_DIR})
 
+function(add_symlink source_file link_name)
+    if (EXISTS ${source_file})
+        file(CREATE_LINK ${source_file} ${link_name} SYMBOLIC)
+    else()
+        message(WARNING "${source_file} not found. Skipping symlink creation for ${link_name}.")
+    endif()
+endfunction()
+
 add_subdirectory(include)
 add_subdirectory(lib)
 add_subdirectory(test)
@@ -18,3 +26,16 @@ if (TRITON_SHARED_BUILD_CPU_BACKEND)
     add_triton_plugin(TritonShared ${CMAKE_CURRENT_SOURCE_DIR}/triton_shared.cc LINK_LIBS TritonSharedAnalysis TritonToLinalg TritonTilingExtIR)
     target_link_libraries(TritonShared PRIVATE Python3::Module pybind11::headers ${Python3_LIBRARIES})
 endif()
+
+# Add symlinks to selected pytest files and the clang-format setting in triton. The tests are imported into triton-shared’s test folder to 
+# run under triton-shared's conftest configuration, and the clang-format link ensures consistent code style enforcement across both repositories.
+cmake_path(APPEND CMAKE_CURRENT_SOURCE_DIR "python" "examples" "test_core.py" OUTPUT_VARIABLE TRITON_SHARED_TEST_CORE)
+cmake_path(APPEND CMAKE_CURRENT_SOURCE_DIR "python" "examples" "test_annotations.py" OUTPUT_VARIABLE TRITON_SHARED_TEST_ANNOTATIONS)
+cmake_path(APPEND CMAKE_CURRENT_SOURCE_DIR ".clang-format" OUTPUT_VARIABLE TRITON_SHARED_CLANG_FORMAT_SETTING)
+cmake_path(APPEND CMAKE_SOURCE_DIR "python" "test" "unit" "language" "test_core.py" OUTPUT_VARIABLE TRITON_TEST_CORE)
+cmake_path(APPEND CMAKE_SOURCE_DIR "python" "test" "unit" "language" "test_annotations.py" OUTPUT_VARIABLE TRITON_TEST_ANNOTATIONS)
+cmake_path(APPEND CMAKE_SOURCE_DIR ".clang-format" OUTPUT_VARIABLE TRITON_CLANG_FORMAT_SETTING)
+
+add_symlink(${TRITON_TEST_CORE} ${TRITON_SHARED_TEST_CORE})
+add_symlink(${TRITON_TEST_ANNOTATIONS} ${TRITON_SHARED_TEST_ANNOTATIONS})
+add_symlink(${TRITON_CLANG_FORMAT_SETTING} ${TRITON_SHARED_CLANG_FORMAT_SETTING})

README.md

@@ -25,8 +25,9 @@ You need to set the `TRITON_PLUGIN_DIRS` environment variable to the location of
 ```
 export TRITON_PLUGIN_DIRS=$(pwd)/triton_shared
 
-git clone --recurse-submodules https://github.com/microsoft/triton-shared.git triton_shared
-cd triton_shared/triton
+git clone https://github.com/microsoft/triton-shared.git triton_shared
+git clone https://github.com/triton-lang/triton.git
+cd triton && git checkout $(cat ../triton_shared/triton-hash.txt)
 ```
 
 To build with Clang:

triton-hash.txt

@@ -0,0 +1 @@
+ec8cb09329cf25ac241a7dee1eea5a5d94daef8a

triton-san/README.md

@@ -38,13 +38,15 @@ git clone --recurse-submodules https://github.com/microsoft/triton-shared.git
 triton-shared/triton-san/build.sh
 ```
 **Note: because LLVM takes a long time to build, the initial run of `build.sh` may exceed 20 minutes. On later runs, the build script reuses the Python virtual environment and LLVM binary to speed up the build of `triton`, `triton-shared`, and `triton-san`.**
+
 After installation, `build.sh` generates the following folders alongside the `triton-shared` repository.
 ```
-llvm  triton-san  triton-shared  venv
+llvm  triton triton-san  triton-shared  venv
 ```
 - `llvm`: the custom LLVM source and binary,
-- `venv`: Python environment with TritonSan-enabled Triton package installed,
+- `triton` the triton source and binary,
 - `triton-san`: TritonSan driver script and associated files.
+- `venv`: Python environment with TritonSan-enabled Triton package installed,
 
 ## Usage
 To use TritonSan, run the TritonSan driver script (`triton-san/triton-san`) with the target Triton program and its corresponding inputs. We also provide two sample Triton programs containing known bugs, which `build.sh` installs into the `triton-san/examples` directory.

triton-san/build.sh

@@ -15,6 +15,7 @@ ROOT="$(realpath "${PARENT_FOLDER}/../..")"
 SCRIPT_FOLDER="$(realpath "${PARENT_FOLDER}/script")"
 VENV_PATH="${ROOT}/venv"
 LLVM_PATH="${ROOT}/llvm"
+TRITON_PATH="${ROOT}/triton"
 LLVM_BUILD_PATH="${LLVM_PATH}/llvm-build"
 LLVM_INSTALL_DIR="${LLVM_PATH}/llvm-install"
 TRITON_SAN_INSTALL_DIR="${ROOT}"
@@ -63,8 +64,14 @@ fi
 echo -e "\n\n\n"
 
 echo "================== Build trion-shared and triton =================="
-"${SCRIPT_FOLDER}/build_triton_shared_with_triton_san.sh" "${LLVM_BUILD_PATH}" "${VENV_PATH}"
+if [ -e "${TRITON_PATH}" ]; then
+  warning_msg=("The path ${TRITON_PATH} already exists and will be overwritten.")
+  print_warning "${warning_msg[@]}"
+  rm -rf "${TRITON_PATH}"
+fi
+mkdir -p "${TRITON_PATH}"
+"${SCRIPT_FOLDER}/build_triton_shared_with_triton_san.sh" "${LLVM_BUILD_PATH}" "${VENV_PATH}" "${TRITON_PATH}"
 echo -e "\n\n\n"
 
 echo "======================== Install triton-san ======================="
-"${SCRIPT_FOLDER}/install_triton_san.sh" "${TRITON_SAN_INSTALL_DIR}" "${VENV_PATH}" "${LLVM_INSTALL_DIR}"
+"${SCRIPT_FOLDER}/install_triton_san.sh" "${TRITON_SAN_INSTALL_DIR}" "${LLVM_INSTALL_DIR}" "${VENV_PATH}" "${TRITON_PATH}"

triton-san/script/build_triton_shared_with_triton_san.sh

@@ -3,15 +3,16 @@
 #!/bin/bash
 set -e
 
-if [ "$#" -lt 2 ]; then
-  echo "Usage: $0 <path to llvm-build directory> <path to Python venv>"
+if [ "$#" -lt 3 ]; then
+  echo "Usage: $0 <path to LLVM build directory> <path to Python venv> <path to Triton source directory>"
   exit 1
 fi
 
 PARENT_FOLDER="$(realpath "$(dirname "$0")")"
 TRITON_SHARED_PATH="$(realpath "${PARENT_FOLDER}/../..")"
 LLVM_BUILD_PATH="$(realpath "$1")"
 VENV_PATH="$(realpath "$2")"
+TRITON_PATH="$(realpath "$3")"
 
 # include utility functions
 source "${PARENT_FOLDER}/utility.inc"
@@ -44,8 +45,14 @@ if [ -e "${TRITON_HOME}" ]; then
 fi
 mkdir -p "${TRITON_HOME}"
 
-# build triton-shared with the custom LLVM
-cd "${TRITON_SHARED_PATH}/triton"
+# build triton-shared and triton with the custom LLVM
+TRITON_HASH_FILE="${TRITON_SHARED_PATH}/triton-hash.txt"
+if [ ! -e "${TRITON_HASH_FILE}" ]; then
+  print_error_and_exit "${TRITON_HASH_FILE} does not exist."
+fi
+
+git clone https://github.com/triton-lang/triton.git "${TRITON_PATH}"
+cd "${TRITON_PATH}" && git checkout $(cat "${TRITON_HASH_FILE}")
 export TRITON_PLUGIN_DIRS="${TRITON_SHARED_PATH}"
 export LLVM_BUILD_DIR="${LLVM_BUILD_PATH}"
 LLVM_INCLUDE_DIRS="${LLVM_BUILD_DIR}/include" LLVM_LIBRARY_DIR="${LLVM_BUILD_DIR}/lib" LLVM_SYSPATH="${LLVM_BUILD_DIR}" TRITON_BUILD_WITH_CLANG_LLD=true TRITON_BUILD_WITH_CCACHE=false python3 -m pip install --no-build-isolation -vvv '.[tests]'

triton-san/script/install_triton_san.sh

@@ -3,17 +3,17 @@
 #!/bin/bash
 set -e
 
-if [ "$#" -lt 3 ]; then
-  echo "Usage: $0 <desired path for triton-san installation> <path to venv> <path to LLVM install dir>"
+if [ "$#" -lt 4 ]; then
+  echo "Usage: $0 <desired path for triton-san installation> <path to LLVM install directory> <path to Python venv> <path to Triton source directory>"
   exit 1
 fi
 
 PARENT_FOLDER="$(realpath "$(dirname "$0")")"
 TRITON_SAN_PATH="$(realpath "${PARENT_FOLDER}/..")"
-TRITON_SHARED_PATH="$(realpath "${PARENT_FOLDER}/../..")"
 TRITON_SAN_INSTALL_DIR="$(realpath "$1")"
-VENV_PATH="$(realpath "$2")"
-LLVM_INSTALL_DIR="$(realpath "$3")"
+LLVM_INSTALL_DIR="$(realpath "$2")"
+VENV_PATH="$(realpath "$3")"
+TRITON_PATH="$(realpath "$4")"
 
 # include utility functions
 source "${PARENT_FOLDER}/utility.inc"
@@ -31,7 +31,7 @@ fi
 mkdir -p "${TRITON_SAN_INSTALL_ROOT}"
 
 # Locate required executables and objects
-locate_file "triton-shared-opt" "TRITON_SHARED_OPT_PATH" "${TRITON_SHARED_PATH}"
+locate_file "triton-shared-opt" "TRITON_SHARED_OPT_PATH" "${TRITON_PATH}"
 locate_file "libclang_rt.asan.so" "ASAN_OBJ_PATH" "${LLVM_INSTALL_DIR}"
 locate_file "libclang_rt.tsan.so" "TSAN_OBJ_PATH" "${LLVM_INSTALL_DIR}"
 locate_file "libarcher.so" "ARCHER_OBJ_PATH" "${LLVM_INSTALL_DIR}"