From 768730df9dce370b414c2fe3409b4a4ffa71d0a1 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 17 Sep 2024 00:07:38 +0200
Subject: [PATCH] python312Packages.bitsandbytes: 0.43.1 -> 0.44.1

Diff:
https://github.com/TimDettmers/bitsandbytes/compare/refs/tags/0.43.1...0.44.1

Changelog:
https://github.com/TimDettmers/bitsandbytes/releases/tag/0.44.1
---
 .../python-modules/bitsandbytes/default.nix   | 99 +++++++++++--------
 1 file changed, 56 insertions(+), 43 deletions(-)

diff --git a/pkgs/development/python-modules/bitsandbytes/default.nix b/pkgs/development/python-modules/bitsandbytes/default.nix
index 9bfc4a89784fe..1768fc1652a93 100644
--- a/pkgs/development/python-modules/bitsandbytes/default.nix
+++ b/pkgs/development/python-modules/bitsandbytes/default.nix
@@ -1,31 +1,34 @@
 {
   lib,
+  torch,
+  symlinkJoin,
   buildPythonPackage,
   fetchFromGitHub,
-  python,
-  pythonOlder,
+  cmake,
   setuptools,
-  wheel,
-  torch,
   scipy,
-  symlinkJoin,
 }:
 
 let
   pname = "bitsandbytes";
-  version = "0.43.1";
+  version = "0.44.1";
 
-  inherit (torch) cudaCapabilities cudaPackages cudaSupport;
-  inherit (cudaPackages) backendStdenv cudaVersion;
+  inherit (torch) cudaPackages cudaSupport;
+  inherit (cudaPackages) cudaVersion;
+
+  cudaVersionString = lib.replaceStrings [ "." ] [ "" ] (lib.versions.majorMinor cudaVersion);
 
   # NOTE: torchvision doesn't use cudnn; torch does!
   #   For this reason it is not included.
   cuda-common-redist = with cudaPackages; [
-    cuda_cccl # <thrust/*>
-    libcublas # cublas_v2.h
+    (lib.getDev cuda_cccl) # <thrust/*>
+    (lib.getDev libcublas) # cublas_v2.h
+    (lib.getLib libcublas)
     libcurand
     libcusolver # cusolverDn.h
-    libcusparse # cusparse.h
+    (lib.getDev libcusparse) # cusparse.h
+    (lib.getLib libcusparse)
+    libcusparse.lib
   ];
 
   cuda-native-redist = symlinkJoin {
@@ -33,7 +36,9 @@ let
     paths =
       with cudaPackages;
       [
-        cuda_cudart # cuda_runtime.h cuda_runtime_api.h
+        (lib.getDev cuda_cudart) # cuda_runtime.h cuda_runtime_api.h
+        (lib.getLib cuda_cudart)
+        (lib.getStatic cuda_cudart)
         cuda_nvcc
       ]
       ++ cuda-common-redist;
@@ -48,47 +53,55 @@ buildPythonPackage {
   inherit pname version;
   pyproject = true;
 
-  disabled = pythonOlder "3.7";
-
   src = fetchFromGitHub {
     owner = "TimDettmers";
     repo = "bitsandbytes";
     rev = "refs/tags/${version}";
-    hash = "sha256-GFbFKPdV96DXPA+PZO4h0zdBclN670fb0PGv4QPHWHU=";
+    hash = "sha256-yvxD5ymMK5p4Xg7Csx/90mPV3yxUC6QUuF/8BKO2p0k=";
   };
 
-  postPatch =
-    ''
-      substituteInPlace Makefile --replace "/usr/bin/g++" "g++" --replace "lib64" "lib"
-      substituteInPlace bitsandbytes/cuda_setup/main.py  \
-        --replace "binary_path = package_dir / self.binary_name"  \
-                  "binary_path = Path('$out/${python.sitePackages}/${pname}')/self.binary_name"
-    ''
-    + lib.optionalString torch.cudaSupport ''
-      substituteInPlace bitsandbytes/cuda_setup/main.py  \
-        --replace "/usr/local/cuda/lib64" "${cuda-native-redist}/lib"
-    '';
+  # By default, which library is loaded depends on the result of `torch.cuda.is_available()`.
+  # When `cudaSupport` is enabled, bypass this check and load the cuda library unconditionnally.
+  # Indeed, in this case, only `libbitsandbytes_cuda124.so` is built. `libbitsandbytes_cpu.so` is not.
+  # Also, hardcode the path to the previously built library instead of relying on
+  # `get_cuda_bnb_library_path(cuda_specs)` which relies on `torch.cuda` too.
+  #
+  # WARNING: The cuda library is currently named `libbitsandbytes_cudaxxy` for cuda version `xx.y`.
+  # This upstream convention could change at some point and thus break the following patch.
+  postPatch = lib.optionalString cudaSupport ''
+    substituteInPlace bitsandbytes/cextension.py \
+      --replace-fail "if cuda_specs:" "if True:" \
+      --replace-fail \
+        "cuda_binary_path = get_cuda_bnb_library_path(cuda_specs)" \
+        "cuda_binary_path = PACKAGE_DIR / 'libbitsandbytes_cuda${cudaVersionString}.so'"
+  '';
 
   CUDA_HOME = "${cuda-native-redist}";
 
-  preBuild =
-    if torch.cudaSupport then
-      with torch.cudaPackages;
-      let
-        cudaVersion = lib.concatStrings (lib.splitVersion torch.cudaPackages.cudaMajorMinorVersion);
-      in
-      ''make CUDA_VERSION=${cudaVersion} cuda${cudaMajorVersion}x''
-    else
-      ''make CUDA_VERSION=CPU cpuonly'';
-
   nativeBuildInputs = [
+    cmake
+  ];
+
+  build-system = [
     setuptools
-    wheel
-  ] ++ lib.optionals torch.cudaSupport [ cuda-native-redist ];
+  ];
+
+  buildInputs = lib.optionals cudaSupport [ cuda-redist ];
+
+  cmakeFlags = [
+    (lib.cmakeFeature "-DCOMPUTE_BACKEND" (if cudaSupport then "cuda" else "cpu"))
+  ];
+  NVCC_PREPEND_FLAGS = lib.optionals cudaSupport [
+    "-I${cuda-native-redist}/include"
+    "-L${cuda-native-redist}/lib"
+  ];
 
-  buildInputs = lib.optionals torch.cudaSupport [ cuda-redist ];
+  preBuild = ''
+    make -j $NIX_BUILD_CORES
+    cd .. # leave /build/source/build
+  '';
 
-  propagatedBuildInputs = [
+  dependencies = [
     scipy
     torch
   ];
@@ -97,11 +110,11 @@ buildPythonPackage {
 
   pythonImportsCheck = [ "bitsandbytes" ];
 
-  meta = with lib; {
+  meta = {
     description = "8-bit CUDA functions for PyTorch";
     homepage = "https://github.com/TimDettmers/bitsandbytes";
     changelog = "https://github.com/TimDettmers/bitsandbytes/releases/tag/${version}";
-    license = licenses.mit;
-    maintainers = with maintainers; [ bcdarwin ];
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ bcdarwin ];
   };
 }