Merge pull request #283777 from ConnorBaker/fix/magma-update-cmake-options

#  supportedGpuTargets: List String

# }

{ blas

{ autoPatchelfHook

, blas

, cmake

, cudaPackages

, cudaSupport ? config.cudaSupport

, libpthreadstubs

, magmaRelease

, ninja

, python3

, config

  # At least one back-end has to be enabled,

  # and we can't default to CUDA since it's unfree

, rocmSupport ? !cudaSupport

, static ? stdenv.hostPlatform.isStatic

, stdenv

, symlinkJoin

}:

let

  inherit (lib) lists strings trivial;

  inherit (cudaPackages) backendStdenv cudaFlags cudaVersion;

  inherit (cudaPackages) cudaAtLeast cudaFlags cudaOlder;

  inherit (magmaRelease) version hash supportedGpuTargets;

  # NOTE: The lists.subtractLists function is perhaps a bit unintuitive. It subtracts the elements

    inherit hash;

  };

  # Magma doesn't have anything which could be run under doCheck, but it does build test suite executables.

  # These are moved to $test/bin/ and $test/lib/ in postInstall.

  outputs = ["out" "test"];

  # Fixup for the python test runners

  postPatch = ''

    patchShebangs ./testing/run_{tests,summarize}.py

    substituteInPlace ./testing/run_tests.py \

      --replace-fail \

        "print >>sys.stderr, cmdp, \"doesn't exist (original name: \" + cmd + \", precision: \" + precision + \")\"" \

        "print(f\"{cmdp} doesn't exist (original name: {cmd}, precision: {precision})\", file=sys.stderr)"

  '';

  nativeBuildInputs = [

    autoPatchelfHook

    cmake

    ninja

    gfortran

    libpthreadstubs

    lapack

    blas

    python3

  ] ++ lists.optionals cudaSupport (with cudaPackages; [

    cuda_cudart.dev # cuda_runtime.h

    cuda_cudart.lib # cudart

    libcublas.lib # cublas

    libcusparse.dev # cusparse.h

    libcusparse.lib # cusparse

  ] ++ lists.optionals (strings.versionOlder cudaVersion "11.8") [

  ] ++ lists.optionals (cudaOlder "11.8") [

    cuda_nvprof.dev # <cuda_profiler_api.h>

  ] ++ lists.optionals (strings.versionAtLeast cudaVersion "11.8") [

  ] ++ lists.optionals (cudaAtLeast "11.8") [

    cuda_profiler_api.dev # <cuda_profiler_api.h>

  ] ++ lists.optionals (strings.versionAtLeast cudaVersion "12.0") [

  ] ++ lists.optionals (cudaAtLeast "12.0") [

    cuda_cccl.dev # <nv/target>

  ]) ++ lists.optionals rocmSupport [

    rocmPackages.clr

  ];

  cmakeFlags = [

    "-DGPU_TARGET=${gpuTargetString}"

    (lib.cmakeBool "MAGMA_ENABLE_CUDA" cudaSupport)

    (lib.cmakeBool "MAGMA_ENABLE_HIP" rocmSupport)

  ] ++ lists.optionals static [

    "-DBUILD_SHARED_LIBS=OFF"

    (strings.cmakeFeature "GPU_TARGET" gpuTargetString)

    (strings.cmakeBool "MAGMA_ENABLE_CUDA" cudaSupport)

    (strings.cmakeBool "MAGMA_ENABLE_HIP" rocmSupport)

    (strings.cmakeBool "BUILD_SHARED_LIBS" (!static))

    # Set the Fortran name mangling scheme explicitly. We must set FORTRAN_CONVENTION manually because it will

    # otherwise not be set in NVCC_FLAGS or DEVCCFLAGS (which we cannot modify).

    # See https://github.com/NixOS/nixpkgs/issues/281656#issuecomment-1902931289

    (strings.cmakeBool "USE_FORTRAN" true)

    (strings.cmakeFeature "CMAKE_C_FLAGS" "-DADD_")

    (strings.cmakeFeature "CMAKE_CXX_FLAGS" "-DADD_")

    (strings.cmakeFeature "FORTRAN_CONVENTION" "-DADD_")

  ] ++ lists.optionals cudaSupport [

    "-DCMAKE_CUDA_ARCHITECTURES=${cudaArchitecturesString}"

    "-DMIN_ARCH=${minArch}" # Disarms magma's asserts

    "-DCMAKE_C_COMPILER=${backendStdenv.cc}/bin/cc"

    "-DCMAKE_CXX_COMPILER=${backendStdenv.cc}/bin/c++"

    (strings.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" cudaArchitecturesString)

    (strings.cmakeFeature "MIN_ARCH" minArch) # Disarms magma's asserts

  ] ++ lists.optionals rocmSupport [

    "-DCMAKE_C_COMPILER=${rocmPackages.clr}/bin/hipcc"

    "-DCMAKE_CXX_COMPILER=${rocmPackages.clr}/bin/hipcc"

  ] ++ lists.optionals (cudaPackages.cudaAtLeast "12.0.0") [

    (lib.cmakeBool "USE_FORTRAN" false)

  ];

  buildFlags = [

    "magma"

    "magma_sparse"

    (strings.cmakeFeature "CMAKE_C_COMPILER" "${rocmPackages.clr}/bin/hipcc")

    (strings.cmakeFeature "CMAKE_CXX_COMPILER" "${rocmPackages.clr}/bin/hipcc")

  ];

  # Magma doesn't have a test suite we can easily run, just loose executables, all of which require a GPU.

  doCheck = false;

  # Copy the files to the test output and fix the RPATHs.

  postInstall =

    # NOTE: The python scripts aren't copied by CMake into the build directory, so we must copy them from the source.

    # TODO(@connorbaker): This should be handled by having CMakeLists.txt install them, but such a patch is

    # out of the scope of the PR which introduces the `test` output: https://github.com/NixOS/nixpkgs/pull/283777.

    # See https://github.com/NixOS/nixpkgs/pull/283777#discussion_r1482125034 for more information.

    ''

      install -Dm755 ../testing/run_{tests,summarize}.py -t "$test/bin/"

    ''

    # Copy core test executables and libraries over to the test output.

    # NOTE: Magma doesn't provide tests for sparse solvers for ROCm, but it does for CUDA -- we put them both in the same

    # install command to avoid the case where a glob would fail to find any files and cause the install command to fail

    # because it has no files to install.

    + ''

      install -Dm755 ./testing/testing_* ./sparse/testing/testing_* -t "$test/bin/"

      install -Dm755 ./lib/libtester.so ./lib/liblapacktest.so -t "$test/lib/"

    ''

    # All of the test executables and libraries will have a reference to the build directory in their RPATH, which we

    # must remove. We do this by shrinking the RPATH to only include the Nix store. The autoPatchelfHook will take care

    # of supplying the correct RPATH for needed libraries (like `libtester.so`).

    + ''

      find "$test" -type f -exec \

        patchelf \

          --shrink-rpath \

          --allowed-rpath-prefixes "$NIX_STORE" \

          {} \;

    '';

  passthru = {

    inherit cudaPackages cudaSupport rocmSupport gpuTargets;

  };

    broken =

      !(cudaSupport || rocmSupport) # At least one back-end enabled

      || (cudaSupport && rocmSupport) # Mutually exclusive

      || (cudaSupport && strings.versionOlder cudaVersion "9");

      || (cudaSupport && cudaOlder "9.0");

  };

}