tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 autoAddDriverRunpath,
3 config,
4 cudaPackages,
5 fetchFromGitHub,
6 lib,
7}:
8let
9 inherit (lib.lists) last map optionals;
10 inherit (lib.trivial) boolToString;
11 inherit (config) cudaSupport;
12 inherit (cudaPackages)
13 backendStdenv
14 cuda_cccl
15 cuda_cudart
16 cuda_nvcc
17 libcublas
18 ;
19 inherit (cudaPackages.flags) cudaCapabilities dropDots isJetsonBuild;
20in
21backendStdenv.mkDerivation {
22 pname = "gpu-burn";
23 version = "0-unstable-2024-04-09";
24
25 strictDeps = true;
26
27 src = fetchFromGitHub {
28 owner = "wilicc";
29 repo = "gpu-burn";
30 rev = "9aefd7c0cc603bbc8c3c102f5338c6af26f8127c";
31 hash = "sha256-Nz0yaoHGfodaYl2HJ7p+1nasqRmxwPJ9aL9oxitXDpM=";
32 };
33
34 postPatch = ''
35 substituteInPlace gpu_burn-drv.cpp \
36 --replace-fail \
37 '#define COMPARE_KERNEL "compare.ptx"' \
38 '#define COMPARE_KERNEL "${placeholder "out"}/share/compare.ptx"'
39 substituteInPlace Makefile \
40 --replace-fail \
41 '${"\${CUDAPATH}/bin/nvcc"}' \
42 '${lib.getExe cuda_nvcc}'
43 '';
44
45 nativeBuildInputs = [
46 autoAddDriverRunpath
47 cuda_nvcc
48 ];
49
50 buildInputs = [
51 cuda_cccl # <nv/target>
52 cuda_cudart # driver_types.h
53 cuda_nvcc # crt/host_defines.h
54 libcublas # cublas_v2.h
55 ];
56
57 makeFlags = [
58 # NOTE: CUDAPATH assumes cuda_cudart is a single output containing all of lib, dev, and stubs.
59 "CUDAPATH=${cuda_cudart}"
60 "COMPUTE=${last (map dropDots cudaCapabilities)}"
61 "IS_JETSON=${boolToString isJetsonBuild}"
62 ];
63
64 installPhase = ''
65 runHook preInstall
66 mkdir -p $out/{bin,share}
67 install -Dm755 gpu_burn $out/bin/
68 install -Dm644 compare.ptx $out/share/
69 runHook postInstall
70 '';
71
72 # NOTE: Certain packages may be missing from cudaPackages on non-Linux platforms. To avoid evaluation failure,
73 # we only include the platforms where we know the package is available -- thus the conditionals setting the
74 # platforms and badPlatforms fields.
75 meta = {
76 badPlatforms = optionals (!cudaSupport) lib.platforms.all;
77 broken = !cudaSupport;
78 description = "Multi-GPU CUDA stress test";
79 homepage = "http://wili.cc/blog/gpu-burn.html";
80 license = lib.licenses.bsd2;
81 mainProgram = "gpu_burn";
82 maintainers = with lib.maintainers; [ connorbaker ];
83 platforms = optionals cudaSupport lib.platforms.linux;
84 };
85}