diff --git a/pkgs/development/rocm-modules/6/aotriton/default.nix b/pkgs/development/rocm-modules/6/aotriton/default.nix index ccadbbf9633e..26b22e61d8c1 100644 --- a/pkgs/development/rocm-modules/6/aotriton/default.nix +++ b/pkgs/development/rocm-modules/6/aotriton/default.nix @@ -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_ + # 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_ - # 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; + }; +}) diff --git a/pkgs/development/rocm-modules/6/default.nix b/pkgs/development/rocm-modules/6/default.nix index a113ce0925ac..73e1b0cb1c20 100644 --- a/pkgs/development/rocm-modules/6/default.nix +++ b/pkgs/development/rocm-modules/6/default.nix @@ -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)