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:
@@ -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]
|
||||
Reference in New Issue
Block a user