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
|
{
lib,
config,
buildPythonPackage,
fetchFromGitHub,
replaceVars,
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";
tag = version;
hash = "sha256-xUMvMBK1UhZaMZfik0Ia6+siyZGpCkBV+LTnQvzt/rw=";
};
patches = [
(replaceVars ./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)
disabledTests =
lib.optionals (!testCudaRuntime) [
"TestCUDADevice"
]
++ lib.optionals (!testRocmRuntime) [
"TestHIPDevice"
];
pytestFlags = 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.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs
'';
# If neither rocmSupport or cudaSupport is enabled, no tests are selected
dontUsePytestCheck = !(rocmSupport || cudaSupport) && (!testOpenclRuntime);
meta = {
description = "Ctypes wrappers for HIP, CUDA, and OpenCL";
homepage = "https://github.com/tinygrad/gpuctypes";
license = lib.licenses.mit;
maintainers = with lib.maintainers; [
GaetanLepage
matthewcroughan
wozeparrot
];
};
}
|
