pkgs/by-name/gp/gpu-burn/package.nix

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