Skip to content

feat: prefer aotrion from release #541

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
drbh merged 3 commits into from
May 19, 2026
+36 −166
Merged

feat: prefer aotrion from release #541

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
318ff1f
feat: prefer aotrion from release
drbh May 11, 2026
6e03fa7
fix: resolve libamdhip
drbh May 12, 2026
42d32c9
fix: remove installCheckPhase
drbh May 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 8 additions & 60 deletions nix-builder/pkgs/aotriton/default.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,21 +1,11 @@
{
callPackage,
fetchFromGitHub,
fetchpatch,
fetchurl,
stdenvNoCC,
}:

let
generic = callPackage ./generic.nix { };
postFetch = ''
cd $out
git reset --hard HEAD
for submodule in $(git config --file .gitmodules --get-regexp path | awk '{print $2}' | grep '^third_party/' | grep -v '^third_party/triton$'); do
git submodule update --init --recursive "$submodule"
done
find "$out" -name .git -print0 | xargs -0 rm -rf
'';
mkImages =
version: srcs:
stdenvNoCC.mkDerivation {
Expand All @@ -35,31 +25,12 @@ in
aotriton_0_11_1 = generic rec {
version = "0.11.1b";

src = fetchFromGitHub {
owner = "ROCm";
repo = "aotriton";
tag = version;
hash = "sha256-F7JjyS+6gMdCpOFLldTsNJdVzzVwd6lwW7+V8ZOZfig=";
leaveDotGit = true;
inherit postFetch;
hashes = {
"7.0" = "sha256-3rgEbp75dsJzn9BWO1AjnhLcAC19T5fBxKGHSstlq8Q=";
"7.1" = "sha256-wWE+2enuzHNZ8EoWJLtSjlT15jaeaC3URuqpNtlFI1g=";
"7.2" = "sha256-VsoxJUwWVfpNUWji2zFZeBwkQtX2sBiCFV6ThZuFzxY=";
};

patches = [
# Fails with: ld.lld: error: unable to insert .comment after .comment
./v0.11.1b-no-ld-script.diff
];

gpuTargets = [
# aotriton GPU support list:
# https://github.com/ROCm/aotriton/blob/main/v2python/gpu_targets.py
"gfx90a"
"gfx942"
"gfx950"
"gfx1100"
"gfx1151"
"gfx1201"
];

images = mkImages version [
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11.1b/aotriton-0.11.1b-images-amd-gfx90a.tar.gz";
Expand All @@ -82,38 +53,17 @@ in
hash = "sha256-Ck/zJL/9rAwv3oeop/cFY9PISoCtTo8xNF8rQKE4TpU=";
})
];

extraPythonDepends = ps: [ ps.pandas ];
};

aotriton_0_11_2 = generic rec {
version = "0.11.2b";

src = fetchFromGitHub {
owner = "ROCm";
repo = "aotriton";
tag = version;
hash = "sha256-VIwwQR1fl40NLNOwO8KhQK/xOK6wb2l8qBugJ1cRjm4=";
leaveDotGit = true;
inherit postFetch;
hashes = {
"7.0" = "sha256-VQGgo7MAiQABtmJfKjU5p7rWDzhvCgYevn1O1coPr7k=";
"7.1" = "sha256-/uNr6z6khM4YFVu6/gJsV3/WcF5EaeWUBbJgvXS4zBA=";
"7.2" = "sha256-zYq/J7u2POxFyUE16bKHRZZgdCY6awVV5YeK4ctqI0k=";
};

patches = [
# Fails with: ld.lld: error: unable to insert .comment after .comment
./v0.11.1b-no-ld-script.diff
];

gpuTargets = [
# aotriton GPU support list:
# https://github.com/ROCm/aotriton/blob/main/v2python/gpu_targets.py
"gfx90a"
"gfx942"
"gfx950"
"gfx1100"
"gfx1151"
"gfx1201"
];

images = mkImages version [
(fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/0.11.2b/aotriton-0.11.2b-images-amd-gfx90a.tar.gz";
Expand All @@ -136,8 +86,6 @@ in
hash = "sha256-Ck/zJL/9rAwv3oeop/cFY9PISoCtTo8xNF8rQKE4TpU=";
})
];

extraPythonDepends = ps: [ ps.pandas ];
};

}
134 changes: 28 additions & 106 deletions nix-builder/pkgs/aotriton/generic.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,137 +1,59 @@
# Vendored from nixpkgs
{
autoPatchelfHook,
clr,
fetchurl,
lib,
rocm-core,
stdenv,
cmake,
jq,
python3,
ninja,
pkg-config,
rocmPackages,
writableTmpDirAsHomeHook,
writeShellScriptBin,
xz,
}:

{
version,
gpuTargets,
patches ? [ ],
src,
images,
extraPythonDepends ? ps: [ ],
hashes,
}:

