python3Packages.torch: enable vulkan support (#350091)

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

@@ -95,6 +95,11 @@
   rocmSupport ? config.rocmSupport,
   rocmPackages_5,
   gpuTargets ? [ ],
+
+  vulkanSupport ? false,
+  vulkan-headers,
+  vulkan-loader,
+  shaderc,
 }:
 
 let
@@ -447,58 +452,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 ‘<unknown>’ 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 ‘<unknown>’ 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 =
     [
@@ -563,31 +576,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 ];

pkgs/top-level/python-packages.nix

@@ -16212,6 +16212,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 { };