tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 cudaPackages,
3 lib,
4 writeGpuTestPython,
5 # Configuration flags
6 openCVFirst,
7 useOpenCVDefaultCuda,
8 useTorchDefaultCuda,
9}:
10let
11 inherit (lib.strings) optionalString;
12
13 openCVBlock = ''
14
15 import cv2
16 print("OpenCV version:", cv2.__version__)
17
18 # Ensure OpenCV can access the GPU.
19 assert cv2.cuda.getCudaEnabledDeviceCount() > 0, "No CUDA devices found for OpenCV"
20 print("OpenCV CUDA device:", cv2.cuda.printCudaDeviceInfo(cv2.cuda.getDevice()))
21
22 # Ensure OpenCV can access the GPU.
23 print(cv2.getBuildInformation())
24
25 a = cv2.cuda.GpuMat(size=(256, 256), type=cv2.CV_32S, s=1)
26 b = cv2.cuda.GpuMat(size=(256, 256), type=cv2.CV_32S, s=1)
27 c = int(cv2.cuda.sum(cv2.cuda.add(a, b))[0]) # OpenCV returns a Scalar float object.
28
29 assert c == 2 * 256 * 256, f"Expected {2 * 256 * 256} OpenCV, got {c}"
30
31 '';
32
33 torchBlock = ''
34
35 import torch
36 print("Torch version:", torch.__version__)
37
38 # Set up the GPU.
39 torch.cuda.init()
40 # Ensure the GPU is available.
41 assert torch.cuda.is_available(), "CUDA is not available to Torch"
42 print("Torch CUDA device:", torch.cuda.get_device_properties(torch.cuda.current_device()))
43
44 a = torch.ones(256, 256, dtype=torch.int32).cuda()
45 b = torch.ones(256, 256, dtype=torch.int32).cuda()
46 c = (a + b).sum().item()
47 assert c == 2 * 256 * 256, f"Expected {2 * 256 * 256} for Torch, got {c}"
48
49 '';
50
51 content = if openCVFirst then openCVBlock + torchBlock else torchBlock + openCVBlock;
52
53 torchName = "torch" + optionalString useTorchDefaultCuda "-with-default-cuda";
54 openCVName = "opencv4" + optionalString useOpenCVDefaultCuda "-with-default-cuda";
55in
56# TODO: Ensure the expected CUDA libraries are loaded.
57# TODO: Ensure GPU access works as expected.
58writeGpuTestPython {
59 name = if openCVFirst then "${openCVName}-then-${torchName}" else "${torchName}-then-${openCVName}";
60 libraries =
61 # NOTE: These are purposefully in this order.
62 pythonPackages:
63 let
64 effectiveOpenCV = pythonPackages.opencv4.override (prevAttrs: {
65 cudaPackages = if useOpenCVDefaultCuda then prevAttrs.cudaPackages else cudaPackages;
66 });
67 effectiveTorch = pythonPackages.torchWithCuda.override (prevAttrs: {
68 cudaPackages = if useTorchDefaultCuda then prevAttrs.cudaPackages else cudaPackages;
69 });
70 in
71 if openCVFirst then
72 [
73 effectiveOpenCV
74 effectiveTorch
75 ]
76 else
77 [
78 effectiveTorch
79 effectiveOpenCV
80 ];
81} content