rocmPackages.aotriton: 0.9.2b -> 0.10b (#439056)
This commit is contained in:
Peder Bergebakken Sundt
@@ -2,226 +2,173 @@
|
||||
lib,
|
||||
stdenv,
|
||||
fetchFromGitHub,
|
||||
fetchpatch,
|
||||
cmake,
|
||||
rocm-cmake,
|
||||
clr,
|
||||
rocblas,
|
||||
rocsolver,
|
||||
gtest,
|
||||
msgpack,
|
||||
libxml2,
|
||||
python3,
|
||||
python3Packages,
|
||||
openmp,
|
||||
hipblas-common,
|
||||
hipblas,
|
||||
nlohmann_json,
|
||||
triton-llvm,
|
||||
rocmlir,
|
||||
lapack-reference,
|
||||
ninja,
|
||||
ncurses,
|
||||
libffi,
|
||||
zlib,
|
||||
zstd,
|
||||
xz,
|
||||
writableTmpDirAsHomeHook,
|
||||
pkg-config,
|
||||
buildTests ? false,
|
||||
buildBenchmarks ? false,
|
||||
buildSamples ? false,
|
||||
gpuTargets ? [
|
||||
gpuTargets ? clr.localGpuTargets or clr.gpuTargets,
|
||||
# for passthru.tests
|
||||
aotriton,
|
||||
hello,
|
||||
}:
|
||||
let
|
||||
supportedTargets = lib.lists.intersectLists [
|
||||
# aotriton GPU support list:
|
||||
# https://github.com/ROCm/aotriton/blob/main/v2python/gpu_targets.py
|
||||
"gfx90a"
|
||||
"gfx942"
|
||||
"gfx950"
|
||||
"gfx1100"
|
||||
"gfx1101"
|
||||
],
|
||||
}:
|
||||
"gfx1151"
|
||||
"gfx1150"
|
||||
"gfx1201"
|
||||
"gfx1200"
|
||||
] gpuTargets;
|
||||
supportedTargets' = lib.concatStringsSep ";" supportedTargets;
|
||||
anySupportedTargets = supportedTargets != [ ];
|
||||
in
|
||||
stdenv.mkDerivation (finalAttrs: {
|
||||
pname = "aotriton";
|
||||
version = "0.10b";
|
||||
|
||||
stdenv.mkDerivation (
|
||||
finalAttrs:
|
||||
let
|
||||
py = python3.withPackages (ps: [
|
||||
ps.pyyaml
|
||||
ps.distutils
|
||||
ps.setuptools
|
||||
ps.packaging
|
||||
ps.numpy
|
||||
ps.wheel
|
||||
ps.filelock
|
||||
ps.iniconfig
|
||||
ps.pluggy
|
||||
ps.pybind11
|
||||
]);
|
||||
gpuTargets' = lib.concatStringsSep ";" gpuTargets;
|
||||
compiler = "amdclang++";
|
||||
cFlags = "-O3 -DNDEBUG";
|
||||
cxxFlags = "${cFlags} -Wno-c++11-narrowing";
|
||||
triton-llvm' = triton-llvm;
|
||||
in
|
||||
{
|
||||
pname = "aotriton";
|
||||
version = "0.9.2b";
|
||||
|
||||
src = fetchFromGitHub {
|
||||
owner = "ROCm";
|
||||
repo = "aotriton";
|
||||
rev = "${finalAttrs.version}";
|
||||
hash = "sha256-1Cf0olD3zRg9JESD6s/WaGifm3kfD12VUvjTZHpmGAE=";
|
||||
fetchSubmodules = true;
|
||||
};
|
||||
env.CXX = compiler;
|
||||
env.ROCM_PATH = "${clr}";
|
||||
requiredSystemFeatures = [ "big-parallel" ];
|
||||
|
||||
outputs = [
|
||||
"out"
|
||||
]
|
||||
++ lib.optionals buildTests [
|
||||
"test"
|
||||
]
|
||||
++ lib.optionals buildBenchmarks [
|
||||
"benchmark"
|
||||
]
|
||||
++ lib.optionals buildSamples [
|
||||
"sample"
|
||||
];
|
||||
|
||||
# Need an empty cuda.h for this to compile
|
||||
# Better than pulling in unfree cuda headers
|
||||
postPatch = ''
|
||||
touch third_party/triton/third_party/nvidia/include/cuda.h
|
||||
src = fetchFromGitHub {
|
||||
owner = "ROCm";
|
||||
repo = "aotriton";
|
||||
tag = finalAttrs.version;
|
||||
hash = "sha256-stAHnsqChkNv69wjlhM/qUetrJpNwI1i7rGnPMwsNz0=";
|
||||
leaveDotGit = true;
|
||||
# fetch all submodules except unused triton submodule that is ~500MB
|
||||
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
|
||||
'';
|
||||
};
|
||||
|
||||
doCheck = false;
|
||||
doInstallCheck = false;
|
||||
cmakeBuildType = "RelWithDebInfo";
|
||||
separateDebugInfo = true;
|
||||
__structuredAttrs = true;
|
||||
strictDeps = true;
|
||||
# Only set big-parallel when we are building kernels, no-image mode build is faster
|
||||
requiredSystemFeatures = if anySupportedTargets then [ "big-parallel" ] else [ ];
|
||||
|
||||
nativeBuildInputs = [
|
||||
cmake
|
||||
rocm-cmake
|
||||
pkg-config
|
||||
py
|
||||
clr
|
||||
ninja
|
||||
];
|
||||
env = {
|
||||
ROCM_PATH = "${clr}";
|
||||
CFLAGS = "-w -g1 -gz -Wno-c++11-narrowing";
|
||||
CXXFLAGS = finalAttrs.env.CFLAGS;
|
||||
};
|
||||
|
||||
buildInputs = [
|
||||
rocblas
|
||||
rocsolver
|
||||
hipblas-common
|
||||
hipblas
|
||||
openmp
|
||||
libffi
|
||||
ncurses
|
||||
xz
|
||||
nlohmann_json
|
||||
rocmlir
|
||||
nativeBuildInputs = [
|
||||
cmake
|
||||
rocm-cmake
|
||||
pkg-config
|
||||
python3
|
||||
ninja
|
||||
clr
|
||||
writableTmpDirAsHomeHook # venv wants to cache in ~
|
||||
];
|
||||
|
||||
msgpack
|
||||
libxml2
|
||||
python3Packages.msgpack
|
||||
zlib
|
||||
zstd
|
||||
]
|
||||
++ lib.optionals buildTests [
|
||||
gtest
|
||||
]
|
||||
++ lib.optionals (buildTests || buildBenchmarks) [
|
||||
lapack-reference
|
||||
];
|
||||
buildInputs = [
|
||||
clr
|
||||
xz
|
||||
]
|
||||
++ (with python3.pkgs; [
|
||||
wheel
|
||||
packaging
|
||||
pyyaml
|
||||
numpy
|
||||
filelock
|
||||
iniconfig
|
||||
pluggy
|
||||
pybind11
|
||||
pandas
|
||||
triton
|
||||
]);
|
||||
|
||||
env.TRITON_OFFLINE_BUILD = 1;
|
||||
env.LLVM_SYSPATH = "${triton-llvm'}";
|
||||
env.JSON_SYSPATH = nlohmann_json;
|
||||
env.MLIR_DIR = "${triton-llvm'}/lib/cmake/mlir";
|
||||
env.CXXFLAGS = "-I/build/source/third_party/triton/third_party/nvidia/backend/include";
|
||||
patches = [
|
||||
# CMakeLists.txt: AOTRITON_INHERIT_SYSTEM_SITE_TRITON flag
|
||||
(fetchpatch {
|
||||
url = "https://github.com/ROCm/aotriton/commit/9734c3e999c412a07d2b35671998650942b26ed4.patch";
|
||||
hash = "sha256-tBmjjhRJmLv3K6F2+4OcMuwf8dH7efPPECMQjh6QdUA=";
|
||||
})
|
||||
];
|
||||
|
||||
# Fix up header issues in triton: https://github.com/triton-lang/triton/pull/3985/files
|
||||
preConfigure = ''
|
||||
mkdir third_party/triton/third_party/nvidia/backend/include/
|
||||
touch third_party/triton/third_party/nvidia/backend/include/cuda.h
|
||||
find third_party/triton -type f -exec sed -i 's|[<]cupti.h[>]|"cupti.h"|g' {} +
|
||||
find third_party/triton -type f -exec sed -i 's|[<]cuda.h[>]|"cuda.h"|g' {} +
|
||||
# Excerpt 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.
|
||||
dontBuild = true;
|
||||
# This builds+installs
|
||||
installPhase = ''
|
||||
runHook preInstall
|
||||
ninja -v install
|
||||
runHook postInstall
|
||||
'';
|
||||
# tests are intended to be ran manually as test/ python scripts and need accelerator
|
||||
doCheck = false;
|
||||
doInstallCheck = false;
|
||||
|
||||
sed -i '2s;^;set(CMAKE_SUPPRESS_DEVELOPER_WARNINGS ON CACHE BOOL "ON")\n;' CMakeLists.txt
|
||||
sed -i '2s;^;set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON")\n;' CMakeLists.txt
|
||||
sed -i '2s;^;set(CMAKE_SUPPRESS_DEVELOPER_WARNINGS ON CACHE BOOL "ON")\n;' third_party/triton/CMakeLists.txt
|
||||
sed -i '2s;^;set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON")\n;' third_party/triton/CMakeLists.txt
|
||||
substituteInPlace third_party/triton/python/setup.py \
|
||||
--replace-fail "from distutils.command.clean import clean" "import setuptools;from distutils.command.clean import clean" \
|
||||
--replace-fail 'system == "Linux"' 'False'
|
||||
# Fix 'ld: error: unable to insert .comment after .comment'
|
||||
substituteInPlace v2python/ld_script.py \
|
||||
--replace-fail 'INSERT AFTER .comment;' ""
|
||||
# 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
|
||||
preConfigure = ''
|
||||
cmakeFlagsArray+=(
|
||||
"-DVENV_DIR=$(pwd)/aotriton-venv/"
|
||||
"-DVENV_BIN_PYTHON=$(pwd)/aotriton-venv/bin/python"
|
||||
)
|
||||
'';
|
||||
|
||||
cmakeFlagsArray+=(
|
||||
'-DCMAKE_C_FLAGS_RELEASE=${cFlags}'
|
||||
'-DCMAKE_CXX_FLAGS_RELEASE=${cxxFlags}'
|
||||
)
|
||||
prependToVar cmakeFlags "-GNinja"
|
||||
mkdir -p /build/tmp-home
|
||||
export HOME=/build/tmp-home
|
||||
'';
|
||||
cmakeFlags = [
|
||||
# Disable building kernels if no supported targets are enabled
|
||||
(lib.cmakeBool "AOTRITON_NOIMAGE_MODE" (!anySupportedTargets))
|
||||
# Use preinstalled triton from our python's site-packages
|
||||
(lib.cmakeBool "AOTRITON_INHERIT_SYSTEM_SITE_TRITON" true)
|
||||
# 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.optionals anySupportedTargets [
|
||||
# Note: build will warn "AMDGPU_TARGETS was not set, and system GPU detection was unsuccsesful."
|
||||
# but this can safely be ignored, aotriton uses a different approach to pass targets
|
||||
(lib.cmakeFeature "AOTRITON_TARGET_ARCH" supportedTargets')
|
||||
];
|
||||
|
||||
# Excerpt 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.
|
||||
dontBuild = true;
|
||||
|
||||
installPhase = ''
|
||||
runHook preInstall
|
||||
ninja -v install
|
||||
runHook postInstall
|
||||
'';
|
||||
|
||||
cmakeFlags = [
|
||||
"-Wno-dev"
|
||||
"-DAOTRITON_NOIMAGE_MODE=ON" # FIXME: Should be able to build with object code but generate_shim is failing
|
||||
"-DCMAKE_BUILD_TYPE=Release"
|
||||
"-DCMAKE_VERBOSE_MAKEFILE=ON"
|
||||
"-DVIRTUALENV_PYTHON_EXENAME=${lib.getExe py}"
|
||||
"-DCMAKE_CXX_COMPILER=${compiler}"
|
||||
# Manually define CMAKE_INSTALL_<DIR>
|
||||
# See: https://github.com/NixOS/nixpkgs/pull/197838
|
||||
"-DCMAKE_INSTALL_BINDIR=bin"
|
||||
"-DCMAKE_INSTALL_LIBDIR=lib"
|
||||
"-DCMAKE_INSTALL_INCLUDEDIR=include"
|
||||
"-DAMDGPU_TARGETS=${gpuTargets'}"
|
||||
"-DGPU_TARGETS=${gpuTargets'}"
|
||||
]
|
||||
++ lib.optionals buildTests [
|
||||
"-DBUILD_CLIENTS_TESTS=ON"
|
||||
]
|
||||
++ lib.optionals buildBenchmarks [
|
||||
"-DBUILD_CLIENTS_BENCHMARKS=ON"
|
||||
]
|
||||
++ lib.optionals buildSamples [
|
||||
"-DBUILD_CLIENTS_SAMPLES=ON"
|
||||
];
|
||||
|
||||
postInstall =
|
||||
lib.optionalString buildTests ''
|
||||
mkdir -p $test/bin
|
||||
mv $out/bin/hipblas-test $test/bin
|
||||
''
|
||||
+ lib.optionalString buildBenchmarks ''
|
||||
mkdir -p $benchmark/bin
|
||||
mv $out/bin/hipblas-bench $benchmark/bin
|
||||
''
|
||||
+ lib.optionalString buildSamples ''
|
||||
mkdir -p $sample/bin
|
||||
mv $out/bin/example-* $sample/bin
|
||||
''
|
||||
+ lib.optionalString (buildTests || buildBenchmarks || buildSamples) ''
|
||||
rmdir $out/bin
|
||||
passthru.tests = {
|
||||
# regression test that aotriton so doesn't crash in static constructor
|
||||
# currently known to fail on rocm toolchain but fine with default stdenv
|
||||
ld-preload-into-hello = stdenv.mkDerivation {
|
||||
name = "aotriton-basic-load-test";
|
||||
nativeBuildInputs = [ hello ];
|
||||
buildCommand = ''
|
||||
set -e
|
||||
LD_PRELOAD=${
|
||||
aotriton.override {
|
||||
gpuTargets = [ ];
|
||||
}
|
||||
}/lib/libaotriton_v2.so ${hello}/bin/hello > /dev/null
|
||||
echo "ld-preload-into-hello" > $out
|
||||
'';
|
||||
meta = with lib; {
|
||||
description = "ROCm Ahead of Time (AOT) Triton Math Library ";
|
||||
homepage = "https://github.com/ROCm/aotriton";
|
||||
license = with licenses; [ mit ];
|
||||
teams = [ teams.rocm ];
|
||||
platforms = platforms.linux;
|
||||
};
|
||||
}
|
||||
)
|
||||
};
|
||||
|
||||
meta = {
|
||||
description = "ROCm Ahead of Time (AOT) Triton Math Library";
|
||||
homepage = "https://github.com/ROCm/aotriton";
|
||||
license = lib.licenses.mit;
|
||||
teams = [ lib.teams.rocm ];
|
||||
platforms = lib.platforms.linux;
|
||||
# ld: error: unable to insert .comment after .comment
|
||||
broken = stdenv.cc.isClang;
|
||||
};
|
||||
})
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
python3Packages,
|
||||
triton-llvm,
|
||||
openmpi,
|
||||
stdenv,
|
||||
rocmGpuArches ? [ ],
|
||||
}:
|
||||
|
||||
@@ -83,7 +84,7 @@ let
|
||||
# Replaces hip, opencl-runtime, and rocclr
|
||||
clr = self.callPackage ./clr { };
|
||||
|
||||
aotriton = self.callPackage ./aotriton { };
|
||||
aotriton = self.callPackage ./aotriton { inherit stdenv; };
|
||||
|
||||
hipify = self.callPackage ./hipify {
|
||||
inherit (llvm)
|
||||
|
||||
Reference in New Issue
Block a user