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

final class

RocmAadKernel

A recorded tape compiled into a HIP replay kernel.

Recording and HIPRTC compilation happen once, in compile; after that a replay costs one hipModuleLaunchKernel plus a small block-partial download, no matter how many scenarios it covers. The generated source is CudaAadCodegen's portable CUDA C — HIPRTC takes it unchanged — and inputs are kernel arguments, so setInput re-values a shifted parameters with no re-record and no re-compile.

For a long adjoint tape (over nablatensor.checkpoint.minNodes, 512 by default; disable with -Dnablatensor.checkpoint=off) the kernel is segment-checkpointed instead of fully unrolled — see AadCheckpointPlan — which needs a per-invocation scratch buffer and two extra kernel arguments.

View source ↗Full module reference on javadoc.io ↗

Methods

static RocmAadKernel compile(AadTape tape, AadOptions options)

String engineName()

double compileSeconds()

protected double defaultMaxChunkSeconds()

An integrated GPU shares the DRM scheduler with the display; a wedged kernel there can take the driver down, so a dispatch is kept well below the seconds-range hang check. Override with -Dnablatensor.maxLaunchSeconds.

AadResult replay(long paths, long pathOffset, long seed)

void close()