python3Packages.triton*: 2.1.0 -> 3.0.0 (#328247)

{ lib

, stdenv

, fetchFromGitHub

, fetchpatch

, pkgsBuildBuild

, pkg-config

, cmake

, libedit

, libffi

, libpfm

, lit

, mpfr

, zlib

, ncurses

  isNative = stdenv.hostPlatform == stdenv.buildPlatform;

in stdenv.mkDerivation (finalAttrs: {

  pname = "triton-llvm";

  version = "17.0.0-c5dede880d17";

  version = "19.1.0-rc1"; # One of the tags at https://github.com/llvm/llvm-project/commit/10dc3a8e916d73291269e5e2b82dd22681489aa1

  outputs = [

    "out"

  src = fetchFromGitHub {

    owner = "llvm";

    repo = "llvm-project";

    rev = "c5dede880d175f7229c9b2923f4753e12702305d";

    hash = "sha256-v4r3+7XVFK+Dzxt/rErZNJ9REqFO3JmGN4X4vZ+77ew=";

    rev = "10dc3a8e916d73291269e5e2b82dd22681489aa1";

    hash = "sha256-9DPvcFmhzw6MipQeCQnr35LktW0uxtEL8axMMPXIfWw=";

  };

  patches = [

    # glibc-2.40 support

    # [llvm-exegesis] Use correct rseq struct size #100804

    # https://github.com/llvm/llvm-project/issues/100791

    (fetchpatch {

      url = "https://github.com/llvm/llvm-project//commit/84837e3cc1cf17ed71580e3ea38299ed2bfaa5f6.patch";

      hash = "sha256-QKa+kyXjjGXwTQTEpmKZx5yYjOyBX8A8NQoIYUaGcIw=";

    })

  ];

  nativeBuildInputs = [

    pkg-config

    doxygen

    sphinx

    python3Packages.recommonmark

    python3Packages.myst-parser

  ];

  buildInputs = [

    ncurses

  ];

  sourceRoot = "${finalAttrs.src.name}/llvm";

  preConfigure = ''

    cd llvm

  '';

  cmakeFlags = [

    (lib.cmakeFeature "LLVM_TARGETS_TO_BUILD" (lib.concatStringsSep ";" llvmTargetsToBuild'))

  postPatch = ''

    # `CMake Error: cannot write to file "/build/source/llvm/build/lib/cmake/mlir/MLIRTargets.cmake": Permission denied`

    chmod +w -R ../mlir

    patchShebangs ../mlir/test/mlir-reduce

    chmod +w -R ./mlir

    patchShebangs ./mlir/test/mlir-reduce

    # FileSystem permissions tests fail with various special bits

    rm test/tools/llvm-objcopy/ELF/mirror-permissions-unix.test

    rm unittests/Support/Path.cpp

    rm llvm/test/tools/llvm-objcopy/ELF/mirror-permissions-unix.test

    rm llvm/unittests/Support/Path.cpp

    substituteInPlace unittests/Support/CMakeLists.txt \

    substituteInPlace llvm/unittests/Support/CMakeLists.txt \

      --replace "Path.cpp" ""

  '' + lib.optionalString stdenv.hostPlatform.isAarch64 ''

    # Not sure why this fails

    rm test/tools/llvm-exegesis/AArch64/latency-by-opcode-name.s

    rm llvm/test/tools/llvm-exegesis/AArch64/latency-by-opcode-name.s

  '';

  postInstall = lib.optionalString (!isNative) ''

  postInstall = ''

    cp ${lib.getExe lit} $out/bin/llvm-lit

  '' + (lib.optionalString (!isNative) ''

    cp -a NATIVE/bin/llvm-config $out/bin/llvm-config-native

  '';

  '');

  doCheck = buildTests;

From c5d4087519eae6f41c80bbd8ffbcc9390db44c7f Mon Sep 17 00:00:00 2001

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

Date: Thu, 10 Oct 2024 19:19:18 +0000

Subject: [PATCH] cmake.py: propagate cmakeFlags from environment

---

 tools/setup_helpers/cmake.py | 6 ++++++

 1 file changed, 6 insertions(+)

diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py

index 4b605fe5975..ea1d6a1ef46 100644

--- a/tools/setup_helpers/cmake.py

+++ b/tools/setup_helpers/cmake.py

@@ -332,6 +332,12 @@ class CMake:

                         file=sys.stderr,

                     )

                     print(e, file=sys.stderr)

+

+        # Nixpkgs compat:

+        if "cmakeFlags" in os.environ:

+            import shlex

+            args.extend(shlex.split(os.environ["cmakeFlags"]))

+

         # According to the CMake manual, we should pass the arguments first,

         # and put the directory as the last element. Otherwise, these flags

         # may not be passed correctly.

-- 

2.46.0

  removeReferencesTo,

  # Build inputs

  darwin,

  numactl,

  Accelerate,

  CoreServices,

  libobjc,

  # Propagated build inputs

  astunparse,

  tritonSupport ? (!stdenv.hostPlatform.isDarwin),

  triton,

  # TODO: 1. callPackage needs to learn to distinguish between the task

  #          of "asking for an attribute from the parent scope" and

  #          the task of "exposing a formal parameter in .override".

  # TODO: 2. We should probably abandon attributes such as `torchWithCuda` (etc.)

  #          as they routinely end up consuming the wrong arguments\

  #          (dependencies without cuda support).

  #          Instead we should rely on overlays and nixpkgsFun.

  # (@SomeoneSerge)

  _tritonEffective ? if cudaSupport then triton-cuda else triton,

  triton-cuda,

  # Unit tests

  hypothesis,

  psutil,

    ;

  inherit (cudaPackages) cudaFlags cudnn nccl;

  triton = throw "python3Packages.torch: use _tritonEffective instead of triton to avoid divergence";

  rocmPackages = rocmPackages_5;

  setBool = v: if v then "1" else "0";

      # Allow setting PYTHON_LIB_REL_PATH with an environment variable.

      # https://github.com/pytorch/pytorch/pull/128419

      ./passthrough-python-lib-rel-path.patch

      ./0001-cmake.py-propagate-cmakeFlags-from-environment.patch

    ]

    ++ lib.optionals cudaSupport [ ./fix-cmake-cuda-toolkit.patch ]

    ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [

    ];

  postPatch =

    lib.optionalString rocmSupport ''

    ''

      substituteInPlace cmake/public/cuda.cmake \

        --replace-fail \

          'message(FATAL_ERROR "Found two conflicting CUDA' \

          'message(WARNING "Found two conflicting CUDA' \

        --replace-warn \

          "set(CUDAToolkit_ROOT" \

          "# Upstream: set(CUDAToolkit_ROOT"

      substituteInPlace third_party/gloo/cmake/Cuda.cmake \

        --replace-warn "find_package(CUDAToolkit 7.0" "find_package(CUDAToolkit"

    ''

    + lib.optionalString rocmSupport ''

      # https://github.com/facebookincubator/gloo/pull/297

      substituteInPlace third_party/gloo/cmake/Hipify.cmake \

        --replace "\''${HIPIFY_COMMAND}" "python \''${HIPIFY_COMMAND}"

  # NB technical debt: building without NNPACK as workaround for missing `six`

  USE_NNPACK = 0;

  cmakeFlags =

    [

      # (lib.cmakeBool "CMAKE_FIND_DEBUG_MODE" true)

      (lib.cmakeFeature "CUDAToolkit_VERSION" cudaPackages.cudaVersion)

    ]

    ++ lib.optionals cudaSupport [

      # Unbreaks version discovery in enable_language(CUDA) when wrapping nvcc with ccache

      # Cf. https://gitlab.kitware.com/cmake/cmake/-/issues/26363

      (lib.cmakeFeature "CMAKE_CUDA_COMPILER_TOOLKIT_VERSION" cudaPackages.cudaVersion)

    ];

  preBuild = ''

    export MAX_JOBS=$NIX_BUILD_CORES

    ${python.pythonOnBuildForHost.interpreter} setup.py build --cmake-only

    ++ lib.optionals (cudaSupport || rocmSupport) [ effectiveMagma ]

    ++ lib.optionals stdenv.hostPlatform.isLinux [ numactl ]

    ++ lib.optionals stdenv.hostPlatform.isDarwin [

      Accelerate

      CoreServices

      libobjc

      darwin.apple_sdk.frameworks.Accelerate

      darwin.apple_sdk.frameworks.CoreServices

      darwin.libobjc

    ]

    ++ lib.optionals tritonSupport [ triton ]

    ++ lib.optionals tritonSupport [ _tritonEffective ]

    ++ lib.optionals MPISupport [ mpi ]

    ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];

    # torch/csrc requires `pybind11` at runtime

    pybind11

  ] ++ lib.optionals tritonSupport [ triton ];

  ] ++ lib.optionals tritonSupport [ _tritonEffective ];

  propagatedCxxBuildInputs =

    [ ] ++ lib.optionals MPISupport [ mpi ] ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];

      thoughtpolice

      tscholak

    ]; # tscholak esp. for darwin-related builds

    platforms = with lib.platforms; linux ++ lib.optionals (!cudaSupport && !rocmSupport) darwin;

    platforms =

      lib.platforms.linux

      ++ lib.optionals (!cudaSupport && !rocmSupport) lib.platforms.darwin;

    broken = builtins.any trivial.id (builtins.attrValues brokenConditions);

  };

}

diff --git a/python/setup.py b/python/setup.py

index 18764ec13..b3bb5b60a 100644

--- a/python/setup.py

+++ b/python/setup.py

@@ -269,10 +269,6 @@ class CMakeBuild(build_ext):

         subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env)

         subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=cmake_dir)

-

-download_and_copy_ptxas()

-

-

 setup(

     name="triton",

     version="2.1.0",

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

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

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

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

---

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

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

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

index d7baeb286..d334dce77 100644

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

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

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

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

     include_dirs = include_dirs + [srcdir, py_include_dir]

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

+

+    # Nixpkgs support branch

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

+    cc_cmd_extra_flags = "@ccCmdExtraFlags@"

+    if cc_cmd_extra_flags != ("@" + "ccCmdExtraFlags@"): # substituteAll hack

+        import shlex

+        cc_cmd.extend(shlex.split(cc_cmd_extra_flags))

+

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

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

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

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

     ret = subprocess.check_call(cc_cmd)

     if ret == 0:

         return so

-- 

2.46.0