|
1 | | -# Vendored from nixpkgs |
2 | 1 | { |
| 2 | + autoPatchelfHook, |
| 3 | + clr, |
| 4 | + fetchurl, |
3 | 5 | lib, |
| 6 | + rocm-core, |
4 | 7 | stdenv, |
5 | | - cmake, |
6 | | - jq, |
7 | | - python3, |
8 | | - ninja, |
9 | | - pkg-config, |
10 | | - rocmPackages, |
11 | | - writableTmpDirAsHomeHook, |
12 | | - writeShellScriptBin, |
13 | 8 | xz, |
14 | 9 | }: |
15 | 10 |
|
16 | 11 | { |
17 | 12 | version, |
18 | | - gpuTargets, |
19 | | - patches ? [ ], |
20 | | - src, |
21 | 13 | images, |
22 | | - extraPythonDepends ? ps: [ ], |
| 14 | + hashes, |
23 | 15 | }: |
24 | 16 |
|
25 | 17 | let |
26 | | - gpuTargets' = lib.concatStringsSep ";" gpuTargets; |
27 | | - compiler = "amdclang++"; |
| 18 | + rocmVersion = lib.versions.majorMinor rocm-core.version; |
| 19 | + hash = |
| 20 | + hashes.${rocmVersion} |
| 21 | + or (throw "aotriton ${version} binary package is not specified for ROCm ${rocmVersion}"); |
28 | 22 | in |
29 | | -stdenv.mkDerivation (finalAttrs: { |
| 23 | +stdenv.mkDerivation { |
30 | 24 | pname = "aotriton"; |
| 25 | + inherit version; |
31 | 26 |
|
32 | | - inherit version src patches; |
33 | | - |
34 | | - env = { |
35 | | - #CXX = compiler; |
36 | | - ROCM_PATH = "${rocmPackages.clr}"; |
37 | | - CFLAGS = "-w -g1 -gz -Wno-c++11-narrowing"; |
38 | | - CXXFLAGS = finalAttrs.env.CFLAGS; |
39 | | - |
40 | | - # aotriton passes a lot of files to the linker. |
41 | | - NIX_LD_USE_RESPONSE_FILE = 1; |
| 27 | + src = fetchurl { |
| 28 | + url = "https://github.com/ROCm/aotriton/releases/download/${version}/aotriton-${version}-manylinux_2_28_x86_64-rocm${rocmVersion}-shared.tar.gz"; |
| 29 | + inherit hash; |
42 | 30 | }; |
43 | 31 |
|
44 | | - requiredSystemFeatures = [ "big-parallel" ]; |
45 | | - |
46 | | - nativeBuildInputs = [ |
47 | | - cmake |
48 | | - jq |
49 | | - rocmPackages.rocm-cmake |
50 | | - pkg-config |
51 | | - python3 |
52 | | - ninja |
53 | | - rocmPackages.clr |
54 | | - writableTmpDirAsHomeHook # venv wants to cache in ~ |
55 | | - (writeShellScriptBin "amdclang++" '' |
56 | | - exec ${rocmPackages.llvm.clang}/bin/clang++ "$@" |
57 | | - '') |
58 | | - ]; |
59 | | - |
| 32 | + nativeBuildInputs = [ autoPatchelfHook ]; |
60 | 33 | buildInputs = [ |
61 | | - rocmPackages.clr |
| 34 | + clr |
| 35 | + stdenv.cc.cc.lib |
62 | 36 | xz |
63 | | - ] |
64 | | - ++ (with python3.pkgs; [ |
65 | | - wheel |
66 | | - packaging |
67 | | - pyyaml |
68 | | - numpy |
69 | | - filelock |
70 | | - iniconfig |
71 | | - pluggy |
72 | | - pybind11 |
73 | | - pandas |
74 | | - triton |
75 | | - ]); |
76 | | - |
77 | | - preConfigure = lib.optionalString (lib.versionAtLeast version "0.11.1") '' |
78 | | - # Since we use pre-built images, we can grab the image SHA from there. |
79 | | - # As of 0.11.1b this doesn't seem to be used for image loading yet, but |
80 | | - # just in case this happens in the future, we set this to the actual |
81 | | - # value and not a stub. |
82 | | - export AOTRITON_CI_SUPPLIED_SHA1=$(jq -r '.["AOTRITON_GIT_SHA1"]' ${images}/lib/aotriton.images/amd-gfx90a/__signature__) |
83 | | -
|
84 | | - # Need to set absolute paths to VENV and its PYTHON or |
85 | | - # build fails with "AOTRITON_INHERIT_SYSTEM_SITE_TRITON is enabled |
86 | | - # but triton is not available … no such file or directory" |
87 | | - # Set via a preConfigure hook so a valid absolute path can be |
88 | | - # picked if nix-shell is used against this package |
89 | | - cmakeFlagsArray+=( |
90 | | - "-DVENV_DIR=$(pwd)/aotriton-venv/" |
91 | | - "-DVENV_BIN_PYTHON=$(pwd)/aotriton-venv/bin/python" |
92 | | - ) |
93 | | - ''; |
| 37 | + ]; |
94 | 38 |
|
95 | | - # From README: |
96 | | - # Note: do not run ninja separately, due to the limit of the current build system, |
97 | | - # ninja install will run the whole build process unconditionally. |
| 39 | + dontConfigure = true; |
98 | 40 | dontBuild = true; |
99 | | - |
| 41 | + dontStrip = true; |
100 | 42 | installPhase = '' |
101 | 43 | runHook preInstall |
102 | | - ninja -v install |
103 | | - ln -sf ${images}/lib/aotriton.images $out/lib/aotriton.images |
104 | | - runHook postInstall |
105 | | - ''; |
106 | 44 |
|
107 | | - doCheck = false; |
108 | | - doInstallCheck = false; |
| 45 | + mkdir -p "$out" |
| 46 | + cp -r include lib "$out/" |
| 47 | + ln -s ${images}/lib/aotriton.images "$out/lib/aotriton.images" |
109 | 48 |
|
110 | | - cmakeFlags = [ |
111 | | - # Disable building kernels if no supported targets are enabled |
112 | | - (lib.cmakeBool "AOTRITON_NOIMAGE_MODE" true) |
113 | | - # Use preinstalled triton from our python's site-packages |
114 | | - (lib.cmakeBool "AOTRITON_INHERIT_SYSTEM_SITE_TRITON" true) |
115 | | - # Avoid kernels being skipped if build host is overloaded |
116 | | - (lib.cmakeFeature "AOTRITON_GPU_BUILD_TIMEOUT" "0") |
117 | | - (lib.cmakeFeature "CMAKE_CXX_COMPILER" compiler) |
118 | | - # Manually define CMAKE_INSTALL_<DIR> |
119 | | - # See: https://github.com/NixOS/nixpkgs/pull/197838 |
120 | | - (lib.cmakeFeature "CMAKE_INSTALL_BINDIR" "bin") |
121 | | - (lib.cmakeFeature "CMAKE_INSTALL_LIBDIR" "lib") |
122 | | - (lib.cmakeFeature "CMAKE_INSTALL_INCLUDEDIR" "include") |
123 | | - (lib.cmakeFeature "AOTRITON_TARGET_ARCH" gpuTargets') |
124 | | - (lib.cmakeBool "AOTRITON_USE_TORCH" false) |
125 | | - ]; |
| 49 | + runHook postInstall |
| 50 | + ''; |
126 | 51 |
|
127 | 52 | meta = with lib; { |
128 | 53 | description = "Ahead of Time (AOT) Triton Math Library"; |
129 | 54 | homepage = "https://github.com/ROCm/aotriton"; |
130 | 55 | license = with licenses; [ mit ]; |
131 | | - platforms = platforms.linux; |
132 | | - sourceProvenance = with sourceTypes; [ |
133 | | - fromSource |
134 | | - binaryNativeCode # aotriton.images |
135 | | - ]; |
| 56 | + platforms = [ "x86_64-linux" ]; |
| 57 | + sourceProvenance = with sourceTypes; [ binaryNativeCode ]; |
136 | 58 | }; |
137 | | -}) |
| 59 | +} |
0 commit comments