GpuKernels

The compute-kernel source shared by the CUDA and ROCm/HIP tensor backends.

The two backends run byte-for-byte the same kernels: the bodies are plain CUDA C that NVRTC compiles to PTX and HIPRTC accepts unchanged and compiles to a GCN code object. Keeping the source here rather than copied into each backend means an op added or a bug fixed lands on both GPU flavours at once. (The Vulkan backend is unrelated — it runs GLSL compute shaders through shaderc.)

This class is a façade. The kernels themselves live one family per file (ElementwiseKernels, MatmulKernels, RandomKernels, ReductionKernels, ConvKernels), each as a GpuKernel that parses its own entry-point name and parameter types out of its source — so TENSOR_KERNEL_NAMES is derived, never hand-maintained, and cannot drift from the code it names.

If a kernel ever needs architecture-specific tuning that CUDA C's shared subset cannot express (wavefront-64 vs warp-32 reductions, differing block sizes, __syncwarp), split that one kernel back out into the backend that needs the variant; everything that stays identical stays here.

View source ↗Full module reference on javadoc.io ↗

Fields

static final String FUSED_KERNEL_NAME

Entry-point name of the kernel fusedSource generates.

Methods

static List<GpuKernel> TENSOR_KERNELS = Stream. of(ElementwiseKernels. KERNELS, MatmulKernels. KERNELS, RandomKernels. KERNELS, ReductionKernels. KERNELS, ConvKernels.KERNELS) .flatMap(List::stream) .toList( ConvKernels.KERNELS) .flatMap(List::stream) .toList()

Every shared kernel, by family, in the order they are emitted and loaded.

static String[] TENSOR_KERNEL_NAMES = BY_NAME. keySet().toArray(String[]:: new)

Kernels in TENSOR_SOURCE, in a stable order for module loading.

static String TENSOR_SOURCE = DevicePrelude.SOURCE + TENSOR_KERNELS. stream().map(GpuKernel::source).collect(Collectors.joining() ).map(GpuKernel::source).collect(Collectors.joining())

Every static tensor kernel, as one CUDA-C translation unit.

static GpuKernel kernel(String name)

Looks up a shared kernel by the entry-point name parsed out of its source.

static String fusedSource(Expr expr, int numInputs)

Emits a whole elementwise expression chain as one fused_kernel translation unit, so a backend can run the chain in a single launch with no per-op intermediate buffers. The text is CUDA C; HIPRTC takes it unchanged.