Update triton to 784b550 (#273)

Changes required to update to
triton-lang/triton@784b550.

- build directory moved from `triton/python` to `triton/`
- setup.py treats all contents of backend `tools/` and `language/` as
packages, so move `tools/CMakeLists.txt` and
`tools/RegisterTritonSharedDialects.h`.
- driver's interface for `load_binary` changed; must be a large enough
integer to not fail the irrelevant check
- `!(self.metadata.num_warps * warp_size > self.n_max_threads)` in
`triton/compiler/compiler.py`

Author: bmyerz0

.github/workflows/test-plugin.yml

@@ -64,7 +64,7 @@ jobs:
         python3 -m pre_commit run --all-files --verbose
 
     - name: Build/Install Triton
-      working-directory: triton_shared/triton/python
+      working-directory: triton_shared/triton
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install cmake==3.24 ninja pytest-xdist pybind11 setuptools
@@ -74,7 +74,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/python
+      working-directory: triton_shared/triton
       run: |
         python3 -m pip install lit
         LIT_TEST_DIR="build/$(ls build | grep -i cmake)/third_party/triton_shared/test"
@@ -89,7 +89,7 @@ jobs:
         python3 -m pip install pytest
 
     - name: Prepare CPU backend environment
-      working-directory: triton_shared/triton/python
+      working-directory: triton_shared/triton
       run: |
         CMAKE_BUILD_DIR=$(ls $(pwd)/build | grep -i cmake)
         LLVM_BINARY_DIR=$(find  ${HOME}/.triton/llvm/ -name bin | head -1)
@@ -138,7 +138,7 @@ jobs:
         python3 -m pre_commit run --all-files --verbose
 
     - name: Build/Install Triton
-      working-directory: triton_shared/triton/python
+      working-directory: triton_shared/triton
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install cmake==3.24 ninja pytest-xdist pybind11 setuptools
@@ -148,7 +148,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/python
+      working-directory: triton_shared/triton
       run: |
         python3 -m pip install lit
         LIT_TEST_DIR="build/$(ls build | grep -i cmake)/third_party/triton_shared/test"
@@ -163,7 +163,7 @@ jobs:
         python3 -m pip install pytest
 
     - name: Prepare CPU backend environment
-      working-directory: triton_shared/triton/python
+      working-directory: triton_shared/triton
       run: |
         CMAKE_BUILD_DIR=$(ls $(pwd)/build | grep -i cmake)
         LLVM_BINARY_DIR=$(find  ${HOME}/.triton/llvm/ -name bin | head -1)

CMakeLists.txt

@@ -11,7 +11,7 @@ include_directories(${pybind11_INCLUDE_DIR})
 add_subdirectory(include)
 add_subdirectory(lib)
 add_subdirectory(test)
-add_subdirectory(tools)
+add_subdirectory(tools/triton-shared-opt)
 
 if (TRITON_SHARED_BUILD_CPU_BACKEND)
     add_triton_plugin(TritonShared ${CMAKE_CURRENT_SOURCE_DIR}/triton_shared.cc LINK_LIBS TritonSharedAnalysis TritonToLinalg TritonTilingExtIR)

README.md

@@ -26,7 +26,7 @@ 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/python
+cd triton_shared/triton
 ```
 
 To build with Clang:
@@ -49,7 +49,7 @@ pip3 install ninja cmake wheel pytest pybind11 setuptools
 pip3 install -e . --no-build-isolation
 ```
 
-The resulting `triton-shared` binaries will be placed under `triton/python/build/{current_cmake_version}/third_party/triton_shared`
+The resulting `triton-shared` binaries will be placed under `triton/build/{current_cmake_version}/third_party/triton_shared`
 
 ### 1. Stand-Alone
 The middle layer can be used as a stand-alone component to convert Triton dialect to the middle layer dialects. This is intended for testing and validation purposes, but could potentially be used before sending the IR to another MLIR complier.
@@ -170,7 +170,7 @@ The prototype was tested on the following triton kernel examples:
 The Python tests are setup to run with Pytest and you will need to set the following environment variables to run them:
 ```
 export LLVM_BINARY_DIR=<path-to-your-llvm-binaries>
-export TRITON_SHARED_OPT_PATH=$TRITON_PLUGIN_DIRS/triton/python/build/<your-cmake-directory>/third_party/triton_shared/tools/triton-shared-opt/triton-shared-opt
+export TRITON_SHARED_OPT_PATH=$TRITON_PLUGIN_DIRS/triton/build/<your-cmake-directory>/third_party/triton_shared/tools/triton-shared-opt/triton-shared-opt
 
 pytest <path-to-triton-shared>/python/examples
 ```

backend/driver.py

@@ -349,7 +349,8 @@ def load_binary(name, kernel_obj, shared, device):
           None,       # module
           kernel_obj, # function
           None,       # n_regs
-          None        # n_spills
+          None,        # n_spills
+          sys.maxsize, # n_max_threads
         )
 
 

tools/CMakeLists.txt

@@ -5,7 +5,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "../RegisterTritonSharedDialects.h"
+#include "RegisterTritonSharedDialects.h"
 
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"