Merge pull request #287914 from GaetanLepage/gpuctypes

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

{ 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 ];

  };

}

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

{ 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 ];

  };

}

  gpsoauth = callPackage ../development/python-modules/gpsoauth { };

  gpuctypes = callPackage ../development/python-modules/gpuctypes { };

  gpustat = callPackage ../development/python-modules/gpustat { };

  gpxpy = callPackage ../development/python-modules/gpxpy { };

  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 { };