1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
{ 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 ];
};
}
|