Unverified Commit 961dbce1 authored by Someone's avatar Someone Committed by GitHub
Browse files

Merge pull request #287914 from GaetanLepage/gpuctypes

python3Packages.tinygrad: init at 0.8.0
parents faeae5d4 24833cd1
Loading
Loading
Loading
Loading
+44 −0
Original line number Diff line number Diff line
From d448321436e8314d3e2a6a09d4017c4bc10f612d Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sat, 17 Feb 2024 17:37:22 +0100
Subject: [PATCH] fix-dlopen-cuda

---
 gpuctypes/cuda.py | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/gpuctypes/cuda.py b/gpuctypes/cuda.py
index acba81c..091f7f7 100644
--- a/gpuctypes/cuda.py
+++ b/gpuctypes/cuda.py
@@ -143,9 +143,25 @@ def char_pointer_cast(string, encoding='utf-8'):
 
 
 
+NAME_TO_PATHS = {
+    "libcuda.so": ["@driverLink@/lib/libcuda.so"],
+    "libnvrtc.so": ["@libnvrtc@"],
+}
+def _try_dlopen(name):
+    try:
+        return ctypes.CDLL(name)
+    except OSError:
+        pass
+    for candidate in NAME_TO_PATHS.get(name, []):
+        try:
+            return ctypes.CDLL(candidate)
+        except OSError:
+	        pass
+    raise RuntimeError(f"{name} not found")
+
 _libraries = {}
-_libraries['libcuda.so'] = ctypes.CDLL(ctypes.util.find_library('cuda'))
-_libraries['libnvrtc.so'] = ctypes.CDLL(ctypes.util.find_library('nvrtc'))
+_libraries['libcuda.so'] = _try_dlopen('libcuda.so')
+_libraries['libnvrtc.so'] = _try_dlopen('libnvrtc.so')
 
 
 cuuint32_t = ctypes.c_uint32
-- 
2.43.0
+127 −0
Original line number Diff line number Diff line
{ lib
, config
, buildPythonPackage
, fetchFromGitHub
, substituteAll
, addDriverRunpath
, cudaSupport ? config.cudaSupport
, rocmSupport ? config.rocmSupport
, cudaPackages
, setuptools
, ocl-icd
, rocmPackages
, pytestCheckHook
, gpuctypes
, testCudaRuntime ? false
, testOpenclRuntime ? false
, testRocmRuntime ? false
}:
assert testCudaRuntime -> cudaSupport;
assert testRocmRuntime -> rocmSupport;

buildPythonPackage rec {
  pname = "gpuctypes";
  version = "0.3.0";
  pyproject = true;

  src = fetchFromGitHub {
    repo = "gpuctypes";
    owner = "tinygrad";
    rev = "refs/tags/${version}";
    hash = "sha256-xUMvMBK1UhZaMZfik0Ia6+siyZGpCkBV+LTnQvzt/rw=";
  };

  patches = [
    (substituteAll {
      src = ./0001-fix-dlopen-cuda.patch;
      inherit (addDriverRunpath) driverLink;
      libnvrtc =
        if cudaSupport
        then "${lib.getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so"
        else "Please import nixpkgs with `config.cudaSupport = true`";
    })
  ];

  nativeBuildInputs = [
    setuptools
  ];

  postPatch = ''
    substituteInPlace gpuctypes/opencl.py \
      --replace "ctypes.util.find_library('OpenCL')" "'${ocl-icd}/lib/libOpenCL.so'"
  ''
  # hipGetDevicePropertiesR0600 is a symbol from rocm-6. We are currently at rocm-5.
  # We are not sure that this works. Remove when rocm gets updated to version 6.
  + lib.optionalString rocmSupport ''
    substituteInPlace gpuctypes/hip.py \
      --replace "/opt/rocm/lib/libamdhip64.so" "${rocmPackages.clr}/lib/libamdhip64.so" \
      --replace "/opt/rocm/lib/libhiprtc.so" "${rocmPackages.clr}/lib/libhiprtc.so" \
      --replace "hipGetDevicePropertiesR0600" "hipGetDeviceProperties"

    substituteInPlace gpuctypes/comgr.py \
      --replace "/opt/rocm/lib/libamd_comgr.so" "${rocmPackages.rocm-comgr}/lib/libamd_comgr.so"
  '';

  pythonImportsCheck = [ "gpuctypes" ];

  nativeCheckInputs = [
    pytestCheckHook
  ];

  disabledTestPaths = lib.optionals (!testOpenclRuntime) [
    "test/test_opencl.py"
  ] ++ lib.optionals (!rocmSupport) [
    "test/test_hip.py"
  ] ++ lib.optionals (!cudaSupport) [
    "test/test_cuda.py"
  ];

  # Require GPU access to run (not available in the sandbox)
  pytestFlagsArray = lib.optionals (!testCudaRuntime) [
    "-k" "'not TestCUDADevice'"
  ] ++ lib.optionals (!testRocmRuntime) [
    "-k" "'not TestHIPDevice'"
  ] ++ lib.optionals (testCudaRuntime || testOpenclRuntime || testRocmRuntime) [
    "-v"
  ];

  # Running these tests requires special configuration on the builder.
  # e.g. https://github.com/NixOS/nixpkgs/pull/256230 implements a nix
  # pre-build hook which exposes the devices and the drivers in the sandbox
  # based on requiredSystemFeatures:
  requiredSystemFeatures = lib.optionals testCudaRuntime [
    "cuda"
  ] ++ lib.optionals testOpenclRuntime [
    "opencl"
  ] ++ lib.optionals testRocmRuntime [
    "rocm"
  ];

  passthru.gpuChecks = {
    cuda = gpuctypes.override {
      cudaSupport = true;
      testCudaRuntime = true;
    };
    opencl = gpuctypes.override {
      testOpenclRuntime = true;
    };
    rocm = gpuctypes.override {
      rocmSupport = true;
      testRocmRuntime = true;
    };
  };

  preCheck = lib.optionalString (cudaSupport && !testCudaRuntime) ''
    addToSearchPath LD_LIBRARY_PATH ${lib.getLib cudaPackages.cuda_cudart}/lib/stubs
  '';

  # If neither rocmSupport or cudaSupport is enabled, no tests are selected
  dontUsePytestCheck = !(rocmSupport || cudaSupport) && (!testOpenclRuntime);

  meta = with lib; {
    description = "Ctypes wrappers for HIP, CUDA, and OpenCL";
    homepage = "https://github.com/tinygrad/gpuctypes";
    license = licenses.mit;
    maintainers = with maintainers; [ GaetanLepage matthewcroughan wozeparrot ];
  };
}
+32 −0
Original line number Diff line number Diff line
diff --git a/gpuctypes/cuda.py b/gpuctypes/cuda.py
index acba81c..aac5fc7 100644
--- a/gpuctypes/cuda.py
+++ b/gpuctypes/cuda.py
@@ -143,9 +143,25 @@ def char_pointer_cast(string, encoding='utf-8'):
 
 
 
