From f0f4196f87459172fbbb55d6375fe93123d7771c Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Tue, 19 May 2026 22:01:38 -0700 Subject: [PATCH] gpu-burn: support multiple capabilities Signed-off-by: Connor Baker --- ...bin-and-make-arch-compute_X-conditio.patch | 121 ++++++++++++++++++ pkgs/by-name/gp/gpu-burn/package.nix | 19 ++- 2 files changed, 134 insertions(+), 6 deletions(-) create mode 100644 pkgs/by-name/gp/gpu-burn/0001-Emit-compare.fatbin-and-make-arch-compute_X-conditio.patch diff --git a/pkgs/by-name/gp/gpu-burn/0001-Emit-compare.fatbin-and-make-arch-compute_X-conditio.patch b/pkgs/by-name/gp/gpu-burn/0001-Emit-compare.fatbin-and-make-arch-compute_X-conditio.patch new file mode 100644 index 000000000000..a0117b075dab --- /dev/null +++ b/pkgs/by-name/gp/gpu-burn/0001-Emit-compare.fatbin-and-make-arch-compute_X-conditio.patch @@ -0,0 +1,121 @@ +From b8dc3f753bd2aefb192f95688dc12797bc041627 Mon Sep 17 00:00:00 2001 +From: Connor Baker +Date: Tue, 19 May 2026 20:16:28 -0700 +Subject: [PATCH] Emit compare.fatbin and make -arch=compute_X conditional + +cuModuleLoad accepts cubin/fatbin/PTX interchangeably, so the C++ side +just needs the filename swap. The benefit is that callers wanting to +target multiple architectures (or to opt into architecture-conditional +sm_90a / family-specific compute_100f features) can now append -gencode +flags via NVCCFLAGS; -ptx output is a single-arch text format with no +way to multiplex more than one architecture into one file. + +To let those callers fully control which architectures end up in the +fatbin, guard the default -arch=compute_$(COMPUTE) entry behind a non- +empty check. 'make COMPUTE=' now skips it; default 'make' and +'make COMPUTE=86' behave exactly as before. + +Co-Authored-By: Claude Opus 4.7 (1M context) +--- + Dockerfile | 2 +- + Makefile | 10 ++++++---- + README.md | 8 ++++++++ + gpu-burn.8 | 2 +- + gpu_burn-drv.cpp | 2 +- + 5 files changed, 17 insertions(+), 7 deletions(-) + +diff --git a/Dockerfile b/Dockerfile +index 08e5202..1c46976 100644 +--- a/Dockerfile ++++ b/Dockerfile +@@ -12,7 +12,7 @@ RUN make + FROM nvidia/cuda:${CUDA_VERSION}-runtime-${IMAGE_DISTRO} + + COPY --from=builder /build/gpu_burn /app/ +-COPY --from=builder /build/compare.ptx /app/ ++COPY --from=builder /build/compare.fatbin /app/ + + WORKDIR /app + +diff --git a/Makefile b/Makefile +index ec54a19..d23425a 100644 +--- a/Makefile ++++ b/Makefile +@@ -32,23 +32,25 @@ IMAGE_DISTRO ?= ubi8 + + override NVCCFLAGS ?= + override NVCCFLAGS += -I${CUDAPATH}/include ++ifneq ($(strip $(COMPUTE)),) + override NVCCFLAGS += -arch=compute_$(subst .,,${COMPUTE}) ++endif + + IMAGE_NAME ?= gpu-burn + + .PHONY: clean + +-gpu_burn: gpu_burn-drv.o compare.ptx ++gpu_burn: gpu_burn-drv.o compare.fatbin + g++ -o $@ $< -O3 ${LDFLAGS} + + %.o: %.cpp + g++ ${CFLAGS} -c $< + +-%.ptx: %.cu +- PATH="${PATH}:${CCPATH}:." ${NVCC} ${NVCCFLAGS} -ptx $< -o $@ ++%.fatbin: %.cu ++ PATH="${PATH}:${CCPATH}:." ${NVCC} ${NVCCFLAGS} -fatbin $< -o $@ + + clean: +- $(RM) *.ptx *.o gpu_burn ++ $(RM) *.fatbin *.o gpu_burn + + image: + docker build --build-arg CUDA_VERSION=${CUDA_VERSION} --build-arg IMAGE_DISTRO=${IMAGE_DISTRO} -t ${IMAGE_NAME} . +diff --git a/README.md b/README.md +index ca95342..1b050bf 100644 +--- a/README.md ++++ b/README.md +@@ -37,6 +37,14 @@ To override this with a different value: + + `make COMPUTE=` + ++`COMPUTE` selects a single virtual architecture for the default `-arch` ++flag. For a fat binary targeting multiple architectures, or to opt into ++architecture-conditional (`sm_90a`) or family-specific (`compute_100f`) ++features, set `COMPUTE=` to suppress the default and drive the build ++entirely from `-gencode` flags via `NVCCFLAGS`: ++ ++`make COMPUTE= NVCCFLAGS='-gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_90,code=sm_90'` ++ + CFLAGS can be added when invoking make to add to the default + list of compiler flags: + +diff --git a/gpu-burn.8 b/gpu-burn.8 +index bf071a4..b1c46c4 100644 +--- a/gpu-burn.8 ++++ b/gpu-burn.8 +@@ -19,7 +19,7 @@ GPU Burn + .br + \fB\-i\fR N Execute only on GPU N + .br +-\fB\-c\fR FILE Use FILE as compare kernel. Default is compare.ptx ++\fB\-c\fR FILE Use FILE as compare kernel. Default is compare.fatbin + .br + \fB\-stts\fR T Set timeout threshold to T seconds for using SIGTERM to abort child processes before using SIGKILL. Default is 30 + .br +diff --git a/gpu_burn-drv.cpp b/gpu_burn-drv.cpp +index eae3532..a66ec9a 100644 +--- a/gpu_burn-drv.cpp ++++ b/gpu_burn-drv.cpp +@@ -30,7 +30,7 @@ + // Matrices are SIZE*SIZE.. POT should be efficiently implemented in CUBLAS + #define SIZE 8192ul + #define USEMEM 0.9 // Try to allocate 90% of memory +-#define COMPARE_KERNEL "compare.ptx" ++#define COMPARE_KERNEL "compare.fatbin" + + // Used to report op/s, measured through Visual Profiler, CUBLAS from CUDA 7.5 + // (Seems that they indeed take the naive dim^3 approach) +-- +2.43.0 + diff --git a/pkgs/by-name/gp/gpu-burn/package.nix b/pkgs/by-name/gp/gpu-burn/package.nix index 6656b62a626c..42ece7eb5b9c 100644 --- a/pkgs/by-name/gp/gpu-burn/package.nix +++ b/pkgs/by-name/gp/gpu-burn/package.nix @@ -7,7 +7,7 @@ nix-update-script, }: let - inherit (lib.lists) last map optionals; + inherit (lib.lists) optionals; inherit (lib.trivial) boolToString; inherit (config) cudaSupport; inherit (cudaPackages) @@ -17,7 +17,7 @@ let cuda_nvcc libcublas ; - inherit (cudaPackages.flags) cudaCapabilities dropDots isJetsonBuild; + inherit (cudaPackages.flags) gencodeString isJetsonBuild; in backendStdenv.mkDerivation { pname = "gpu-burn"; @@ -32,11 +32,14 @@ backendStdenv.mkDerivation { hash = "sha256-zaGzwpdvF9dw3RypBO+g6FhjOFN8/F9+yI1+lLxLjgs="; }; + # TODO: drop once https://github.com/wilicc/gpu-burn/pull/148 lands. + patches = [ ./0001-Emit-compare.fatbin-and-make-arch-compute_X-conditio.patch ]; + postPatch = '' substituteInPlace gpu_burn-drv.cpp \ --replace-fail \ - '#define COMPARE_KERNEL "compare.ptx"' \ - '#define COMPARE_KERNEL "${placeholder "out"}/share/compare.ptx"' + '#define COMPARE_KERNEL "compare.fatbin"' \ + '#define COMPARE_KERNEL "${placeholder "out"}/share/compare.fatbin"' substituteInPlace Makefile \ --replace-fail \ '${"\${CUDAPATH}/bin/nvcc"}' \ @@ -58,15 +61,19 @@ backendStdenv.mkDerivation { makeFlags = [ # NOTE: CUDAPATH assumes cuda_cudart is a single output containing all of lib, dev, and stubs. "CUDAPATH=${cuda_cudart}" - "COMPUTE=${last (map dropDots cudaCapabilities)}" "IS_JETSON=${boolToString isJetsonBuild}" + # Empty COMPUTE suppresses the Makefile's default -arch=compute_$(COMPUTE); + # gencodeString below is the single source of truth for architectures. + "COMPUTE=" ]; + env.NVCCFLAGS = gencodeString; + installPhase = '' runHook preInstall mkdir -p $out/{bin,share} install -Dm755 gpu_burn $out/bin/ - install -Dm644 compare.ptx $out/share/ + install -Dm644 compare.fatbin $out/share/ runHook postInstall '';