let
gpuTargets' = lib.concatStringsSep ";" gpuTargets;
compiler = "amdclang++";
rocmVersion = lib.versions.majorMinor rocm-core.version;
hash =
hashes.${rocmVersion}
or (throw "aotriton ${version} binary package is not specified for ROCm ${rocmVersion}");
in
stdenv.mkDerivation (finalAttrs: {
stdenv.mkDerivation {
pname = "aotriton";
inherit version;

inherit version src patches;

env = {
#CXX = compiler;
ROCM_PATH = "${rocmPackages.clr}";
CFLAGS = "-w -g1 -gz -Wno-c++11-narrowing";
CXXFLAGS = finalAttrs.env.CFLAGS;

# aotriton passes a lot of files to the linker.
NIX_LD_USE_RESPONSE_FILE = 1;
src = fetchurl {
url = "https://github.com/ROCm/aotriton/releases/download/${version}/aotriton-${version}-manylinux_2_28_x86_64-rocm${rocmVersion}-shared.tar.gz";
inherit hash;
};

requiredSystemFeatures = [ "big-parallel" ];

nativeBuildInputs = [
cmake
jq
rocmPackages.rocm-cmake
pkg-config
python3
ninja
rocmPackages.clr
writableTmpDirAsHomeHook # venv wants to cache in ~
(writeShellScriptBin "amdclang++" ''
exec ${rocmPackages.llvm.clang}/bin/clang++ "$@"
'')
];

nativeBuildInputs = [ autoPatchelfHook ];
buildInputs = [
rocmPackages.clr
clr
stdenv.cc.cc.lib
xz
]
++ (with python3.pkgs; [
wheel
packaging
pyyaml
numpy
filelock
iniconfig
pluggy
pybind11
pandas
triton
]);

preConfigure = lib.optionalString (lib.versionAtLeast version "0.11.1") ''
# Since we use pre-built images, we can grab the image SHA from there.
# As of 0.11.1b this doesn't seem to be used for image loading yet, but
# just in case this happens in the future, we set this to the actual
# value and not a stub.
export AOTRITON_CI_SUPPLIED_SHA1=$(jq -r '.["AOTRITON_GIT_SHA1"]' ${images}/lib/aotriton.images/amd-gfx90a/__signature__)

# Need to set absolute paths to VENV and its PYTHON or
# build fails with "AOTRITON_INHERIT_SYSTEM_SITE_TRITON is enabled
# but triton is not available … no such file or directory"
# Set via a preConfigure hook so a valid absolute path can be
# picked if nix-shell is used against this package
cmakeFlagsArray+=(
"-DVENV_DIR=$(pwd)/aotriton-venv/"
"-DVENV_BIN_PYTHON=$(pwd)/aotriton-venv/bin/python"
)
'';
];

# From README:
# Note: do not run ninja separately, due to the limit of the current build system,
# ninja install will run the whole build process unconditionally.
dontConfigure = true;
dontBuild = true;

dontStrip = true;
installPhase = ''
runHook preInstall
ninja -v install
ln -sf ${images}/lib/aotriton.images $out/lib/aotriton.images
runHook postInstall
'';

doCheck = false;
doInstallCheck = false;
mkdir -p "$out"
cp -r include lib "$out/"
ln -s ${images}/lib/aotriton.images "$out/lib/aotriton.images"

cmakeFlags = [
# Disable building kernels if no supported targets are enabled
(lib.cmakeBool "AOTRITON_NOIMAGE_MODE" true)
# Use preinstalled triton from our python's site-packages
(lib.cmakeBool "AOTRITON_INHERIT_SYSTEM_SITE_TRITON" true)
# Avoid kernels being skipped if build host is overloaded
(lib.cmakeFeature "AOTRITON_GPU_BUILD_TIMEOUT" "0")
(lib.cmakeFeature "CMAKE_CXX_COMPILER" compiler)
# Manually define CMAKE_INSTALL_<DIR>
# See: https://github.com/NixOS/nixpkgs/pull/197838
(lib.cmakeFeature "CMAKE_INSTALL_BINDIR" "bin")
(lib.cmakeFeature "CMAKE_INSTALL_LIBDIR" "lib")
(lib.cmakeFeature "CMAKE_INSTALL_INCLUDEDIR" "include")
(lib.cmakeFeature "AOTRITON_TARGET_ARCH" gpuTargets')
(lib.cmakeBool "AOTRITON_USE_TORCH" false)
];
runHook postInstall
'';

meta = with lib; {
description = "Ahead of Time (AOT) Triton Math Library";
homepage = "https://github.com/ROCm/aotriton";
license = with licenses; [ mit ];
platforms = platforms.linux;
sourceProvenance = with sourceTypes; [
fromSource
binaryNativeCode # aotriton.images
];
platforms = [ "x86_64-linux" ];
sourceProvenance = with sourceTypes; [ binaryNativeCode ];
};
})
}
Loading