commit 8eae4972240f69c834bf56d5a5c9b39adb1c02c4 · pyrox.dev/nixpkgs

+6 -22

pkgs/development/python-modules/triton/0001-_build-allow-extra-cc-flags.patch

···

       1
       1
       -
       From 2751c5de5c61c90b56e3e392a41847f4c47258fd Mon Sep 17 00:00:00 2001

     

       2
       2
       -
       From: SomeoneSerge <else+aalto@someonex.net>

     

       3
       3
       -
       Date: Sun, 13 Oct 2024 14:16:48 +0000

     

       4
       4
       -
       Subject: [PATCH 1/3] _build: allow extra cc flags

     

       5
       5
       -
       

     

       6
       6
       -
       ---

     

       7
       7
       -
        python/triton/runtime/build.py | 10 +++++++++-

     

       8
       8
       -
        1 file changed, 9 insertions(+), 1 deletion(-)

     

       9
       9
       -
       

     

       10
       1
        
       diff --git a/python/triton/runtime/build.py b/python/triton/runtime/build.py

     

       11
       11
       -
       index d7baeb286..d334dce77 100644

     

       2
       2
       +
       index 1b76548d4..1edbfd3da 100644

     

       12
       3
        
       --- a/python/triton/runtime/build.py

     

       13
       4
        
       +++ b/python/triton/runtime/build.py

     

       14
       14
       -
       @@ -42,9 +42,17 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):

     

       15
       15
       -
            py_include_dir = sysconfig.get_paths(scheme=scheme)["include"]

     

       16
       16
       -
            include_dirs = include_dirs + [srcdir, py_include_dir]

     

       17
       17
       -
            cc_cmd = [cc, src, "-O3", "-shared", "-fPIC", "-o", so]

     

       5
       5
       +
       @@ -30,6 +30,14 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):

     

       6
       6
       +
            include_dirs = include_dirs + [srcdir, py_include_dir, *custom_backend_dirs]

     

       7
       7
       +
            # for -Wno-psabi, see https://gcc.gnu.org/bugzilla/show_bug.cgi?id=111047

     

       8
       8
       +
            cc_cmd = [cc, src, "-O3", "-shared", "-fPIC", "-Wno-psabi", "-o", so]

     

       18
       9
        
       +

     

       19
       10
        
       +    # Nixpkgs support branch

     

       20
       11
        
       +    # Allows passing e.g. extra -Wl,-rpath

     
···

       25
       16
        
       +

     

       26
       17
        
            cc_cmd += [f'-l{lib}' for lib in libraries]

     

       27
       18
        
            cc_cmd += [f"-L{dir}" for dir in library_dirs]

     

       28
       28
       -
       -    cc_cmd += [f"-I{dir}" for dir in include_dirs]

     

       29
       29
       -
       +    cc_cmd += [f"-I{dir}" for dir in include_dirs if dir is not None]

     

       30
       30
       -
            ret = subprocess.check_call(cc_cmd)

     

       31
       31
       -
            if ret == 0:

     

       32
       32
       -
                return so

     

       33
       33
       -
       -- 

     

       34
       34
       -
       2.46.0

     

       35
       35
       -
       

     

       19
       19
       +
            cc_cmd += [f"-I{dir}" for dir in include_dirs if dir is not None]

+14 -29

pkgs/development/python-modules/triton/0002-nvidia-amd-driver-short-circuit-before-ldconfig.patch

