Merge branch 'proton-dev' into proton_lowering

triton-lang · Feb 7, 2025 · 1706ecf · 1706ecf
2 parents 2e9a9f4 + 58225df
commit 1706ecf
Show file tree

Hide file tree

Showing 215 changed files with 7,153 additions and 5,507 deletions.
diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml
@@ -26,6 +26,7 @@ env:
   TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE"
   TRITON_DISABLE_LINE_INFO: 1
   PROTON_SKIP_PC_SAMPLING_TEST: 1
+  PYTHON: "python3"
   CCACHE_COMPRESS: "true"
 jobs:
   Runner-Preparation:
@@ -233,68 +234,28 @@ jobs:
       - name: Update PATH
         run: |
           echo "$HOME/.local/bin" >> $GITHUB_PATH
-      - name: Install pip dependencies
-        run: |
-          python3 -m pip install --upgrade pip
-          python3 -m pip install cython setuptools wheel cmake==3.24 ninja lit
       - name: Install Triton
         env:
           CUDA_HOME: "/usr/local/cuda"
         run: |
           echo "PATH is '$PATH'"
-          cd python
           ccache --zero-stats
-          python3 -m pip install -v '.[tests]'
+          make dev-install
       - name: CCache Stats
         run: ccache --print-stats
       - name: Run lit tests
-        run: |
-          cd python
-          LIT_TEST_DIR="build/$(ls build | grep -i cmake)/test"
-          if [ ! -d "${LIT_TEST_DIR}" ]; then
-            echo "Could not find '${LIT_TEST_DIR}'" ; exit -1
-          fi
-          lit -v "${LIT_TEST_DIR}"
+        run: make test-lit
       - name: Run python tests on CUDA
-        run: |
-          INSTRUMENTATION_LIB_DIR="${GITHUB_WORKSPACE}/python/build/$(ls python/build | grep -i lib)/triton/instrumentation"
-          if [ ! -d "${INSTRUMENTATION_LIB_DIR}" ]; then
-            echo "Could not find '${INSTRUMENTATION_LIB_DIR}'" ; exit -1
-          fi
-          cd python/test/unit
-          python3 -m pytest -s -n 8 --ignore=cuda/test_flashattention.py --ignore=language/test_line_info.py --ignore=language/test_subprocess.py --ignore=test_debug.py
-          python3 -m pytest -s -n 8 language/test_subprocess.py
-          python3 -m pytest -s -n 8 test_debug.py --forked
-          # Run test_line_info.py separately with TRITON_DISABLE_LINE_INFO=0
-          TRITON_DISABLE_LINE_INFO=0 python3 -m pytest -s language/test_line_info.py
-          # Run cuda/test_flashattention.py separately to avoid out of gpu memory
-          python3 -m pytest -s cuda/test_flashattention.py
-          TRITON_ALWAYS_COMPILE=1 TRITON_DISABLE_LINE_INFO=0 LLVM_PASS_PLUGIN_PATH=${INSTRUMENTATION_LIB_DIR}/libGPUInstrumentationTestLib.so \
-          python3 -m pytest --capture=tee-sys -rfs -vvv instrumentation/test_gpuhello.py
+        run: make test-unit
       - name: Run interpreter tests
         if: ${{ matrix.runner[0] == 'h100-runner-set' }}
-        env:
-          TRITON_INTERPRET: "1"
-        run: |
-          cd python/test/unit
-          python3 -m pytest -s -n 16 -m interpreter language/test_core.py language/test_standard.py \
-           language/test_random.py language/test_block_pointer.py language/test_subprocess.py language/test_line_info.py \
-           runtime/test_autotuner.py::test_kwargs[False]\
-           ../../tutorials/06-fused-attention.py::test_op --device cpu
+        run: make test-interpret
       - name: Run regression tests
-        run: |
-          cd python/test/regression
-          python3 -m pytest -s -n 8 .
+        run: make test-regression
       - name: Run C++ unittests
-        run: |
-          cd python
-          cd "build/$(ls build | grep -i cmake)"
-          ctest -j32
+        run: make test-cpp
       - name: Run Proton tests
-        run: |
-          cd third_party/proton/test
-          python3 -m pytest -s .
-          cd ..
+        run: make test-proton
       - name: Inspect cache directories
         run: |
           mkdir -p ~/.triton
