Unverified Commit f6a1b39d authored by figsoda's avatar figsoda Committed by GitHub
Browse files

Merge pull request #270448 from NixOS/backport-268746-to-release-23.11 

[Backport release-23.11] python311Packages.torch: choose magma at the expression level
parents d3aefc17 cb1d60fa

nixos/doc/manual/release-notes/rl-2311.section.md

+3 −0
Original line number Diff line number Diff line
@@ -605,6 +605,9 @@ The module update takes care of the new config syntax and the data itself (user 


- `python3.pkgs.flitBuildHook` has been removed. Use `flit-core` and `format = "pyproject"` instead.



-  Now `magma` defaults to `magma-hip` instead of `magma-cuda`. It also

   respects the `config.cudaSupport` and `config.rocmSupport` options.



- The `extend` function of `llvmPackages` has been removed due it coming from the `tools` attrset thus only extending the `tool` attrset. A possible replacement is to construct the set from `libraries` and `tools`, or patch nixpkgs.



- The `qemu-vm.nix` module now supports disabling overriding `fileSystems` with

pkgs/development/libraries/science/math/magma/generic.nix

+14 −13
Original line number Diff line number Diff line
@@ -8,12 +8,7 @@ 
{ blas

, cmake

, cudaPackages

  # FIXME: cuda being unfree means ofborg won't eval "magma".

  # respecting config.cudaSupport -> false by default

  # -> ofborg eval -> throws "no GPU targets specified".

  # Probably should delete everything but "magma-cuda" and "magma-hip"

  # from all-packages.nix

, cudaSupport ? true

, cudaSupport ? config.cudaSupport

, fetchurl

, gfortran

, cudaCapabilities ? cudaPackages.cudaFlags.cudaCapabilities
@@ -25,8 +20,10 @@ 
, magmaRelease

, ninja

, config

, rocmSupport ? config.rocmSupport

, static ? false

  # 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

}:
@@ -133,6 +130,8 @@ stdenv.mkDerivation { 


  cmakeFlags = [

    "-DGPU_TARGET=${gpuTargetString}"

    (lib.cmakeBool "MAGMA_ENABLE_CUDA" cudaSupport)

    (lib.cmakeBool "MAGMA_ENABLE_HIP" rocmSupport)

  ] ++ lists.optionals static [

    "-DBUILD_SHARED_LIBS=OFF"

  ] ++ lists.optionals cudaSupport [
@@ -140,11 +139,9 @@ stdenv.mkDerivation { 
    "-DMIN_ARCH=${minArch}" # Disarms magma's asserts

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

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

    "-DMAGMA_ENABLE_CUDA=ON"

  ] ++ lists.optionals rocmSupport [

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

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

    "-DMAGMA_ENABLE_HIP=ON"

  ];



  buildFlags = [
@@ -155,7 +152,7 @@ stdenv.mkDerivation { 
  doCheck = false;



  passthru = {

    inherit cudaPackages cudaSupport;

    inherit cudaPackages cudaSupport rocmSupport gpuTargets;

  };



  meta = with lib; {
@@ -164,7 +161,11 @@ stdenv.mkDerivation { 
    homepage = "http://icl.cs.utk.edu/magma/index.html";

    platforms = platforms.unix;

    maintainers = with maintainers; [ connorbaker ];

    # CUDA and ROCm are mutually exclusive

    broken = cudaSupport && rocmSupport || cudaSupport && strings.versionOlder cudaVersion "9";



    # Cf. https://bitbucket.org/icl/magma/src/fcfe5aa61c1a4c664b36a73ebabbdbab82765e9f/CMakeLists.txt#lines-20

    broken =

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

      || (cudaSupport && rocmSupport) # Mutually exclusive

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

  };

}

pkgs/development/python-modules/torch/default.nix

+11 −4
Original line number Diff line number Diff line 
{ stdenv, lib, fetchFromGitHub, fetchpatch, buildPythonPackage, python,

  config, cudaSupport ? config.cudaSupport, cudaPackages, magma,

  config, cudaSupport ? config.cudaSupport, cudaPackages,

  effectiveMagma ?

  if cudaSupport then magma-cuda-static

  else if rocmSupport then magma-hip

  else magma,

  magma,

  magma-hip,

  magma-cuda-static,

  useSystemNccl ? true,

  MPISupport ? false, mpi,

  buildDocs ? false,
@@ -111,11 +118,11 @@ let 
  };



  brokenConditions = attrsets.filterAttrs (_: cond: cond) {

    "CUDA and ROCm are not mutually exclusive" = cudaSupport && rocmSupport;

    "CUDA and ROCm are mutually exclusive" = cudaSupport && rocmSupport;

    "CUDA is not targeting Linux" = cudaSupport && !stdenv.isLinux;

    "Unsupported CUDA version" = cudaSupport && !(builtins.elem cudaPackages.cudaMajorVersion [ "11" "12" ]);

    "MPI cudatoolkit does not match cudaPackages.cudatoolkit" = MPISupport && cudaSupport && (mpi.cudatoolkit != cudaPackages.cudatoolkit);

    "Magma cudaPackages does not match cudaPackages" = cudaSupport && (magma.cudaPackages != cudaPackages);

    "Magma cudaPackages does not match cudaPackages" = cudaSupport && (effectiveMagma.cudaPackages != cudaPackages);

  };

in buildPythonPackage rec {

  pname = "torch";
@@ -359,7 +366,7 @@ in buildPythonPackage rec { 
      cuda_profiler_api.dev # <cuda_profiler_api.h>

    ])

    ++ lib.optionals rocmSupport [ rocmPackages.llvm.openmp ]

    ++ lib.optionals (cudaSupport || rocmSupport) [ magma ]

    ++ lib.optionals (cudaSupport || rocmSupport) [ effectiveMagma ]

    ++ lib.optionals stdenv.isLinux [ numactl ]

    ++ lib.optionals stdenv.isDarwin [ Accelerate CoreServices libobjc ];

pkgs/top-level/python-packages.nix

+0 −6
Original line number Diff line number Diff line
@@ -14125,10 +14125,6 @@ self: super: with self; { 
  toposort = callPackage ../development/python-modules/toposort { };















  torch = callPackage ../development/python-modules/torch {













    magma =













      if pkgs.config.cudaSupport













      then pkgs.magma-cuda-static













      else pkgs.magma;















    inherit (pkgs.darwin.apple_sdk.frameworks) Accelerate CoreServices;















    inherit (pkgs.darwin) libobjc;















  };









@@ -14138,7 +14134,6 @@ self: super: with self; {













  };





























  torchWithCuda = self.torch.override {













    magma = pkgs.magma-cuda-static;















    openai-triton = self.openai-triton-cuda;















    cudaSupport = true;















    rocmSupport = false;









@@ -14149,7 +14144,6 @@ self: super: with self; {













  };





























  torchWithRocm = self.torch.override {













    magma = pkgs.magma-hip;















    openai-triton = self.openai-triton-no-cuda;















    rocmSupport = true;















    cudaSupport = false;