dev/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Merge pull request #266081 from ConnorBaker/fix/torch-jetson

Browse Source 
python3Packages.torch: patch cpp_extension.py for Jetson support
This commit is contained in:
Connor Baker 2023-11-09 10:28:19 -05:00 committed by GitHub
commit 47f07caec9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
pkgs/development/python-modules/torch/default.nix
View File

@ -48,7 +48,10 @@
let
inherit (lib) attrsets lists strings trivial;
inherit (cudaPackages) cudaFlags cudnn nccl;
inherit (cudaPackages) cudaFlags cudnn;
# Some packages are not available on all platforms
nccl = cudaPackages.nccl or null;
setBool = v: if v then "1" else "0";
@ -178,6 +181,13 @@ in buildPythonPackage rec {
'message(FATAL_ERROR "Found NCCL header version and library version' \
'message(WARNING "Found NCCL header version and library version'
''
# TODO(@connorbaker): Remove this patch after 2.1.0 lands.
+ lib.optionalString cudaSupport ''
substituteInPlace torch/utils/cpp_extension.py \
--replace \
"'8.6', '8.9'" \
"'8.6', '8.7', '8.9'"
''
# error: no member named 'aligned_alloc' in the global namespace; did you mean simply 'aligned_alloc'
# This lib overrided aligned_alloc hence the error message. Tltr: his function is linkable but not in header.
+ lib.optionalString (stdenv.isDarwin && lib.versionOlder stdenv.targetPlatform.darwinSdkVersion "11.0") ''
@ -253,6 +263,7 @@ in buildPythonPackage rec {
PYTORCH_BUILD_VERSION = version;
PYTORCH_BUILD_NUMBER = 0;
USE_NCCL = setBool (nccl != null);
USE_SYSTEM_NCCL = setBool useSystemNccl; # don't build pytorch's third_party NCCL
USE_STATIC_NCCL = setBool useSystemNccl;
@ -316,6 +327,8 @@ in buildPythonPackage rec {
libcusolver.lib
libcusparse.dev
libcusparse.lib
] ++ lists.optionals (nccl != null) [
# Some platforms do not support NCCL (i.e., Jetson)
nccl.dev # Provides nccl.h AND a static copy of NCCL!
] ++ lists.optionals (strings.versionOlder cudaVersion "11.8") [
cuda_nvprof.dev # <cuda_profiler_api.h>