From 5cd68d33bbcd73c7e2d852117a470addd2e8cf14 Mon Sep 17 00:00:00 2001 From: Isidor Zeuner Date: Mon, 21 Oct 2024 09:40:16 +0200 Subject: [PATCH] python3Packages.torch: enable vulkan support --- .../python-modules/torch/default.nix | 167 ++++++++++-------- pkgs/top-level/python-packages.nix | 4 + 2 files changed, 97 insertions(+), 74 deletions(-) diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index cd43c04d67d4..c14bcf3d37ea 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -91,6 +91,11 @@ rocmSupport ? config.rocmSupport, rocmPackages_5, gpuTargets ? [ ], + + vulkanSupport ? false, + vulkan-headers, + vulkan-loader, + shaderc, }: let @@ -412,58 +417,66 @@ buildPythonPackage rec { # https://github.com/pytorch/pytorch/commit/3d617333e PYTHON_LIB_REL_PATH = "${placeholder "out"}/${python.sitePackages}"; - # Suppress a weird warning in mkl-dnn, part of ideep in pytorch - # (upstream seems to have fixed this in the wrong place?) - # https://github.com/intel/mkl-dnn/commit/8134d346cdb7fe1695a2aa55771071d455fae0bc - # https://github.com/pytorch/pytorch/issues/22346 - # - # Also of interest: pytorch ignores CXXFLAGS uses CFLAGS for both C and C++: - # https://github.com/pytorch/pytorch/blob/v1.11.0/setup.py#L17 - env.NIX_CFLAGS_COMPILE = toString ( - ( - lib.optionals (blas.implementation == "mkl") [ "-Wno-error=array-bounds" ] - # Suppress gcc regression: avx512 math function raises uninitialized variable warning - # https://gcc.gnu.org/bugzilla/show_bug.cgi?id=105593 - # See also: Fails to compile with GCC 12.1.0 https://github.com/pytorch/pytorch/issues/77939 - ++ lib.optionals (stdenv.cc.isGNU && lib.versionAtLeast stdenv.cc.version "12.0.0") [ - "-Wno-error=maybe-uninitialized" - "-Wno-error=uninitialized" - ] - # Since pytorch 2.0: - # gcc-12.2.0/include/c++/12.2.0/bits/new_allocator.h:158:33: error: ‘void operator delete(void*, std::size_t)’ - # ... called on pointer ‘’ with nonzero offset [1, 9223372036854775800] [-Werror=free-nonheap-object] - ++ lib.optionals (stdenv.cc.isGNU && lib.versions.major stdenv.cc.version == "12") [ - "-Wno-error=free-nonheap-object" - ] - # .../source/torch/csrc/autograd/generated/python_functions_0.cpp:85:3: - # error: cast from ... to ... converts to incompatible function type [-Werror,-Wcast-function-type-strict] - ++ lib.optionals (stdenv.cc.isClang && lib.versionAtLeast stdenv.cc.version "16") [ - "-Wno-error=cast-function-type-strict" - # Suppresses the most spammy warnings. - # This is mainly to fix https://github.com/NixOS/nixpkgs/issues/266895. - ] - ++ lib.optionals rocmSupport [ - "-Wno-#warnings" - "-Wno-cpp" - "-Wno-unknown-warning-option" - "-Wno-ignored-attributes" - "-Wno-deprecated-declarations" - "-Wno-defaulted-function-deleted" - "-Wno-pass-failed" - ] - ++ [ - "-Wno-unused-command-line-argument" - "-Wno-uninitialized" - "-Wno-array-bounds" - "-Wno-free-nonheap-object" - "-Wno-unused-result" - ] - ++ lib.optionals stdenv.cc.isGNU [ - "-Wno-maybe-uninitialized" - "-Wno-stringop-overflow" - ] - ) - ); + env = + { + # Suppress a weird warning in mkl-dnn, part of ideep in pytorch + # (upstream seems to have fixed this in the wrong place?) + # https://github.com/intel/mkl-dnn/commit/8134d346cdb7fe1695a2aa55771071d455fae0bc + # https://github.com/pytorch/pytorch/issues/22346 + # + # Also of interest: pytorch ignores CXXFLAGS uses CFLAGS for both C and C++: + # https://github.com/pytorch/pytorch/blob/v1.11.0/setup.py#L17 + NIX_CFLAGS_COMPILE = toString ( + ( + lib.optionals (blas.implementation == "mkl") [ "-Wno-error=array-bounds" ] + # Suppress gcc regression: avx512 math function raises uninitialized variable warning + # https://gcc.gnu.org/bugzilla/show_bug.cgi?id=105593 + # See also: Fails to compile with GCC 12.1.0 https://github.com/pytorch/pytorch/issues/77939 + ++ lib.optionals (stdenv.cc.isGNU && lib.versionAtLeast stdenv.cc.version "12.0.0") [ + "-Wno-error=maybe-uninitialized" + "-Wno-error=uninitialized" + ] + # Since pytorch 2.0: + # gcc-12.2.0/include/c++/12.2.0/bits/new_allocator.h:158:33: error: ‘void operator delete(void*, std::size_t)’ + # ... called on pointer ‘’ with nonzero offset [1, 9223372036854775800] [-Werror=free-nonheap-object] + ++ lib.optionals (stdenv.cc.isGNU && lib.versions.major stdenv.cc.version == "12") [ + "-Wno-error=free-nonheap-object" + ] + # .../source/torch/csrc/autograd/generated/python_functions_0.cpp:85:3: + # error: cast from ... to ... converts to incompatible function type [-Werror,-Wcast-function-type-strict] + ++ lib.optionals (stdenv.cc.isClang && lib.versionAtLeast stdenv.cc.version "16") [ + "-Wno-error=cast-function-type-strict" + # Suppresses the most spammy warnings. + # This is mainly to fix https://github.com/NixOS/nixpkgs/issues/266895. + ] + ++ lib.optionals rocmSupport [ + "-Wno-#warnings" + "-Wno-cpp" + "-Wno-unknown-warning-option" + "-Wno-ignored-attributes" + "-Wno-deprecated-declarations" + "-Wno-defaulted-function-deleted" + "-Wno-pass-failed" + ] + ++ [ + "-Wno-unused-command-line-argument" + "-Wno-uninitialized" + "-Wno-array-bounds" + "-Wno-free-nonheap-object" + "-Wno-unused-result" + ] + ++ lib.optionals stdenv.cc.isGNU [ + "-Wno-maybe-uninitialized" + "-Wno-stringop-overflow" + ] + ) + ); + + USE_VULKAN = setBool vulkanSupport; + } + // lib.optionalAttrs vulkanSupport { + VULKAN_SDK = shaderc.bin; + }; nativeBuildInputs = [ @@ -530,31 +543,37 @@ buildPythonPackage rec { pythonRelaxDeps = [ "sympy" ]; - dependencies = [ - astunparse - cffi - click - numpy - pyyaml + dependencies = + [ + astunparse + cffi + click + numpy + pyyaml - # From install_requires: - fsspec - filelock - typing-extensions - sympy - networkx - jinja2 + # From install_requires: + fsspec + filelock + typing-extensions + sympy + networkx + jinja2 - # the following are required for tensorboard support - pillow - six - future - tensorboard - protobuf + # the following are required for tensorboard support + pillow + six + future + tensorboard + protobuf - # torch/csrc requires `pybind11` at runtime - pybind11 - ] ++ lib.optionals tritonSupport [ _tritonEffective ]; + # torch/csrc requires `pybind11` at runtime + pybind11 + ] + ++ lib.optionals tritonSupport [ _tritonEffective ] + ++ lib.optionals vulkanSupport [ + vulkan-headers + vulkan-loader + ]; propagatedCxxBuildInputs = [ ] ++ lib.optionals MPISupport [ mpi ] ++ lib.optionals rocmSupport [ rocmtoolkit_joined ]; diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 9186cb823ea3..547f2e45ef97 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -15810,6 +15810,10 @@ self: super: with self; { rocmSupport = false; }; + torchWithVulkan = self.torch.override { + vulkanSupport = true; + }; + torch-tb-profiler = callPackage ../development/python-modules/torch-tb-profiler/default.nix { }; torchaudio = callPackage ../development/python-modules/torchaudio { };