···

       1
       1
       -
       From 7407cb03eec82768e333909d87b7668b633bfe86 Mon Sep 17 00:00:00 2001

     

       2
       2
       -
       From: SomeoneSerge <else+aalto@someonex.net>

     

       3
       3
       -
       Date: Sun, 13 Oct 2024 14:28:48 +0000

     

       4
       4
       -
       Subject: [PATCH 2/3] {nvidia,amd}/driver: short-circuit before ldconfig

     

       5
       5
       -
       

     

       6
       6
       -
       ---

     

       7
       7
       -
        python/triton/runtime/build.py       | 6 +++---

     

       8
       8
       -
        third_party/amd/backend/driver.py    | 7 +++++++

     

       9
       9
       -
        third_party/nvidia/backend/driver.py | 3 +++

     

       10
       10
       -
        3 files changed, 13 insertions(+), 3 deletions(-)

     

       11
       11
       -
       

     

       12
       1
        
       diff --git a/python/triton/runtime/build.py b/python/triton/runtime/build.py

     

       13
       13
       -
       index d334dce77..a64e98da0 100644

     

       2
       2
       +
       index 1edbfd3da..2756dccdb 100644

     

       14
       3
        
       --- a/python/triton/runtime/build.py

     

       15
       4
        
       +++ b/python/triton/runtime/build.py

     

       16
       16
       -
       @@ -42,6 +42,9 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):

     

       17
       17
       -
            py_include_dir = sysconfig.get_paths(scheme=scheme)["include"]

     

       18
       18
       -
            include_dirs = include_dirs + [srcdir, py_include_dir]

     

       19
       19
       -
            cc_cmd = [cc, src, "-O3", "-shared", "-fPIC", "-o", so]

     

       5
       5
       +
       @@ -30,6 +30,9 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):

     

       6
       6
       +
            include_dirs = include_dirs + [srcdir, py_include_dir, *custom_backend_dirs]

     

       7
       7
       +
            # for -Wno-psabi, see https://gcc.gnu.org/bugzilla/show_bug.cgi?id=111047

     

       8
       8
       +
            cc_cmd = [cc, src, "-O3", "-shared", "-fPIC", "-Wno-psabi", "-o", so]

     

       20
       9
        
       +    cc_cmd += [f'-l{lib}' for lib in libraries]

     

       21
       10
        
       +    cc_cmd += [f"-L{dir}" for dir in library_dirs]

     

       22
       11
        
       +    cc_cmd += [f"-I{dir}" for dir in include_dirs if dir is not None]

     

       23
       12
        
        

     

       24
       13
        
            # Nixpkgs support branch

     

       25
       14
        
            # Allows passing e.g. extra -Wl,-rpath

     

       26
       26
       -
       @@ -50,9 +53,6 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):

     

       15
       15
       +
       @@ -38,8 +41,5 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries):

     

       27
       16
        
                import shlex

     

       28
       17
        
                cc_cmd.extend(shlex.split(cc_cmd_extra_flags))

     

       29
       18
        
        

     

       30
       19
        
       -    cc_cmd += [f'-l{lib}' for lib in libraries]

     

       31
       20
        
       -    cc_cmd += [f"-L{dir}" for dir in library_dirs]

     

       32
       21
        
       -    cc_cmd += [f"-I{dir}" for dir in include_dirs if dir is not None]

     

       33
       33
       -
            ret = subprocess.check_call(cc_cmd)

     

       34
       34
       -
            if ret == 0:

     

       35
       35
       -
                return so

     

       22
       22
       +
            subprocess.check_call(cc_cmd, stdout=subprocess.DEVNULL)

     

       23
       23
       +
            return so

     

       36
       24
        
       diff --git a/third_party/amd/backend/driver.py b/third_party/amd/backend/driver.py

     

       37
       37
       -
       index 0a8cd7bed..aab8805f6 100644

     

       25
       25
       +
       index b99ff86c8..ea8bc103d 100644

     

       38
       26
        
       --- a/third_party/amd/backend/driver.py

     

       39
       27
        
       +++ b/third_party/amd/backend/driver.py

     

       40
       40
       -
       @@ -24,6 +24,13 @@ def _get_path_to_hip_runtime_dylib():

     

       41
       41
       -
                    return env_libhip_path

     

       42
       42
       -
                raise RuntimeError(f"TRITON_LIBHIP_PATH '{env_libhip_path}' does not point to a valid {lib_name}")

     

       28
       28
       +
       @@ -79,6 +79,13 @@ def _get_path_to_hip_runtime_dylib():

     

       29
       29
       +
                    return mmapped_path

     

       30
       30
       +
                raise RuntimeError(f"memory mapped '{mmapped_path}' in process does not point to a valid {lib_name}")

     

       43
       31
        
        

     

       44
       32
        
       +    # ...on release/3.1.x:

     

       45
       33
        
       +    #         return mmapped_path

     