+NAME_TO_PATHS = {
+    "libcuda.so": ["@driverLink@/lib/libcuda.so"],
+    "libnvrtc.so": ["@libnvrtc@"],
+}
+def _try_dlopen(name):
+    try:
+        return ctypes.CDLL(name)
+    except OSError:
+        pass
+    for candidate in NAME_TO_PATHS.get(name, []):
+        try:
+            return ctypes.CDLL(candidate)
+        except OSError:
+	     pass
+    raise RuntimeError(f"{name} not found")
+
 _libraries = {}
-_libraries['libcuda.so'] = ctypes.CDLL(ctypes.util.find_library('cuda'))
-_libraries['libnvrtc.so'] = ctypes.CDLL(ctypes.util.find_library('nvrtc'))
+_libraries['libcuda.so'] = _try_dlopen('libcuda.so')
+_libraries['libnvrtc.so'] = _try_dlopen('libnvrtc.so')
 
 
 cuuint32_t = ctypes.c_uint32
+111 −0
Original line number Diff line number Diff line
{ lib
, buildPythonPackage
, fetchFromGitHub
, setuptools
, wheel
, gpuctypes
, numpy
, tqdm
, hypothesis
, librosa
, onnx
, pillow
, pytest-xdist
, pytestCheckHook
, safetensors
, sentencepiece
, tiktoken
, torch
, transformers
}:

buildPythonPackage rec {
  pname = "tinygrad";
  version = "0.8.0";
  pyproject = true;

  src = fetchFromGitHub {
    owner = "tinygrad";
    repo = "tinygrad";
    rev = "refs/tags/v${version}";
    hash = "sha256-QAccZ79qUbe27yUykIf22WdkxYUlOffnMlShakKfp60=";
  };

  nativeBuildInputs = [
    setuptools
    wheel
  ];

  propagatedBuildInputs = [
    gpuctypes
    numpy
    tqdm
  ];

  pythonImportsCheck = [ "tinygrad" ];

  nativeCheckInputs = [
    hypothesis
    librosa
    onnx
    pillow
    pytest-xdist
    pytestCheckHook
    safetensors
    sentencepiece
    tiktoken
    torch
    transformers
  ];

  preCheck = ''
    export HOME=$(mktemp -d)
  '';

  disabledTests = [
    # Require internet access
    "test_benchmark_openpilot_model"
    "test_bn_alone"
    "test_bn_linear"
    "test_bn_mnist"
    "test_car"
    "test_chicken"
    "test_chicken_bigbatch"
    "test_conv_mnist"
    "testCopySHMtoDefault"
    "test_data_parallel_resnet"
    "test_e2e_big"
    "test_fetch_small"
    "test_huggingface_enet_safetensors"
    "test_linear_mnist"
    "test_load_convnext"
    "test_load_enet"
    "test_load_enet_alt"
    "test_load_llama2bfloat"
    "test_load_resnet"
    "test_openpilot_model"
    "test_resnet"
    "test_shufflenet"
    "test_transcribe_batch12"
    "test_transcribe_batch21"
    "test_transcribe_file1"
    "test_transcribe_file2"
    "test_transcribe_long"
    "test_transcribe_long_no_batch"
    "test_vgg7"
  ];

  disabledTestPaths = [
    "test/extra/test_lr_scheduler.py"
    "test/models/test_mnist.py"
    "test/models/test_real_world.py"
  ];

  meta = with lib; {
    description = "A simple and powerful neural network framework";
    homepage = "https://github.com/tinygrad/tinygrad";
    changelog = "https://github.com/tinygrad/tinygrad/releases/tag/v${version}";
    license = licenses.mit;
    maintainers = with maintainers; [ GaetanLepage ];
  };
}
+4 −0
Original line number Diff line number Diff line
@@ -4939,6 +4939,8 @@ self: super: with self; {
  gpsoauth = callPackage ../development/python-modules/gpsoauth { };
  gpuctypes = callPackage ../development/python-modules/gpuctypes { };
  gpustat = callPackage ../development/python-modules/gpustat { };
  gpxpy = callPackage ../development/python-modules/gpxpy { };
@@ -14763,6 +14765,8 @@ self: super: with self; {
  tinydb = callPackage ../development/python-modules/tinydb { };
  tinygrad = callPackage ../development/python-modules/tinygrad { };
  tinyobjloader-py = callPackage ../development/python-modules/tinyobjloader-py { };
  tinyrecord = callPackage ../development/python-modules/tinyrecord { };