Merge pull request #238465 from ConnorBaker/feat/torch-use-static-magma

python3Packages.torch: statically link against magma
This commit is contained in:
Samuel Ainsworth 2023-06-22 07:55:20 -07:00 committed by GitHub
commit 83805599a9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 12 additions and 1 deletions

View file

@ -28,6 +28,7 @@
, ninja
, openmp
, rocmSupport ? false
, static ? false
, stdenv
, symlinkJoin
}:
@ -145,6 +146,8 @@ stdenv.mkDerivation {
cmakeFlags = [
"-DGPU_TARGET=${gpuTargetString}"
] ++ lists.optionals static [
"-DBUILD_SHARED_LIBS=OFF"
] ++ lists.optionals cudaSupport [
"-DCMAKE_CUDA_ARCHITECTURES=${cudaArchitecturesString}"
"-DMIN_ARCH=${minArch}" # Disarms magma's asserts

View file

@ -38230,6 +38230,10 @@ with pkgs;
rocmSupport = false;
};
magma-cuda-static = magma-cuda.override {
static = true;
};
# TODO:AMD won't compile with anything newer than 2.6.2 -- it fails at the linking stage.
magma-hip = magma_2_6_2.override {
cudaSupport = false;

View file

@ -12317,6 +12317,10 @@ self: super: with self; {
torch = callPackage ../development/python-modules/torch {
cudaSupport = pkgs.config.cudaSupport or false;
magma =
if pkgs.config.cudaSupport or false
then pkgs.magma-cuda-static
else pkgs.magma;
inherit (pkgs.darwin.apple_sdk.frameworks) Accelerate CoreServices;
inherit (pkgs.darwin) libobjc;
inherit (pkgs.llvmPackages_rocm) openmp;
@ -12327,7 +12331,7 @@ self: super: with self; {
};
torchWithCuda = self.torch.override {
magma = pkgs.magma-cuda;
magma = pkgs.magma-cuda-static;
cudaSupport = true;
rocmSupport = false;
};