diff --git a/pkgs/development/rocm-modules/5/default.nix b/pkgs/development/rocm-modules/5/default.nix
index 8bb9f03355b6..6e263998dbe6 100644
--- a/pkgs/development/rocm-modules/5/default.nix
+++ b/pkgs/development/rocm-modules/5/default.nix
@@ -193,7 +193,7 @@ in rec {
};
rocblas = callPackage ./rocblas {
- inherit rocmUpdateScript rocm-cmake clr tensile;
+ inherit rocblas rocmUpdateScript rocm-cmake clr tensile;
inherit (llvm) openmp;
stdenv = llvm.rocmClangStdenv;
};
diff --git a/pkgs/development/rocm-modules/5/rocblas/default.nix b/pkgs/development/rocm-modules/5/rocblas/default.nix
index 76dc38850d57..17732fd51681 100644
--- a/pkgs/development/rocm-modules/5/rocblas/default.nix
+++ b/pkgs/development/rocm-modules/5/rocblas/default.nix
@@ -1,4 +1,5 @@
-{ lib
+{ rocblas
+, lib
, stdenv
, fetchFromGitHub
, rocmUpdateScript
@@ -25,147 +26,54 @@
, tensileLibFormat ? "msgpack"
, gpuTargets ? [ "all" ]
}:
+
let
- rocblas = stdenv.mkDerivation (finalAttrs: {
- pname = "rocblas";
- version = "5.7.0";
-
- outputs = [
- "out"
- ] ++ lib.optionals buildTests [
- "test"
- ] ++ lib.optionals buildBenchmarks [
- "benchmark"
+ # NOTE: Update the default GPU targets on every update
+ gfx80 = (rocblas.override {
+ gpuTargets = [
+ "gfx803"
];
+ }).overrideAttrs { pname = "rocblas-tensile-gfx80"; };
- src = fetchFromGitHub {
- owner = "ROCmSoftwarePlatform";
- repo = "rocBLAS";
- rev = "rocm-${finalAttrs.version}";
- hash = "sha256-3wKnwvAra8u9xqlC05wUD+gSoBILTVJFU2cIV6xv3Lk=";
- };
-
- nativeBuildInputs = [
- cmake
- rocm-cmake
- clr
+ gfx90 = (rocblas.override {
+ gpuTargets = [
+ "gfx900"
+ "gfx906:xnack-"
+ "gfx908:xnack-"
+ "gfx90a:xnack+"
+ "gfx90a:xnack-"
];
+ }).overrideAttrs { pname = "rocblas-tensile-gfx90"; };
- buildInputs = [
- python3
- ] ++ lib.optionals buildTensile [
- msgpack
- libxml2
- python3Packages.msgpack
- python3Packages.joblib
- ] ++ lib.optionals buildTests [
- gtest
- ] ++ lib.optionals (buildTests || buildBenchmarks) [
- gfortran
- openmp
- amd-blis
- ] ++ lib.optionals (buildTensile || buildTests || buildBenchmarks) [
- python3Packages.pyyaml
+ gfx94 = (rocblas.override {
+ gpuTargets = [
+ "gfx940"
+ "gfx941"
+ "gfx942"
];
+ }).overrideAttrs { pname = "rocblas-tensile-gfx94"; };
- cmakeFlags = [
- "-DCMAKE_C_COMPILER=hipcc"
- "-DCMAKE_CXX_COMPILER=hipcc"
- "-Dpython=python3"
- "-DAMDGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
- "-DBUILD_WITH_TENSILE=${if buildTensile then "ON" else "OFF"}"
- # Manually define CMAKE_INSTALL_
- # See: https://github.com/NixOS/nixpkgs/pull/197838
- "-DCMAKE_INSTALL_BINDIR=bin"
- "-DCMAKE_INSTALL_LIBDIR=lib"
- "-DCMAKE_INSTALL_INCLUDEDIR=include"
- ] ++ lib.optionals buildTensile [
- "-DVIRTUALENV_HOME_DIR=/build/source/tensile"
- "-DTensile_TEST_LOCAL_PATH=/build/source/tensile"
- "-DTensile_ROOT=/build/source/tensile/lib/python${python3.pythonVersion}/site-packages/Tensile"
- "-DTensile_LOGIC=${tensileLogic}"
- "-DTensile_CODE_OBJECT_VERSION=${tensileCOVersion}"
- "-DTensile_SEPARATE_ARCHITECTURES=${if tensileSepArch then "ON" else "OFF"}"
- "-DTensile_LAZY_LIBRARY_LOADING=${if tensileLazyLib then "ON" else "OFF"}"
- "-DTensile_LIBRARY_FORMAT=${tensileLibFormat}"
- ] ++ lib.optionals buildTests [
- "-DBUILD_CLIENTS_TESTS=ON"
- ] ++ lib.optionals buildBenchmarks [
- "-DBUILD_CLIENTS_BENCHMARKS=ON"
- ] ++ lib.optionals (buildTests || buildBenchmarks) [
- "-DCMAKE_CXX_FLAGS=-I${amd-blis}/include/blis"
+ gfx10 = (rocblas.override {
+ gpuTargets = [
+ "gfx1010"
+ "gfx1012"
+ "gfx1030"
];
+ }).overrideAttrs { pname = "rocblas-tensile-gfx10"; };
- # Tensile REALLY wants to write to the nix directory if we include it normally
- postPatch = lib.optionalString buildTensile ''
- cp -a ${tensile} tensile
- chmod +w -R tensile
+ gfx11 = (rocblas.override {
+ gpuTargets = [
+ "gfx1100"
+ "gfx1101"
+ "gfx1102"
+ ];
+ }).overrideAttrs { pname = "rocblas-tensile-gfx11"; };
- # Rewrap Tensile
- substituteInPlace tensile/bin/{.t*,.T*,*} \
- --replace "${tensile}" "/build/source/tensile"
-
- substituteInPlace CMakeLists.txt \
- --replace "include(virtualenv)" "" \
- --replace "virtualenv_install(\''${Tensile_TEST_LOCAL_PATH})" ""
- '';
-
- postInstall = lib.optionalString buildTests ''
- mkdir -p $test/bin
- cp -a $out/bin/* $test/bin
- rm $test/bin/*-bench || true
- '' + lib.optionalString buildBenchmarks ''
- mkdir -p $benchmark/bin
- cp -a $out/bin/* $benchmark/bin
- rm $benchmark/bin/*-test || true
- '' + lib.optionalString (buildTests || buildBenchmarks ) ''
- rm -rf $out/bin
- '';
-
- passthru.updateScript = rocmUpdateScript {
- name = finalAttrs.pname;
- owner = finalAttrs.src.owner;
- repo = finalAttrs.src.repo;
- };
-
- requiredSystemFeatures = [ "big-parallel" ];
-
- meta = with lib; {
- description = "BLAS implementation for ROCm platform";
- homepage = "https://github.com/ROCmSoftwarePlatform/rocBLAS";
- license = with licenses; [ mit ];
- maintainers = teams.rocm.members;
- platforms = platforms.linux;
- broken = versions.minor finalAttrs.version != versions.minor stdenv.cc.version;
- };
- });
-
- gfx80 = runCommand "rocblas-gfx80" { preferLocalBuild = true; } ''
- mkdir -p $out/lib/rocblas/library
- cp -a ${rocblas}/lib/rocblas/library/*gfx80* $out/lib/rocblas/library
- '';
-
- gfx90 = runCommand "rocblas-gfx90" { preferLocalBuild = true; } ''
- mkdir -p $out/lib/rocblas/library
- cp -a ${rocblas}/lib/rocblas/library/*gfx90* $out/lib/rocblas/library
- '';
-
- gfx94 = runCommand "rocblas-gfx94" { preferLocalBuild = true; } ''
- mkdir -p $out/lib/rocblas/library
- cp -a ${rocblas}/lib/rocblas/library/*gfx94* $out/lib/rocblas/library
- '';
-
- gfx10 = runCommand "rocblas-gfx10" { preferLocalBuild = true; } ''
- mkdir -p $out/lib/rocblas/library
- cp -a ${rocblas}/lib/rocblas/library/*gfx10* $out/lib/rocblas/library
- '';
-
- gfx11 = runCommand "rocblas-gfx11" { preferLocalBuild = true; } ''
- mkdir -p $out/lib/rocblas/library
- cp -a ${rocblas}/lib/rocblas/library/*gfx11* $out/lib/rocblas/library
- '';
+ # Unfortunately, we have to do two full builds, otherwise we get overlapping _fallback.dat files
+ fallbacks = rocblas.overrideAttrs { pname = "rocblas-tensile-fallbacks"; };
in stdenv.mkDerivation (finalAttrs: {
- inherit (rocblas) pname version src passthru meta;
+ pname = "rocblas";
+ version = "5.7.0";
outputs = [
"out"
@@ -175,26 +83,127 @@ in stdenv.mkDerivation (finalAttrs: {
"benchmark"
];
- dontUnpack = true;
- dontPatch = true;
- dontConfigure = true;
- dontBuild = true;
+ src = fetchFromGitHub {
+ owner = "ROCmSoftwarePlatform";
+ repo = "rocBLAS";
+ rev = "rocm-${finalAttrs.version}";
+ hash = "sha256-3wKnwvAra8u9xqlC05wUD+gSoBILTVJFU2cIV6xv3Lk=";
+ };
- installPhase = ''
- runHook preInstall
+ nativeBuildInputs = [
+ cmake
+ rocm-cmake
+ clr
+ ];
- mkdir -p $out
- cp -a --no-preserve=mode ${rocblas}/* $out
- ln -sf ${gfx80}/lib/rocblas/library/* $out/lib/rocblas/library
- ln -sf ${gfx90}/lib/rocblas/library/* $out/lib/rocblas/library
- ln -sf ${gfx94}/lib/rocblas/library/* $out/lib/rocblas/library
- ln -sf ${gfx10}/lib/rocblas/library/* $out/lib/rocblas/library
- ln -sf ${gfx11}/lib/rocblas/library/* $out/lib/rocblas/library
- '' + lib.optionalString buildTests ''
- cp -a ${rocblas.test} $test
- '' + lib.optionalString buildBenchmarks ''
- cp -a ${rocblas.benchmark} $benchmark
- '' + ''
- runHook postInstall
+ buildInputs = [
+ python3
+ ] ++ lib.optionals buildTensile [
+ msgpack
+ libxml2
+ python3Packages.msgpack
+ python3Packages.joblib
+ ] ++ lib.optionals buildTests [
+ gtest
+ ] ++ lib.optionals (buildTests || buildBenchmarks) [
+ gfortran
+ openmp
+ amd-blis
+ ] ++ lib.optionals (buildTensile || buildTests || buildBenchmarks) [
+ python3Packages.pyyaml
+ ];
+
+ cmakeFlags = [
+ "-DCMAKE_C_COMPILER=hipcc"
+ "-DCMAKE_CXX_COMPILER=hipcc"
+ "-Dpython=python3"
+ "-DAMDGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
+ "-DBUILD_WITH_TENSILE=${if buildTensile then "ON" else "OFF"}"
+ # Manually define CMAKE_INSTALL_
+ # See: https://github.com/NixOS/nixpkgs/pull/197838
+ "-DCMAKE_INSTALL_BINDIR=bin"
+ "-DCMAKE_INSTALL_LIBDIR=lib"
+ "-DCMAKE_INSTALL_INCLUDEDIR=include"
+ ] ++ lib.optionals buildTensile [
+ "-DVIRTUALENV_HOME_DIR=/build/source/tensile"
+ "-DTensile_TEST_LOCAL_PATH=/build/source/tensile"
+ "-DTensile_ROOT=/build/source/tensile/${python3.sitePackages}/Tensile"
+ "-DTensile_LOGIC=${tensileLogic}"
+ "-DTensile_CODE_OBJECT_VERSION=${tensileCOVersion}"
+ "-DTensile_SEPARATE_ARCHITECTURES=${if tensileSepArch then "ON" else "OFF"}"
+ "-DTensile_LAZY_LIBRARY_LOADING=${if tensileLazyLib then "ON" else "OFF"}"
+ "-DTensile_LIBRARY_FORMAT=${tensileLibFormat}"
+ ] ++ lib.optionals buildTests [
+ "-DBUILD_CLIENTS_TESTS=ON"
+ ] ++ lib.optionals buildBenchmarks [
+ "-DBUILD_CLIENTS_BENCHMARKS=ON"
+ ] ++ lib.optionals (buildTests || buildBenchmarks) [
+ "-DCMAKE_CXX_FLAGS=-I${amd-blis}/include/blis"
+ ];
+
+ postPatch = lib.optionalString (finalAttrs.pname != "rocblas") ''
+ # Return early and install tensile files manually
+ substituteInPlace library/src/CMakeLists.txt \
+ --replace "set_target_properties( TensileHost PROPERTIES OUTPUT_NAME" "return()''\nset_target_properties( TensileHost PROPERTIES OUTPUT_NAME"
+ '' + lib.optionalString (buildTensile && finalAttrs.pname == "rocblas") ''
+ # Link the prebuilt Tensile files
+ mkdir -p build/Tensile/library
+
+ for path in ${gfx80} ${gfx90} ${gfx94} ${gfx10} ${gfx11} ${fallbacks}; do
+ ln -s $path/lib/rocblas/library/* build/Tensile/library
+ done
+
+ unlink build/Tensile/library/TensileManifest.txt
+ '' + lib.optionalString buildTensile ''
+ # Tensile REALLY wants to write to the nix directory if we include it normally
+ cp -a ${tensile} tensile
+ chmod +w -R tensile
+
+ # Rewrap Tensile
+ substituteInPlace tensile/bin/{.t*,.T*,*} \
+ --replace "${tensile}" "/build/source/tensile"
+
+ substituteInPlace CMakeLists.txt \
+ --replace "include(virtualenv)" "" \
+ --replace "virtualenv_install(\''${Tensile_TEST_LOCAL_PATH})" ""
'';
+
+ postInstall = lib.optionalString (finalAttrs.pname == "rocblas") ''
+ ln -sf ${fallbacks}/lib/rocblas/library/TensileManifest.txt $out/lib/rocblas/library
+ '' + lib.optionalString (finalAttrs.pname != "rocblas") ''
+ mkdir -p $out/lib/rocblas/library
+ rm -rf $out/share
+ '' + lib.optionalString (finalAttrs.pname != "rocblas" && finalAttrs.pname != "rocblas-tensile-fallbacks") ''
+ rm Tensile/library/{TensileManifest.txt,*_fallback.dat}
+ mv Tensile/library/* $out/lib/rocblas/library
+ '' + lib.optionalString (finalAttrs.pname == "rocblas-tensile-fallbacks") ''
+ mv Tensile/library/{TensileManifest.txt,*_fallback.dat} $out/lib/rocblas/library
+ '' + lib.optionalString buildTests ''
+ mkdir -p $test/bin
+ cp -a $out/bin/* $test/bin
+ rm $test/bin/*-bench || true
+ '' + lib.optionalString buildBenchmarks ''
+ mkdir -p $benchmark/bin
+ cp -a $out/bin/* $benchmark/bin
+ rm $benchmark/bin/*-test || true
+ '' + lib.optionalString (buildTests || buildBenchmarks ) ''
+ rm -rf $out/bin
+ '';
+
+ passthru.updateScript = rocmUpdateScript {
+ name = finalAttrs.pname;
+ owner = finalAttrs.src.owner;
+ repo = finalAttrs.src.repo;
+ };
+
+ requiredSystemFeatures = [ "big-parallel" ];
+
+ meta = with lib; {
+ description = "BLAS implementation for ROCm platform";
+ homepage = "https://github.com/ROCmSoftwarePlatform/rocBLAS";
+ license = with licenses; [ mit ];
+ maintainers = teams.rocm.members;
+ platforms = platforms.linux;
+ broken = versions.minor finalAttrs.version != versions.minor stdenv.cc.version;
+ };
})