Skip to content

feat(cuda): selectable compilation target via -K cuda_arch, traced by Lake - #27

Open
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:cuda-arch-target
Open

feat(cuda): selectable compilation target via -K cuda_arch, traced by Lake#27
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:cuda-arch-target

Conversation

@NicolasRouquette

Copy link
Copy Markdown
Contributor

Problem

The native kernels are compiled with no -arch/-gencode, so nvcc applies its built-in
default. That default belongs to the toolkit rather than to the machine — CUDA 13.0 emits
sm_75 SASS plus compute_75 PTX. A binary built that way runs at full speed on that one
architecture, reaches a newer GPU only through forward PTX just-in-time compilation, and does
not load on an older one at all. Nothing in the build reports which of those happened, so the
cost surfaces only as unexplained throughput.

nvcc does read NVCC_APPEND_FLAGS, which is how a target can be forced today. But Lake never
sees that variable, so in a warm tree it judges the existing objects current and links kernels
compiled for the previous target. Verified, not assumed — with NVCC_APPEND_FLAGS=-arch=sm_86
set on an up-to-date build, nvcc is not invoked at all.

Change

-K cuda_arch=<spec>, with TORCHLEAN_CUDA_ARCH as a fallback so a container image or
CI job can select a target without rewriting its build command.

lake -R -K cuda=true -K cuda_arch=sm_86  build   # an A10G or an RTX A4500
lake -R -K cuda=true -K cuda_arch=native build   # whatever GPU this machine has
lake -R -K cuda=true \
  -K cuda_arch="-gencode arch=compute_75,code=sm_75 \
                -gencode arch=compute_90,code=[sm_90,compute_90]" build

A bare spec becomes -arch=<spec> (covering sm_86, compute_86, native, all,
all-major); a spec starting with - is passed to nvcc verbatim. The verbatim form is how a
multi-architecture binary is requested, and it deliberately keeps this package free of any policy
about which architecture carries the PTX.

The trace fix that makes the option real. buildNativeBackendLib passed every compiler
argument as buildO's weakArgs, which buildO excludes from the trace by design; traceArgs
was #[]. A changed optimization level or compilation target therefore left the existing objects
looking up to date. Arguments are now split the way buildO intends: include paths stay weak,
where a moved checkout does not invalidate every object, and the flags that change what the
compiler emits are traced.

Evidence

scripts/checks/cuda_arch_target.sh puts a recording stand-in for nvcc first on PATH, builds
one extern library repeatedly, and asserts the argument vector and the recompilation count:

case expected observed
no target no -arch, argv unchanged from today
-K cuda_arch=sm_86 -arch=sm_86
same target again no recompilation
-K cuda_arch=sm_89 recompiles
TORCHLEAN_CUDA_ARCH=sm_90 -arch=sm_90
option + environment option wins
verbatim -gencode list passed through, split correctly

It needs neither a CUDA toolkit nor a GPU, so it is wired into CI. It removes the stand-in's
objects on exit and restores the stored build configuration, since every build in it enables CUDA.

Checks

lake build NN (4166 jobs, no errors, warnings, or sorry) · lake test · lake -R lint ·
scripts/checks/cuda_arch_target.sh (7/7). No real-CUDA build was run: the machine used here has
no toolkit, which is precisely why the check is written against a stand-in.

Expected consequence

Moving the flags into the trace changes the trace of every object built by
buildNativeBackendLib, so the first build after this lands recompiles those four objects once —
the four .cu kernels under CUDA, or the four C stubs on a CPU-only machine.

… Lake

Device code is compiled *for* a GPU architecture, but the native kernels were compiled
with no -arch/-gencode at all, leaving nvcc on its built-in default. That default is a
property of the toolkit, not of the machine: CUDA 13.0 emits sm_75 SASS plus compute_75
PTX. On any other GPU the kernels then run through forward PTX JIT rather than native
SASS, or on an older architecture do not load at all, and nothing about the build says so.

Add -K cuda_arch=<spec>, with TORCHLEAN_CUDA_ARCH as a fallback so a container image or CI
job can select a target without rewriting its build command. A bare spec becomes
-arch=<spec>, which covers sm_86, compute_86, native, all, and all-major; a spec starting
with `-` is passed to nvcc verbatim, which is how a multi-architecture binary is requested
and which keeps this package free of any policy about which architecture carries the PTX.

Selecting a target is only meaningful if changing it rebuilds, so also split
buildNativeBackendLib's compiler arguments the way buildO intends: include paths stay in
weakArgs, where a moved checkout does not invalidate every object, and the flags that
change what the compiler emits move to traceArgs, which buildO hashes. Previously every
argument sat in weakArgs, so a changed optimization level or compilation target left the
existing objects looking current. nvcc's own NVCC_APPEND_FLAGS has the same hole and Lake
cannot close it — the option and the environment variable added here both participate in
the trace, and the documentation says which to prefer.

scripts/checks/cuda_arch_target.sh asserts the whole surface — flags reaching nvcc,
recompilation on a changed target, no recompilation on an unchanged one, the environment
fallback, its precedence, and verbatim pass-through — against a recording stand-in for
nvcc, so it needs neither a CUDA toolkit nor a GPU and runs in CI.

One consequence worth expecting: moving the flags into the trace changes the trace of
every object built by buildNativeBackendLib, so the first build after this lands
recompiles those four objects once — the four .cu kernels under CUDA, or the four C
stubs on a CPU-only machine.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant