b12f/nixpkgs
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

Merge pull request #113207 from danieldk/pytorch-cuda-11

Browse source 
python3Packages.pytorch: add compute capabilities for CUDA 11
This commit is contained in:
Daniël de Kok 2021-03-05 08:01:50 +01:00 committed by GitHub
parent 2c0e6a0e69 6038b56de8
commit e9b3e36f44
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 26 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
pkgs/development/python-modules/pytorch/default.nix
View file

@ -74,27 +74,35 @@ let
# (allowing FBGEMM to be built in pytorch-1.1), and may future proof this # (allowing FBGEMM to be built in pytorch-1.1), and may future proof this
# derivation. # derivation.
brokenArchs = [ "3.0" ]; # this variable is only used as documentation. brokenArchs = [ "3.0" ]; # this variable is only used as documentation.
cuda9ArchList = [
"3.5" cudaCapabilities = rec {
"5.0" cuda9 = [
"5.2" "3.5"
"6.0" "5.0"
"6.1" "5.2"
"7.0" "6.0"
"7.0+PTX" # I am getting a "undefined architecture compute_75" on cuda 9 "6.1"
# which leads me to believe this is the final cuda-9-compatible architecture. "7.0"
]; "7.0+PTX" # I am getting a "undefined architecture compute_75" on cuda 9
cuda10ArchList = cuda9ArchList ++ [ # which leads me to believe this is the final cuda-9-compatible architecture.
"7.5" ];
"7.5+PTX" # < most recent architecture as of cudatoolkit_10_0 and pytorch-1.2.0
]; cuda10 = cuda9 ++ [
"7.5"
"7.5+PTX" # < most recent architecture as of cudatoolkit_10_0 and pytorch-1.2.0
];
cuda11 = cuda10 ++ [
"8.0"
"8.0+PTX" # < CUDA toolkit 11.0
"8.6"
"8.6+PTX" # < CUDA toolkit 11.1
];
};
final_cudaArchList = final_cudaArchList =
if !cudaSupport || cudaArchList != null if !cudaSupport || cudaArchList != null
then cudaArchList then cudaArchList
else else cudaCapabilities."cuda${lib.versions.major cudatoolkit.version}";
if lib.versions.major cudatoolkit.version == "9"
then cuda9ArchList
else cuda10ArchList; # the assert above removes any ambiguity here.
# Normally libcuda.so.1 is provided at runtime by nvidia-x11 via # Normally libcuda.so.1 is provided at runtime by nvidia-x11 via
# LD_LIBRARY_PATH=/run/opengl-driver/lib. We only use the stub # LD_LIBRARY_PATH=/run/opengl-driver/lib. We only use the stub