flybrain/services/flysim/crates/flybrain-core/cuda/cu2ptx.c
acamilo 660c3cf00d
Some checks failed
ci / node 22 (test + typecheck) (push) Has been cancelled
ci / rust stable (cargo test --workspace --release) (push) Has been cancelled
ci / infra/tests/lint.sh (push) Has been cancelled
ci / playwright apps/stage (allowed to fail) (push) Has been cancelled
flybrain v0.4.0: public tree (history retained privately)
2026-09-21 15:09:46 +00:00

105 lines
3.6 KiB
C

/*
* Compile a .cu file to PTX with NVRTC, which is the only CUDA compiler this project can install
* without root: the `nvidia-cuda-nvcc-cu12` wheel ships `ptxas` and `libnvvm` but no `nvcc`
* frontend and no `cicc`, while `nvidia-cuda-nvrtc-cu12` ships a complete compiler as a library.
*
* cc -o cu2ptx cu2ptx.c -ldl
* LIBNVRTC=~/cuda/nvidia/cuda_nvrtc/lib/libnvrtc.so.12 ./cu2ptx lif.cu lif.ptx compute_75
*
* `--fmad=false` is not optional: the bit-exactness contract needs every multiply and add to round
* separately, exactly as the CPU kernel's f64 expressions do.
*/
#include <dlfcn.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
typedef void *nvrtcProgram;
typedef int (*create_fn)(nvrtcProgram *, const char *, const char *, int, const char **,
const char **);
typedef int (*compile_fn)(nvrtcProgram, int, const char **);
typedef int (*size_fn)(nvrtcProgram, size_t *);
typedef int (*get_fn)(nvrtcProgram, char *);
typedef const char *(*str_fn)(int);
int main(int argc, char **argv) {
if (argc < 4) {
fprintf(stderr, "usage: cu2ptx <input.cu> <output.ptx> <compute_arch>\n");
return 2;
}
const char *libname = getenv("LIBNVRTC");
if (libname == NULL) {
libname = "libnvrtc.so.12";
}
void *lib = dlopen(libname, RTLD_NOW);
if (lib == NULL) {
fprintf(stderr, "cu2ptx: cannot load %s: %s\n", libname, dlerror());
return 1;
}
create_fn create = (create_fn)dlsym(lib, "nvrtcCreateProgram");
compile_fn compile = (compile_fn)dlsym(lib, "nvrtcCompileProgram");
size_fn log_size = (size_fn)dlsym(lib, "nvrtcGetProgramLogSize");
get_fn get_log = (get_fn)dlsym(lib, "nvrtcGetProgramLog");
size_fn ptx_size = (size_fn)dlsym(lib, "nvrtcGetPTXSize");
get_fn get_ptx = (get_fn)dlsym(lib, "nvrtcGetPTX");
str_fn error_string = (str_fn)dlsym(lib, "nvrtcGetErrorString");
if (!create || !compile || !log_size || !get_log || !ptx_size || !get_ptx || !error_string) {
fprintf(stderr, "cu2ptx: %s is missing NVRTC entry points\n", libname);
return 1;
}
FILE *input = fopen(argv[1], "rb");
if (input == NULL) {
fprintf(stderr, "cu2ptx: cannot open %s\n", argv[1]);
return 1;
}
fseek(input, 0, SEEK_END);
long length = ftell(input);
fseek(input, 0, SEEK_SET);
char *source = malloc((size_t)length + 1);
if (fread(source, 1, (size_t)length, input) != (size_t)length) {
fprintf(stderr, "cu2ptx: short read on %s\n", argv[1]);
return 1;
}
source[length] = '\0';
fclose(input);
nvrtcProgram program = NULL;
int status = create(&program, source, argv[1], 0, NULL, NULL);
if (status != 0) {
fprintf(stderr, "cu2ptx: nvrtcCreateProgram: %s\n", error_string(status));
return 1;
}
char arch[64];
snprintf(arch, sizeof(arch), "--gpu-architecture=%s", argv[3]);
const char *options[] = {arch, "--fmad=false", "--std=c++17", "-default-device"};
status = compile(program, 4, options);
size_t log_bytes = 0;
log_size(program, &log_bytes);
if (log_bytes > 1) {
char *log = malloc(log_bytes);
get_log(program, log);
fputs(log, stderr);
free(log);
}
if (status != 0) {
fprintf(stderr, "cu2ptx: nvrtcCompileProgram: %s\n", error_string(status));
return 1;
}
size_t bytes = 0;
ptx_size(program, &bytes);
char *ptx = malloc(bytes);
get_ptx(program, ptx);
FILE *output = fopen(argv[2], "wb");
if (output == NULL) {
fprintf(stderr, "cu2ptx: cannot write %s\n", argv[2]);
return 1;
}
fwrite(ptx, 1, bytes - 1, output);
fclose(output);
fprintf(stderr, "cu2ptx: wrote %s (%zu bytes) for %s\n", argv[2], bytes - 1, argv[3]);
return 0;
}