diff --git a/pkgs/by-name/tr/triton-llvm/package.nix b/pkgs/by-name/tr/triton-llvm/package.nix index 9f75c47099ec..1ede101acda8 100644 --- a/pkgs/by-name/tr/triton-llvm/package.nix +++ b/pkgs/by-name/tr/triton-llvm/package.nix @@ -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 = [ diff --git a/pkgs/development/python-modules/triton/0003-nvidia-cudart-a-systempath.patch b/pkgs/development/python-modules/triton/0003-nvidia-cudart-a-systempath.patch index da5ff0c9de7c..7f479cf967c6 100644 --- a/pkgs/development/python-modules/triton/0003-nvidia-cudart-a-systempath.patch +++ b/pkgs/development/python-modules/triton/0003-nvidia-cudart-a-systempath.patch @@ -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 #include + #include #include - #include 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 - #include - #include - #include diff --git a/pkgs/development/python-modules/triton/0005-add-gcn5-gfx906-target.patch b/pkgs/development/python-modules/triton/0005-add-gcn5-gfx906-target.patch deleted file mode 100644 index 6fcc6a22bc42..000000000000 --- a/pkgs/development/python-modules/triton/0005-add-gcn5-gfx906-target.patch +++ /dev/null @@ -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(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: diff --git a/pkgs/development/python-modules/triton/default.nix b/pkgs/development/python-modules/triton/default.nix index cd81a8c6de0f..747f2baae44e 100644 --- a/pkgs/development/python-modules/triton/default.nix +++ b/pkgs/development/python-modules/triton/default.nix @@ -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; diff --git a/pkgs/development/python-modules/triton/inject-nix-cmakeFlags.patch b/pkgs/development/python-modules/triton/inject-nix-cmakeFlags.patch new file mode 100644 index 000000000000..7d0757107b6d --- /dev/null +++ b/pkgs/development/python-modules/triton/inject-nix-cmakeFlags.patch @@ -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]