···

       52
       40
        
        

     

       53
       41
        
            import site

     

       54
       42
        
       diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py

     

       55
       55
       -
       index 90f71138b..30fbadb2a 100644

     

       43
       43
       +
       index 5f2621ae5..e7762a3ec 100644

     

       56
       44
        
       --- a/third_party/nvidia/backend/driver.py

     

       57
       45
        
       +++ b/third_party/nvidia/backend/driver.py

     

       58
       58
       -
       @@ -21,6 +21,9 @@ def libcuda_dirs():

     

       46
       46
       +
       @@ -23,6 +23,9 @@ def libcuda_dirs():

     

       59
       47
        
            if env_libcuda_path:

     

       60
       48
        
                return [env_libcuda_path]

     

       61
       49
        
        

     
···

       65
       53
        
            libs = subprocess.check_output(["/sbin/ldconfig", "-p"]).decode()

     

       66
       54
        
            # each line looks like the following:

     

       67
       55
        
            # libcuda.so.1 (libc6,x86-64) => /lib/x86_64-linux-gnu/libcuda.so.1

     

       68
       68
       -
       -- 

     

       69
       69
       -
       2.46.0

     

       70
       70
       -

+19 -26

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

···

       16
       16
        
         pybind11,

     

       17
       17
        
         python,

     

       18
       18
        
         pytestCheckHook,

     

       19
       19
       +
         writableTmpDirAsHomeHook,

     

       19
       20
        
         stdenv,

     

       20
       21
        
         replaceVars,

     

       21
       22
        
         setuptools,

     
