From a92626dfa5d8c1bdc21c87b0d5c1c1bcdd0e82f5 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 18 Jun 2024 21:40:37 +0000
Subject: [PATCH 1/2] python311Packages.openai-triton: move postPatch up to src

---
 .../python-modules/openai-triton/default.nix  | 61 ++++++++++---------
 1 file changed, 31 insertions(+), 30 deletions(-)

diff --git a/pkgs/development/python-modules/openai-triton/default.nix b/pkgs/development/python-modules/openai-triton/default.nix
index 1b4d713311ee..6869704a1bdb 100644
--- a/pkgs/development/python-modules/openai-triton/default.nix
+++ b/pkgs/development/python-modules/openai-triton/default.nix
@@ -57,36 +57,6 @@ buildPythonPackage rec {
       ./0001-ptxas-disable-version-key-for-non-cuda-targets.patch
     ];
 
-  nativeBuildInputs = [
-    setuptools
-    pythonRelaxDepsHook
-    # pytestCheckHook # Requires torch (circular dependency) and probably needs GPUs:
-    cmake
-    ninja
-
-    # Note for future:
-    # These *probably* should go in depsTargetTarget
-    # ...but we cannot test cross right now anyway
-    # because we only support cudaPackages on x86_64-linux atm
-    lit
-    llvm
-  ];
-
-  buildInputs = [
-    gtest
-    libxml2.dev
-    ncurses
-    pybind11
-    zlib
-  ];
-
-  propagatedBuildInputs = [
-    filelock
-    # openai-triton uses setuptools at runtime:
-    # https://github.com/NixOS/nixpkgs/pull/286763/#discussion_r1480392652
-    setuptools
-  ];
-
   postPatch =
     let
       # Bash was getting weird without linting,
@@ -123,6 +93,37 @@ buildPythonPackage rec {
         --replace '${oldStr}' '${newStr}'
     '';
 
+  nativeBuildInputs = [
+    setuptools
+    pythonRelaxDepsHook
+    # pytestCheckHook # Requires torch (circular dependency) and probably needs GPUs:
+    cmake
+    ninja
+
+    # Note for future:
+    # These *probably* should go in depsTargetTarget
+    # ...but we cannot test cross right now anyway
+    # because we only support cudaPackages on x86_64-linux atm
+    lit
+    llvm
+  ];
+
+  buildInputs = [
+    gtest
+    libxml2.dev
+    ncurses
+    pybind11
+    zlib
+  ];
+
+  propagatedBuildInputs = [
+    filelock
+    # openai-triton uses setuptools at runtime:
+    # https://github.com/NixOS/nixpkgs/pull/286763/#discussion_r1480392652
+    setuptools
+  ];
+
+
   # Avoid GLIBCXX mismatch with other cuda-enabled python packages
   preConfigure =
     ''

From 530d5769d57e88bde975c5fefb9ecfcddba89742 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 18 Jun 2024 22:03:08 +0000
Subject: [PATCH 2/2] python311Packages.torch: prevent libcuda_dirs() from
 running ldconfig

---
 .../python-modules/openai-triton/default.nix  | 39 ++++++++++++-------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/pkgs/development/python-modules/openai-triton/default.nix b/pkgs/development/python-modules/openai-triton/default.nix
index 6869704a1bdb..d1aaed52052d 100644
--- a/pkgs/development/python-modules/openai-triton/default.nix
+++ b/pkgs/development/python-modules/openai-triton/default.nix
@@ -1,10 +1,10 @@
 {
   lib,
   config,
+  addDriverRunpath,
   buildPythonPackage,
   fetchFromGitHub,
   fetchpatch,
-  addOpenGLRunpath,
   setuptools,
   pytestCheckHook,
   pythonRelaxDepsHook,
@@ -59,19 +59,23 @@ buildPythonPackage rec {
 
   postPatch =
     let
-      # Bash was getting weird without linting,
-      # but basically upstream contains [cc, ..., "-lcuda", ...]
-      # and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]
-      old = [ "-lcuda" ];
-      new = [
-        "-lcuda"
-        "-L${addOpenGLRunpath.driverLink}"
-        "-L${cudaPackages.cuda_cudart}/lib/stubs/"
-      ];
-
       quote = x: ''"${x}"'';
-      oldStr = lib.concatMapStringsSep ", " quote old;
-      newStr = lib.concatMapStringsSep ", " quote new;
+      subs.ldFlags =
+        let
+          # Bash was getting weird without linting,
+          # but basically upstream contains [cc, ..., "-lcuda", ...]
+          # and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]
+          old = [ "-lcuda" ];
+          new = [
+            "-lcuda"
+            "-L${addDriverRunpath.driverLink}"
+            "-L${cudaPackages.cuda_cudart}/lib/stubs/"
+          ];
+        in
+        {
+          oldStr = lib.concatMapStringsSep ", " quote old;
+          newStr = lib.concatMapStringsSep ", " quote new;
+        };
     in
     ''
       # Use our `cmakeFlags` instead and avoid downloading dependencies
@@ -86,11 +90,17 @@ buildPythonPackage rec {
       substituteInPlace unittest/CMakeLists.txt \
         --replace "include (\''${CMAKE_CURRENT_SOURCE_DIR}/googletest.cmake)" ""\
         --replace "include(GoogleTest)" "find_package(GTest REQUIRED)"
+
+      cat << \EOF > python/triton/common/build.py
+
+      def libcuda_dirs():
+          return [ "${addDriverRunpath.driverLink}/lib" ]
+      EOF
     ''
     + lib.optionalString cudaSupport ''
       # Use our linker flags
       substituteInPlace python/triton/common/build.py \
-        --replace '${oldStr}' '${newStr}'
+        --replace '${subs.ldFlags.oldStr}' '${subs.ldFlags.newStr}'
     '';
 
   nativeBuildInputs = [
@@ -123,7 +133,6 @@ buildPythonPackage rec {
     setuptools
   ];
 
-
   # Avoid GLIBCXX mismatch with other cuda-enabled python packages
   preConfigure =
     ''