rocmPackages: cleanup and remove double building for rocblas

, gpuTargets ? [ "all" ]

}:

let

  # NOTE: Update the default GPU targets on every update

  gfx80 = (rocblas.override {

    gpuTargets = [

      "gfx803"

    ];

  }).overrideAttrs { pname = "rocblas-tensile-gfx80"; };

  gfx90 = (rocblas.override {

    gpuTargets = [

      "gfx900"

      "gfx906:xnack-"

      "gfx908:xnack-"

      "gfx90a:xnack+"

      "gfx90a:xnack-"

    ];

  }).overrideAttrs { pname = "rocblas-tensile-gfx90"; };

  gfx94 = (rocblas.override {

    gpuTargets = [

      "gfx940"

      "gfx941"

      "gfx942"

    ];

  }).overrideAttrs { pname = "rocblas-tensile-gfx94"; };

  gfx10 = (rocblas.override {

    gpuTargets = [

      "gfx1010"

      "gfx1012"

      "gfx1030"

    ];

  }).overrideAttrs { pname = "rocblas-tensile-gfx10"; };

  gfx11 = (rocblas.override {

    gpuTargets = [

      "gfx1100"

      "gfx1101"

      "gfx1102"

    ];

  }).overrideAttrs { pname = "rocblas-tensile-gfx11"; };

  # 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: {

stdenv.mkDerivation (finalAttrs: {

  pname = "rocblas";

  version = "6.0.2";

    cmake

    rocm-cmake

    clr

    tensile

  ];

  buildInputs = [

    (lib.cmakeFeature "AMDGPU_TARGETS" (lib.concatStringsSep ";" gpuTargets))

    (lib.cmakeBool "BUILD_WITH_TENSILE" buildTensile)

    (lib.cmakeBool "ROCM_SYMLINK_LIBS" false)

    # # 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"

    (lib.cmakeFeature "ROCBLAS_TENSILE_LIBRARY_DIR" "lib/rocblas")

    (lib.cmakeBool "BUILD_CLIENTS_TESTS" buildTests)

    (lib.cmakeBool "BUILD_CLIENTS_BENCHMARKS" buildBenchmarks)

  ] ++ lib.optionals buildTensile [

    (lib.cmakeFeature "VIRTUALENV_HOME_DIR" "/build/source/tensile")

    (lib.cmakeFeature "Tensile_TEST_LOCAL_PATH" "/build/source/tensile")

    (lib.cmakeFeature "Tensile_ROOT" "/build/source/tensile/${python3.sitePackages}/Tensile")

    (lib.cmakeBool "BUILD_WITH_PIP" false)

    (lib.cmakeFeature "Tensile_LOGIC" tensileLogic)

    (lib.cmakeFeature "Tensile_CODE_OBJECT_VERSION" tensileCOVersion)

    (lib.cmakeBool "Tensile_SEPARATE_ARCHITECTURES" tensileSepArch)

    })

  ];

  # Pass $NIX_BUILD_CORES to Tensile

  postPatch = ''

    substituteInPlace cmake/build-options.cmake \

      --replace-fail 'Tensile_CPU_THREADS ""' 'Tensile_CPU_THREADS "$ENV{NIX_BUILD_CORES}"'

  '' + 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 {

      url = "https://github.com/GZGavinZhao/Tensile/commit/855cb15839849addb0816a6dde45772034a3e41f.patch";

      hash = "sha256-d+fVf/vz+sxGqJ96vuxe0jRMgbC5K6j5FQ5SJ1e3Sl8=";

    })

    (fetchpatch {

      name = "Don-t-copy-file-twice-in-copyStaticFiles.patch";

      url = "https://github.com/GZGavinZhao/Tensile/commit/9e14d5a00a096bddac605910a0e4dfb4c35bb0d5.patch";

      hash = "sha256-gOzjJyD1K056OFQ+hK5nbUeBhxLTIgQLoT+0K12SypI=";

    })

  ];

  doCheck = false; # Too many errors, not sure how to set this up properly