@@ -405,19 +366,12 @@ jobs:
         run: |
           echo "PATH is '$PATH'"
           pip uninstall -y triton pytorch-triton-rocm
-          cd python
           ccache --zero-stats
-          pip install -v -e '.[tests]'
+          make dev-install
       - name: CCache Stats
         run: ccache --print-stats
       - name: Run lit tests
-        run: |
-          cd python
-          LIT_TEST_DIR="build/$(ls build | grep -i cmake)/test"
-          if [ ! -d "${LIT_TEST_DIR}" ]; then
-            echo "Could not find '${LIT_TEST_DIR}'" ; exit -1
-          fi
-          lit -v "${LIT_TEST_DIR}"
+        run: make test-lit
       - name: Run python tests on HIP
         run: |
           INSTRUMENTATION_LIB_DIR="${GITHUB_WORKSPACE}/python/triton/instrumentation"
@@ -455,15 +409,9 @@ jobs:
           cd python/test/regression
           python3 -m pytest -s -n 8 ./test_cast_matmul.py
       - name: Run Proton tests
-        run: |
-          cd third_party/proton/test
-          python3 -m pytest -s .
-          cd ..
+        run: make test-proton
       - name: Run C++ unittests
-        run: |
-          cd python
-          cd "build/$(ls build | grep -i cmake)"
-          ctest -j32
+        run: make test-cpp
       - name: Inspect cache directories
         run: |
           mkdir -p ~/.triton
@@ -575,12 +523,11 @@ jobs:
         run: |
           echo "$HOME/.local/bin" >> $GITHUB_PATH
           echo "/opt/homebrew/opt/llvm/bin" >> $GITHUB_PATH
-      - name: Install pip dependencies
+      - name: Create venv
         run: |
           python3 -m venv ~/.venv
           source ~/.venv/bin/activate
           python3 -m pip install --upgrade pip
-          python3 -m pip install cython setuptools wheel cmake==3.24 ninja lit pybind11
       - name: Install Triton
         env:
           TRITON_BUILD_WITH_O1: "true"
@@ -590,9 +537,8 @@ jobs:
         run: |
           source ~/.venv/bin/activate
           echo "PATH is '$PATH'"
-          cd python
           ccache --zero-stats
-          python3 -m pip install -v --no-build-isolation .
+          make dev-install
       - name: CCache Stats
         run: ccache --print-stats
       - name: Inspect cache directories

diff --git a/.github/workflows/integration-tests.yml.in b/.github/workflows/integration-tests.yml.in
@@ -28,6 +28,7 @@ env:
   TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE"
   TRITON_DISABLE_LINE_INFO: 1
   PROTON_SKIP_PC_SAMPLING_TEST: 1
+  PYTHON: "python3"
   CCACHE_COMPRESS: "true"
 
 jobs:
@@ -265,79 +266,39 @@ jobs:
         run: |
           echo "$HOME/.local/bin" >> $GITHUB_PATH
 
-      - name: Install pip dependencies
-        run: |
-          python3 -m pip install --upgrade pip
-          python3 -m pip install cython setuptools wheel cmake==3.24 ninja lit
-
       - name: Install Triton
         env:
           CUDA_HOME: "/usr/local/cuda"
         run: |
           echo "PATH is '$PATH'"
-          cd python
           ccache --zero-stats
-          python3 -m pip install -v '.[tests]'
+          make dev-install
 
       - &print-ccache-stats
         name: CCache Stats
         run: ccache --print-stats
 
       - &run-lit-tests-step
         name: Run lit tests
-        run: |
-          cd python
-          LIT_TEST_DIR="build/$(ls build | grep -i cmake)/test"
-          if [ ! -d "${LIT_TEST_DIR}" ]; then
-            echo "Could not find '${LIT_TEST_DIR}'" ; exit -1
-          fi
-          lit -v "${LIT_TEST_DIR}"
+        run: make test-lit
 
       - name: Run python tests on CUDA
-        run: |
-          INSTRUMENTATION_LIB_DIR="${GITHUB_WORKSPACE}/python/build/$(ls python/build | grep -i lib)/triton/instrumentation"
-          if [ ! -d "${INSTRUMENTATION_LIB_DIR}" ]; then
-            echo "Could not find '${INSTRUMENTATION_LIB_DIR}'" ; exit -1
-          fi
-          cd python/test/unit
-          python3 -m pytest -s -n 8 --ignore=cuda/test_flashattention.py --ignore=language/test_line_info.py --ignore=language/test_subprocess.py --ignore=test_debug.py
-          python3 -m pytest -s -n 8 language/test_subprocess.py
-          python3 -m pytest -s -n 8 test_debug.py --forked
-          # Run test_line_info.py separately with TRITON_DISABLE_LINE_INFO=0
-          TRITON_DISABLE_LINE_INFO=0 python3 -m pytest -s language/test_line_info.py
-          # Run cuda/test_flashattention.py separately to avoid out of gpu memory
-          python3 -m pytest -s cuda/test_flashattention.py
-          TRITON_ALWAYS_COMPILE=1 TRITON_DISABLE_LINE_INFO=0 LLVM_PASS_PLUGIN_PATH=${INSTRUMENTATION_LIB_DIR}/libGPUInstrumentationTestLib.so \
-          python3 -m pytest --capture=tee-sys -rfs -vvv instrumentation/test_gpuhello.py
+        run: make test-unit
+
       - name: Run interpreter tests
         if: ${{ matrix.runner[0] == 'h100-runner-set' }}
-        env:
-          TRITON_INTERPRET: "1"
-        run: |
-          cd python/test/unit
-          python3 -m pytest -s -n 16 -m interpreter language/test_core.py language/test_standard.py \
-           language/test_random.py language/test_block_pointer.py language/test_subprocess.py language/test_line_info.py \
-           runtime/test_autotuner.py::test_kwargs[False]\
-           ../../tutorials/06-fused-attention.py::test_op --device cpu
+        run: make test-interpret
 
       - name: Run regression tests
-        run: |
-          cd python/test/regression
-          python3 -m pytest -s -n 8 .
+        run: make test-regression
 
       - &run-cpp-unittests-step
         name: Run C++ unittests
-        run: |
-          cd python
-          cd "build/$(ls build | grep -i cmake)"
-          ctest -j32
+        run: make test-cpp
 
       - &run-proton-tests-step
         name: Run Proton tests
-        run: |
-          cd third_party/proton/test
-          python3 -m pytest -s .
-          cd ..
+        run: make test-proton
 
       - *inspect-cache-directories-step
 
@@ -397,9 +358,8 @@ jobs:
         run: |
           echo "PATH is '$PATH'"
           pip uninstall -y triton pytorch-triton-rocm
-          cd python
           ccache --zero-stats
-          pip install -v -e '.[tests]'
+          make dev-install
 
       - *print-ccache-stats
       - *run-lit-tests-step
@@ -484,12 +444,11 @@ jobs:
         run: |
           echo "$HOME/.local/bin" >> $GITHUB_PATH
           echo "/opt/homebrew/opt/llvm/bin" >> $GITHUB_PATH
-      - name: Install pip dependencies
+      - name: Create venv
         run: |
           python3 -m venv ~/.venv
           source ~/.venv/bin/activate
           python3 -m pip install --upgrade pip
-          python3 -m pip install cython setuptools wheel cmake==3.24 ninja lit pybind11
       - name: Install Triton
         env:
           TRITON_BUILD_WITH_O1: "true"
@@ -499,9 +458,8 @@ jobs:
         run: |
           source ~/.venv/bin/activate
           echo "PATH is '$PATH'"
-          cd python
           ccache --zero-stats
-          python3 -m pip install -v --no-build-isolation .
+          make dev-install
 
       - *print-ccache-stats
       - *inspect-cache-directories-step

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -39,7 +39,8 @@ The core maintainers as a group have the power to veto any decision made at a Mo
 
 The core maintainers should publicly articulate their decision-making, and share the reasoning behind their decisions, vetoes, and dispute resolution.
 
-__List of core maintainers (as of 08/30/2024, alphabetical order):__
+__List of core maintainers (as of 01/30/2025, alphabetical order):__
+* Jeff Niu
 * Keren Zhou
 * Mario Lezcano-Casado
 * Pawel Szczerbuk

diff --git a/Makefile b/Makefile
@@ -2,9 +2,12 @@
 # Make sure to first initialize the build system with:
 #     make dev-install
 