···

       27
       28
        
         triton,

     

       28
       29
        
       }:

     

       29
       30
        
       

     

       30
       30
       -
       buildPythonPackage {

     

       31
       31
       +
       buildPythonPackage rec {

     

       31
       32
        
         pname = "triton";

     

       32
       32
       -
         version = "3.1.0";

     

       33
       33
       +
         version = "3.3.1";

     

       33
       34
        
         pyproject = true;

     

       34
       35
        
       

     

       35
       36
        
         src = fetchFromGitHub {

     

       36
       37
        
           owner = "triton-lang";

     

       37
       38
        
           repo = "triton";

     

       38
       38
       -
           # latest branch commit from https://github.com/triton-lang/triton/commits/release/3.1.x/

     

       39
       39
       -
           rev = "cf34004b8a67d290a962da166f5aa2fc66751326";

     

       40
       40
       -
           hash = "sha256-233fpuR7XXOaSKN+slhJbE/CMFzAqCRCE4V4rIoJZrk=";

     

       39
       39
       +
           tag = "v${version}";

     

       40
       40
       +
           hash = "sha256-XLw7s5K0j4mfIvNMumlHkUpklSzVSTRyfGazZ4lLpn0=";

     

       41
       41
        
         };

     

       42
       42
        
       

     

       43
       43
        
         patches =

     

       44
       44
        
           [

     

       45
       45
       -
             ./0001-setup.py-introduce-TRITON_OFFLINE_BUILD.patch

     

       45
       45
       +
             # ./0001-setup.py-introduce-TRITON_OFFLINE_BUILD.patch

     

       46
       46
        
             (replaceVars ./0001-_build-allow-extra-cc-flags.patch {

     

       47
       47
        
               ccCmdExtraFlags = "-Wl,-rpath,${addDriverRunpath.driverLink}/lib";

     

       48
       48
        
             })

     
···

       65
       65
        
           # Use our `cmakeFlags` instead and avoid downloading dependencies

     

       66
       66
        
           # remove any downloads

     

       67
       67
        
           substituteInPlace python/setup.py \

     

       68
       68
       -
             --replace-fail "get_json_package_info(), get_pybind11_package_info()" ""\

     

       69
       69
       -
             --replace-fail "get_pybind11_package_info(), get_llvm_package_info()" ""\

     

       70
       70
       -
             --replace-fail 'packages += ["triton/profiler"]' ""\

     

       71
       71
       -
             --replace-fail "curr_version != version" "False"

     

       68
       68
       +
             --replace-fail "[get_json_package_info()]" "[]"\

     

       69
       69
       +
             --replace-fail "[get_llvm_package_info()]" "[]"\

     

       70
       70
       +
             --replace-fail 'packages += ["triton/profiler"]' "pass"\

     

       71
       71
       +
             --replace-fail "curr_version.group(1) != version" "False"

     

       72
       72
        
       

     

       73
       73
        
           # Don't fetch googletest

     

       74
       74
       -
           substituteInPlace unittest/CMakeLists.txt \

     

       75
       75
       -
             --replace-fail "include (\''${CMAKE_CURRENT_SOURCE_DIR}/googletest.cmake)" ""\

     

       74
       74
       +
           substituteInPlace cmake/AddTritonUnitTest.cmake \

     

       75
       75
       +
             --replace-fail "include(\''${PROJECT_SOURCE_DIR}/unittest/googletest.cmake)" ""\

     

       76
       76
        
             --replace-fail "include(GoogleTest)" "find_package(GTest REQUIRED)"

     

       77
       77
       -
       

     

       78
       78
       -
           # Patch the source code to make sure it doesn't specify a non-existent PTXAS version.

     

       79
       79
       -
           # CUDA 12.6 (the current default/max) tops out at PTXAS version 8.5.

     

       80
       80
       -
           # NOTE: This is fixed in `master`:

     

       81
       81
       -
           # https://github.com/triton-lang/triton/commit/f48dbc1b106c93144c198fbf3c4f30b2aab9d242

     

       82
       82
       -
           substituteInPlace "$NIX_BUILD_TOP/$sourceRoot/third_party/nvidia/backend/compiler.py" \

     

       83
       83
       -
             --replace-fail \

     

       84
       84
       -
               'return 80 + minor' \

     

       85
       85
       -
               'return 80 + min(minor, 5)'

     

       86
       77
        
         '';

     

       87
       78
        
       

     

       88
       79
        
         build-system = [ setuptools ];

     
···

       97
       88
        
           # because we only support cudaPackages on x86_64-linux atm

     

       98
       89
        
           lit

     

       99
       90
        
           llvm

     

       91
       91
       +
       

     

       92
       92
       +
           # Upstream's setup.py tries to write cache somewhere in ~/

     

       93
       93
       +
           writableTmpDirAsHomeHook

     

       100
       94
        
         ];

     

       101
       95
        
       

     

       102
       96
        
         buildInputs = [

     
···

       124
       118
        
         preConfigure = ''

     

       125
       119
        
           # Ensure that the build process uses the requested number of cores

     

       126
       120
        
           export MAX_JOBS="$NIX_BUILD_CORES"

     

       127
       127
       -
       

     

       128
       128
       -
           # Upstream's setup.py tries to write cache somewhere in ~/

     

       129
       129
       -
           export HOME=$(mktemp -d)

     

       130
       121
        
       

     

       131
       122
        
           # Upstream's github actions patch setup.cfg to write base-dir. May be redundant

     

       132
       123
        
           echo "

     
···

       193
       184
        
           ];

     

       194
       185
        
       

     

       195
       186
        
           dontBuild = true;

     

       196
       196
       -
           nativeCheckInputs = [ pytestCheckHook ];

     

       187
       187
       +
           nativeCheckInputs = [

     

       188
       188
       +
             pytestCheckHook

     

       189
       189
       +
             writableTmpDirAsHomeHook

     

       190
       190
       +
           ];

     

       197
       191
        
       

     

       198
       192
        
           doCheck = true;

     

       199
       193
        
       

     

       200
       194
        
           preCheck = ''

     

       201
       195
        
             cd python/test/unit

     

       202
       202
       -
             export HOME=$TMPDIR

     

       203
       196
        
           '';

     

       204
       197
        
           checkPhase = "pytestCheckPhase";

     

       205
       198