Docs / jax4j-core / com.nablatensor.jax4j.engine

final class

CudaAadCodegen

Translates a recorded AadTape into CUDA C source implementing one scenario per thread: a fully unrolled forward sweep followed by the adjoint sweep, both in registers, with the per-scenario random stream generated in-thread so a replay touches no memory beyond the handful of inputs and the per-block reduction.

The emitted source is portable CUDA C: AadKernel compiles it with NVRTC, and the ROCm engine (com.nablatensor.jax4j.engine.rocm) compiles the exact same string with HIPRTC, which accepts it unchanged. Kept public for that second consumer; not part of the supported API.

View source ↗Full module reference on javadoc.io ↗

Fields

static final String KERNEL_NAME

static final int BLOCK

Methods

static String generate(AadTape tape, AadOptions options)

static String generateCheckpointed(AadTape tape, AadOptions options, AadCheckpointPlan plan)

The checkpointed counterpart of generate: same maths, but the forward runs segment by segment and the reverse recomputes each segment from a small boundary checkpoint instead of keeping every v_i live. The kernel takes two extra parameters — the per-invocation scratch buffer and the invocation count that strides it. See AadCheckpointPlan.