Engine Internals & GPU Performance

Inspect GPU Kernel Inputs in the Java Tape

TeaVM records the actual differentiable Java tape and displays the operation sequence that a GPU code generator could consume. The browser CPU path does not emit WGSL or PTX.

Market
Java source
GpuKernelAnatomyRiskStudio.java

This exact source runs in TeaVM. Form changes update its Java literals and reset manual edits.

import com.nablatensor.engine.ADouble;
import com.nablatensor.engine.AadRecorder;
import com.nablatensor.engine.AadTape;
import com.nablatensor.quant.EquityMarket;

public final class GpuKernelAnatomyRiskStudio {
  public static void main(String[] args) {
    double maturity = 1;
    int steps = 1;
    EquityMarket market = EquityMarket.of()
        .spot(100)
        .strike(100)
        .vol(.2)
        .rate(.03)
        .maturity(maturity)
        .build();
    AadTape tape = AadRecorder.record(rec -> {
      ADouble spot = rec.input("spot", market.spot()); ADouble strike = rec.input("strike", market.strike()); ADouble vol = rec.input("vol",
          market.vol()); ADouble rate = rec.input("rate", market.rate()); ADouble t = rec.input("maturity",
          market.maturity()); ADouble dt = t.div(steps); ADouble drift = rate.sub(vol.mul(vol)
          .mul(.5))
          .mul(dt); ADouble diffusion = vol.mul(dt.sqrt()); ADouble terminal = spot; for (int i = 0; i < steps; i++) terminal = terminal.mul(drift.add(diffusion.mul(rec.randn()))
          .exp()); rec.output(terminal.sub(strike)
          .max(0)
          .mul(rate.neg()
          .mul(t)
          .exp()));
    });
    int active = 0;
    for (int i = 0; i < tape.size(); i++) {
      if (tape.isActive(i)) active++;
      String op = tape.op(i)
          .name();
      int a = tape.argA(i), b = tape.argB(i);
      System.out.println("NODE|" + i + "|" + op + "|" + a + "|" + b + "|" + tape.isActive(i));
    }
    System.out.println("RESULT|" + tape.size() + "|" + active + "|" + steps + "|" + maturity);
  }
}
TeaVM compiles and runs the Java source above in this browser.
Implementation guide

Inspecting generated device code

Viewing a generated kernel makes the mapping from valuation graph to device arithmetic reviewable instead of opaque.

Core mechanism

The generated source declares buffers, maps tape nodes to local expressions, evaluates the primal graph and performs reverse accumulation according to backend conventions.

Practical workflow

Use source inspection to diagnose unsupported operations, excessive temporary state or unexpected precision paths; then validate with CPU parity and device-level profiling.

Key details

nablatensor-tensor/src/main/java/com/nablatensor/tensor/spi/GpuKernels.java is the single place where the CUDA and ROCm/HIP tensor backends get their device code. It is shared on purpose — NVRTC and HIPRTC both accept the same CUDA-C subset, so an op added here lands on both GPU flavours at once.

1. No structure. Finding reduce_axis_argmax means scrolling a 360-line string literal; there is no navigation, no outline, no per-kernel Javadoc. 2. Two sources of truth. The kernel names live both in the C text and in TENSOR_KERNEL_NAMES. A typo or a forgotten entry is only caught at module-load time on a machine with a GPU. 3. Zero tooling. No syntax highlighting, no clang-format, no per-kernel test, no way to say "this kernel needs block dim 16×16" other than remembering it at the call site.

What follows is seven independent approaches. They are not a sequence — pick one, or combine 1+4, or 2+6. Each section states the shape, a sketch, what it buys, what it costs, and a rough risk rating.

Shape. Delete the text block. Each kernel becomes a real file under nablatensor-tensor/src/main/resources/com/nablatensor/tensor/kernels/*.cu. Java keeps only the manifest and the concatenation.

Buys. Real .cu files get IDE syntax highlighting, clang-format, nvcc --dryrun compile-checks in CI without a GPU, and sane git blame/diffs. GpuKernels.java shrinks to about 60 lines of actual Java. The array becomes the manifest, and a missing file fails loudly at class-init instead of at cuModuleGetFunction.

Costs. Device code leaves the .java file, so a reader of GpuKernels no longer sees the kernels inline. Resource loading must survive shading/jlink/native-image packaging — a module-info/maven-shade misconfiguration turns a compile-time certainty into a runtime NullPointerException. Ordering is now the array's job, so #include-style dependencies between kernels have to be modelled explicitly.

Buys. This is the most "Java-like" answer with the least ceremony. Each kernel gets a name in the IDE outline, its own Javadoc, and — crucially — its launch geometry travels with it instead of being a magic 16 at the call site. The name/source consistency check moves into the record's compact constructor, so it runs at class-init on every JVM, GPU or not. Backends can iterate TENSOR_KERNELS and read blockDimX() rather than hardcoding.

Scope and review point

Readable generated source aids debugging, but semantic correctness still rests on the tape rules, code generator tests and backend conformance suite.