Skip to content

feat(libtorch): report device identity per index - #35

Open
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:libtorch-device-identity
Open

NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:libtorch-device-identity

Conversation

@NicolasRouquette

Copy link
Copy Markdown
Contributor

What

Runtime.Autograd.LibTorch.deviceInfo (index : UInt32) : IO DeviceInfo and
currentDeviceInfo, the LibTorch follow-up to #30: name, compute capability (major * 10 + minor, so 86 reads as sm_86), SM count, core and memory clocks, memory bus width, total
memory, and the driver and runtime versions, plus DeviceInfo.peakBandwidthGBs and a one-line
DeviceInfo.format for a benchmark header. Buffer.allocatorStats answers how much device
memory there is; this answers which device.

The review finding from #30

#30 cached one cudaDeviceProp under a process-wide pthread_once, so after switching devices
it kept reporting device 0 and could combine fields from two cards. Here every query takes an
explicit index and reads at::cuda::getDeviceProperties(index), the SDK's own per-device cache,
so the answer never depends on which device was selected earlier; currentDeviceInfo is
deviceInfo (← getDevice). The test switches through every visible device in a fresh process
and checks that currentDeviceInfo reports each index in turn.

Clocks and bus width come from cudaDeviceGetAttribute, since CUDA 13 removed clockRate and
memoryClockRate from cudaDeviceProp and the attribute enums exist in 12.x and 13.x. An index
at or above device_count is a TORCH_CHECK failure that surfaces as an IO error, like every
other runtime control; the build without LibTorch fails the same way through unavailable.c
rather than returning a record of zeroes.

What is not carried over

#30 also reported the architectures the binary held native code for (__CUDA_ARCH_LIST__) and
warned when the device was being served through PTX JIT. TorchLean no longer compiles device
code, and ATen's C++ API does not expose the SDK's packaged architecture list (only the Python
package reads it back), so those fields are gone. The module docstring says so.

Tests

NN/Tests/Runtime/Cuda/DeviceInfo.lean, called from NN/Tests/Suite.lean under every runtime
status:

  • with a device: every visible index reports itself, a nonempty name, a nonzero capability, SM
    count, and memory; currentDeviceInfo agrees with getDevice; the index at the visible count
    is rejected; and a child process (TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch, the pattern of the
    memory probes) walks setDevice through every device before any buffer exists, because
    setDevice refuses while wrappers are live;
  • without LibTorch: deviceInfo 0 fails.

Checks run

RTX A4500 (sm_86, driver 580.126), pip torch 2.11.0+cu128 as the SDK, CUDA 12.8 toolkit for
SDK discovery, Lean v4.34.0:

scripts/libtorch_build.py (torchlean.cpp, C++20, GNU 11.5)    Built target torchlean_libtorch
lake -R -Kcuda=true build nn_tests_suite                       Build completed successfully
TORCHLEAN_REQUIRE_CUDA=1 nn_tests_suite                        (see the device lines below)
lake build NN NNTests nn_tests_suite (default)                 Build completed successfully (10949 jobs)
nn_tests_suite (default build)                                 device identity: queries are rejected without LibTorch
                                                               == TorchLean: all curated tests passed ==
python3 scripts/checks/repo_lint.py                            OK: no issues found.

GPU suite output of the new section:

== LibTorch device identity ==
  NVIDIA RTX A4500 (device 0, sm_86, 56 SMs, 19.548828 GiB, ~640.080000 GB/s peak, driver 13.0, runtime 12.8)
  device switching: 1 device(s) report under their own index
== TorchLean: all curated tests passed ==

Not run: the sanitizer harness and the elementwise C++ harness. This host has one GPU, so the
switching probe exercised one index; the multi-device case is covered by construction (per-index
reads) rather than by an observation.

`LibTorch.deviceInfo index` reads one device's name, compute capability, SM count, clocks, bus
width, total memory, and the driver and runtime versions through the SDK's per-device cached
properties; `currentDeviceInfo` reads the device `getDevice` selects. Querying by explicit index
keeps the answer independent of the selected device, so switching devices cannot report a stale
record. Without LibTorch, the queries fail through `IO` like the other controls.

`NN/Tests/Runtime/Cuda/DeviceInfo.lean` checks every visible device under its own index, the
current-device agreement, index rejection, and, in a fresh process selected by
`TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch`, that switching devices changes the reading.

Verified in the default build only: the `torchlean.cpp` half has not been compiled against a
LibTorch SDK yet.
NicolasRouquette added a commit to NicolasRouquette/TorchLean that referenced this pull request Oct 3, 2026
Brings in the upstream PR branch (lean-dojo#35) at 65fd9d9, one
commit on upstream main b062b9a: Runtime.Autograd.LibTorch.deviceInfo and
currentDeviceInfo read at::cuda::getDeviceProperties(index) and
cudaDeviceGetAttribute per index, with DeviceInfo.format for benchmark
headers.

Successor of the closed lean-dojo#30. The compiled-architecture fields are gone with
the custom kernels.
@Robertboy18

Copy link
Copy Markdown
Member

Hey Nicolas, thanks! Per-device identity is useful. Could you check this together with #36? The CUDA property queries need CPU-only SDK handling, and the tests should not assume nativeAvailable means a CUDA device exists. Also, the statement in DeviceInfo.lean that TorchLean compiles no device code of its own will need updating when our pending custom-kernel/NVRTC work lands. Please keep the SDK-packaged architecture list separate from runtime-compiled kernels in that explanation. Keeping this open for those integration checks.

This branch has not been deployed

No deployments
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.

2 participants