python3Packages.triton: 3.6.0 -> 3.7.0

Diff: https://github.com/triton-lang/triton/compare/v3.6.0...v3.7.0

Changelog: https://github.com/triton-lang/triton/releases/tag/v3.7.0
This commit is contained in:
Gaetan Lepage
2026-06-03 07:34:08 +00:00
parent 22807bf080
commit 30472d5544
5 changed files with 41 additions and 160 deletions
+5 -3
View File
@@ -68,7 +68,9 @@ let
in
stdenv.mkDerivation (finalAttrs: {
pname = "triton-llvm";
version = "22.0.0-unstable-2025-09-26"; # See https://github.com/llvm/llvm-project/blob/main/cmake/Modules/LLVMVersion.cmake
version = "23.0.0-unstable-2026-01-29"; # See https://github.com/llvm/llvm-project/blob/main/cmake/Modules/LLVMVersion.cmake
__structuredAttrs = true;
outputs = [
"out"
@@ -84,8 +86,8 @@ stdenv.mkDerivation (finalAttrs: {
src = fetchFromGitHub {
owner = "llvm";
repo = "llvm-project";
rev = "f6ded0be897e2878612dd903f7e8bb85448269e5";
hash = "sha256-T76zHZZ2bp3Ye9GTV+MgbKqMbtmMGElMFsWuCkiWqrM=";
rev = "ac5dc54d509169d387fcfd495d71853d81c46484";
hash = "sha256-tA1KcZqyPsgfxQs9tbNhX11oFcNGJefxWmfCpYqdI9M=";
};
nativeBuildInputs = [
@@ -1,18 +1,18 @@
diff --git a/third_party/nvidia/backend/driver.c b/third_party/nvidia/backend/driver.c
index bff09d8c1..a5c341711 100644
index f51d7b37c1..8971588282 100644
--- a/third_party/nvidia/backend/driver.c
+++ b/third_party/nvidia/backend/driver.c
@@ -1,4 +1,4 @@
-#include "cuda.h"
+#include <cuda.h>
#include <dlfcn.h>
#include <stdalign.h>
#include <stdbool.h>
#include <stdlib.h>
diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py
index 2b39fea29..3346eb954 100644
index 21e9045b25..3660e32546 100644
--- a/third_party/nvidia/backend/driver.py
+++ b/third_party/nvidia/backend/driver.py
@@ -12,7 +12,8 @@ from triton.backends.compiler import GPUTarget
@@ -11,7 +11,8 @@ from triton.backends.compiler import GPUTarget
from triton.backends.driver import GPUDriver
dirname = os.path.dirname(os.path.realpath(__file__))
@@ -20,14 +20,5 @@ index 2b39fea29..3346eb954 100644
+import shlex
+include_dirs = [*shlex.split("@cudaToolkitIncludeDirs@"), os.path.join(dirname, "include")]
libdevice_dir = os.path.join(dirname, "lib")
libraries = ['cuda']
libraries = ['libcuda.so.1']
PyCUtensorMap = None
@@ -265,7 +266,7 @@ def make_launcher(constants, signature, tensordesc_meta):
params.append("&global_scratch")
params.append("&profile_scratch")
src = f"""
-#include \"cuda.h\"
+#include <cuda.h>
#include <dlfcn.h>
#include <stdbool.h>
#include <stdlib.h>
@@ -1,131 +0,0 @@
diff --git a/test/Conversion/amd/tritongpu_to_llvm.mlir b/test/Conversion/amd/tritongpu_to_llvm.mlir
index 068e931a6..e651f7b0a 100644
--- a/test/Conversion/amd/tritongpu_to_llvm.mlir
+++ b/test/Conversion/amd/tritongpu_to_llvm.mlir
@@ -1,5 +1,6 @@
// RUN: triton-opt %s -split-input-file --allocate-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx942 --convert-builtin-func-to-llvm | FileCheck %s
// RUN: triton-opt %s -split-input-file --allocate-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx950 | FileCheck %s --check-prefix=GFX950
+// RUN: triton-opt %s -split-input-file --allocate-shared-memory --convert-triton-amdgpu-to-llvm=arch=gfx906 | FileCheck %s --check-prefix=GFX906
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
// CHECK-LABEL: atomic_add_f32_scalar
@@ -633,3 +634,27 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.thr
tt.return
}
}
+
+// -----
+
+// GFX906-LABEL: v_dot_fp16_gfx906
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [8, 8], warpsPerCTA = [2, 2], order = [1, 0]}>
+module attributes {"ttg.target" = "hip:gfx906", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 64 : i32} {
+ tt.func @v_dot_fp16_gfx906(%arg0: tensor<16x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>, %arg1: tensor<16x16xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked}>>, %arg2: tensor<16x16xf32, #blocked>) {
+ // GFX906-COUNT-8: llvm.call_intrinsic "llvm.amdgcn.fdot2"
+ %0 = tt.dot %arg0, %arg1, %arg2, inputPrecision = ieee : tensor<16x16xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked}>> * tensor<16x16xf16, #ttg.dot_op<{opIdx = 1, parent = #blocked}>> -> tensor<16x16xf32, #blocked>
+ tt.return
+ }
+}
+
+// -----
+
+// GFX906-LABEL: v_dot_i8_gfx906
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [8, 8], warpsPerCTA = [2, 2], order = [1, 0]}>
+module attributes {"ttg.target" = "hip:gfx906", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 64 : i32} {
+ tt.func @v_dot_i8_gfx906(%arg0: tensor<16x16xi8, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>, %arg1: tensor<16x16xi8, #ttg.dot_op<{opIdx = 1, parent = #blocked}>>, %arg2: tensor<16x16xi32, #blocked>) {
+ // GFX906-COUNT-4: llvm.call_intrinsic "llvm.amdgcn.sdot4"
+ %0 = tt.dot %arg0, %arg1, %arg2, inputPrecision = ieee : tensor<16x16xi8, #ttg.dot_op<{opIdx = 0, parent = #blocked}>> * tensor<16x16xi8, #ttg.dot_op<{opIdx = 1, parent = #blocked}>> -> tensor<16x16xi32, #blocked>
+ tt.return
+ }
+}
diff --git a/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h b/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h
index a689fe438..5d84110dc 100644
--- a/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h
+++ b/third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h
@@ -8,6 +8,7 @@ namespace mlir::triton::AMD {
// A list of ISA families we care about.
enum class ISAFamily {
Unknown,
+ GCN5_1,
CDNA1,
CDNA2,
CDNA3,
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
index af7b4b1f7..01cacc9de 100644
--- a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
+++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
@@ -71,6 +71,7 @@ llvm::AMDGPU::GPUKind TargetInfo::getGPUKind() const {
int TargetInfo::getWarpSize() const {
switch (getISAFamily()) {
+ case ISAFamily::GCN5_1:
case ISAFamily::CDNA1:
case ISAFamily::CDNA2:
case ISAFamily::CDNA3:
@@ -335,10 +336,10 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
return true;
if (numLaneToReduce != getWarpSize())
return false;
- if (isCDNA(getISAFamily()) && getISAFamily() == ISAFamily::CDNA1)
- return false;
- if (isRDNA(getISAFamily()) &&
- llvm::is_contained({ISAFamily::RDNA1, ISAFamily::RDNA2}, getISAFamily()))
+ // DPP warp reduce requires gfx90a+ (CDNA2+) or gfx11+ (RDNA3+).
+ // Pre-CDNA2 GFX9 (gfx906/gfx908) and GFX10 (RDNA1/2) are excluded.
+ auto v = getIsaVersion();
+ if (!((v.Major == 9 && (v.Minor > 0 || v.Stepping >= 0xa)) || v.Major >= 11))
return false;
Operation *reduxOp = op.getSingleCombiner();
@@ -438,7 +439,7 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
buf = createDppReduxOpWithBoundCtrl(valType, buf, 1 + dppCtrlRowShr,
allRows, allBanks);
- if (isCDNA(getISAFamily())) {
+ if (isCDNA(getISAFamily()) || getISAFamily() == ISAFamily::GCN5_1) {
// row_bcast:15 row_mask:0xa
buf = createDppReduxOpWithBoundCtrl(
valType, buf, static_cast<uint32_t>(DppCtrl::BCAST15), 0xa, allBanks);
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp.rej b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp.rej
new file mode 100644
index 000000000..b75ab8f11
--- /dev/null
+++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp.rej
@@ -0,0 +1,16 @@
+diff a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp (rejected hunks)
+@@ -395,10 +396,10 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
+ return true;
+ if (reduceLaneIdMask != (getWarpSize() - 1))
+ return false;
+- if (isCDNA(getISAFamily()) && getISAFamily() == ISAFamily::CDNA1)
+- return false;
+- if (isRDNA(getISAFamily()) &&
+- llvm::is_contained({ISAFamily::RDNA1, ISAFamily::RDNA2}, getISAFamily()))
++ // DPP warp reduce requires gfx90a+ (CDNA2+) or gfx11+ (RDNA3+).
++ // Pre-CDNA2 GFX9 (gfx906/gfx908) and GFX10 (RDNA1/2) are excluded.
++ auto v = getIsaVersion();
++ if (!((v.Major == 9 && (v.Minor > 0 || v.Stepping >= 0xa)) || v.Major >= 11))
+ return false;
+
+ Operation *reduxOp = op.getSingleCombiner();
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp
index b79841ab5..afb80d9fb 100644
--- a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp
+++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetUtils.cpp
@@ -12,6 +12,9 @@ ISAFamily deduceISAFamily(llvm::StringRef arch) {
if (kind == llvm::AMDGPU::GK_GFX1250)
return ISAFamily::GFX1250;
+ if (kind == llvm::AMDGPU::GK_GFX906)
+ return ISAFamily::GCN5_1;
+
// CDNA ISA cases
switch (kind) {
case llvm::AMDGPU::GK_GFX950:
@@ -41,6 +44,7 @@ ISAFamily deduceISAFamily(llvm::StringRef arch) {
bool supportsVDot(llvm::StringRef arch) {
switch (deduceISAFamily(arch)) {
+ case AMD::ISAFamily::GCN5_1:
case AMD::ISAFamily::CDNA1:
case AMD::ISAFamily::CDNA2:
case AMD::ISAFamily::CDNA3:
@@ -44,15 +44,16 @@
buildPythonPackage (finalAttrs: {
pname = "triton";
version = "3.6.0";
version = "3.7.0";
pyproject = true;
__structuredAttrs = true;
# Remember to bump triton-llvm as well!
src = fetchFromGitHub {
owner = "triton-lang";
repo = "triton";
tag = "v${finalAttrs.version}";
hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
hash = "sha256-FxbBY1lPq7765MqAPR7UljzPsmjOhKKbYExlKgeudew=";
};
patches = [
@@ -63,8 +64,9 @@ buildPythonPackage (finalAttrs: {
libcudaStubsDir =
if cudaSupport then "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs" else null;
})
# Backport of https://github.com/triton-lang/triton/pull/9628 (does not apply cleanly)
./0005-add-gcn5-gfx906-target.patch
# Use our `cmakeFlags` instead and avoid downloading dependencies.
./inject-nix-cmakeFlags.patch
]
++ lib.optionals cudaSupport [
(replaceVars ./0003-nvidia-cudart-a-systempath.patch {
@@ -90,13 +92,6 @@ buildPythonPackage (finalAttrs: {
--replace-fail 'yield ("triton.profiler", "third_party/proton/proton")' 'pass' \
--replace-fail "curr_version.group(1) != version" "False"
''
# Use our `cmakeFlags` instead and avoid downloading dependencies
+ ''
substituteInPlace setup.py \
--replace-fail \
"cmake_args.extend(thirdparty_cmake_args)" \
"cmake_args.extend(thirdparty_cmake_args + os.environ.get('cmakeFlags', \"\").split())"
''
# Don't fetch googletest
+ ''
substituteInPlace cmake/AddTritonUnitTest.cmake \
@@ -135,7 +130,7 @@ buildPythonPackage (finalAttrs: {
# `find_package` is called with `NO_DEFAULT_PATH`
# https://cmake.org/cmake/help/latest/command/find_package.html
# https://github.com/triton-lang/triton/blob/c3c476f357f1e9768ea4e45aa5c17528449ab9ef/third_party/amd/CMakeLists.txt#L6
(lib.cmakeFeature "LLD_DIR" "${lib.getLib llvm}")
(lib.cmakeFeature "LLD_DIR" "${lib.getLib llvm}/lib/cmake/lld")
];
buildInputs = [
@@ -159,6 +154,14 @@ buildPythonPackage (finalAttrs: {
export MAX_JOBS="$NIX_BUILD_CORES"
'';
# `examples/plugins` (an MLIR example dialect plugin and a unit-test helper lib) is built
# unconditionally with the Python module and shipped into `triton/plugins/`.
# It is unused at runtime and keeps a forbidden RPATH reference to the build directory, which
# fails the fixup phase.
postInstall = ''
rm -rf "$out/${python.sitePackages}/triton/plugins"
'';
env = {
TRITON_BUILD_PROTON = "OFF";
TRITON_OFFLINE_BUILD = true;
@@ -0,0 +1,16 @@
diff --git a/setup.py b/setup.py
index d8e0860291..bc09dada5a 100644
--- a/setup.py
+++ b/setup.py
@@ -476,6 +476,11 @@ class CMakeBuild(build_ext):
cmake_args.append("-DLLVM_EXTERNAL_LIT=" + lit_dir)
cmake_args.extend(thirdparty_cmake_args)
+ import json
+
+ nix_attrs_json = json.load(open(os.environ["NIX_ATTRS_JSON_FILE"]))
+ cmake_args.extend(thirdparty_cmake_args + nix_attrs_json.get("cmakeFlags", []))
+
# configuration
cfg = get_build_type()
build_args = ["--config", cfg]