-PYTHON := python
+PYTHON ?= python
 BUILD_DIR := $(shell cd python; $(PYTHON) -c 'from build_helpers import get_cmake_dir; print(get_cmake_dir())')
 TRITON_OPT := $(BUILD_DIR)/bin/triton-opt
+PYTEST := $(PYTHON) -m pytest
+
+# Incremental builds
 
 .PHONY: all
 all:
@@ -14,6 +17,8 @@ all:
 triton-opt:
 	ninja -C $(BUILD_DIR) triton-opt
 
+# Testing
+
 .PHONY: test-lit
 test-lit:
 	ninja -C $(BUILD_DIR) check-triton-lit-tests
@@ -23,20 +28,65 @@ test-cpp:
 	ninja -C $(BUILD_DIR) check-triton-unit-tests
 
 .PHONY: test-python
-test-python: all
-	$(PYTHON) -m pytest python/test/unit
+test-unit: all
+	cd python/test/unit && $(PYTEST) -s -n 8 --ignore=cuda/test_flashattention.py \
+		--ignore=language/test_line_info.py --ignore=language/test_subprocess.py --ignore=test_debug.py
+	$(PYTEST) -s -n 8 python/test/unit/language/test_subprocess.py
+	$(PYTEST) -s -n 8 python/test/unit/test_debug.py --forked
+	TRITON_DISABLE_LINE_INFO=0 $(PYTEST) -s python/test/unit/language/test_line_info.py
+	# Run cuda/test_flashattention.py separately to avoid out of gpu memory
+	$(PYTEST) -s python/test/unit/cuda/test_flashattention.py
+	TRITON_ALWAYS_COMPILE=1 TRITON_DISABLE_LINE_INFO=0 LLVM_PASS_PLUGIN_PATH=python/triton/instrumentation/libGPUInstrumentationTestLib.so \
+		$(PYTEST) --capture=tee-sys -rfs -vvv python/test/unit/instrumentation/test_gpuhello.py
+
+.PHONY: test-regression
+test-regression: all
+	$(PYTEST) -s -n 8 python/test/regression
+
+.PHONY: test-interpret
+test-interpret: all
+	cd python/test/unit && TRITON_INTERPRET=1 $(PYTEST) -s -n 16 -m interpreter cuda language/test_core.py language/test_standard.py \
+		language/test_random.py language/test_block_pointer.py language/test_subprocess.py language/test_line_info.py \
+		language/test_tuple.py runtime/test_autotuner.py::test_kwargs[False] \
+		../../tutorials/06-fused-attention.py::test_op --device=cpu
+
+.PHONY: test-proton
+test-proton: all
+	$(PYTEST) -s -n 8 third_party/proton/test
+
+.PHONY: test-python
+test-python: test-unit test-regression test-interpret test-proton
+
+.PHONY: test-nogpu
+test-nogpu: test-lit test-cpp
 
 .PHONY: test
 test: test-lit test-cpp test-python
 
-.PHONY: dev-install
-dev-install:
-	# build-time dependencies
-	$(PYTHON) -m pip install ninja cmake wheel pybind11
-	# test dependencies
-	$(PYTHON) -m pip install scipy numpy torch pytest lit pandas matplotlib
+# pip install-ing
+
+.PHONY: dev-install-requires
+dev-install-requires:
+	$(PYTHON) -m pip install -r python/requirements.txt
+	$(PYTHON) -m pip install -r python/test-requirements.txt
+
+
+.PHONY: dev-install-torch
+dev-install-torch:
+	# install torch but ensure pytorch-triton isn't installed
+	$(PYTHON) -m pip install torch
+	$(PYTHON) -m pip uninstall triton pytorch-triton -y
+
+.PHONY: dev-install-triton
+dev-install-triton:
 	$(PYTHON) -m pip install -e python --no-build-isolation -v
 
+.PHONY: dev-install
+.NOPARALLEL: dev-install
+dev-install: dev-install-requires dev-install-triton
+
+# Updating lit tests
+
 .PHONY: golden-samples
 golden-samples: triton-opt
 	$(TRITON_OPT) test/TritonGPU/samples/simulated-grouped-gemm.mlir.in -tritongpu-loop-scheduling -tritongpu-pipeline -canonicalize | \