NanoComputeShader.kt

src/jvmMain/kotlin/afield/lang/nano/NanoComputeShader.kt · 200 lines

package afield.lang.nano

import java.nio.ByteBuffer
import java.nio.ByteOrder

The host side of nano_compute.glsl, the GLSL compute port of the Nano VM: the buffer layout a caller has to honour, the mapping struct, and the packing helpers for getting bytes in and out.

The shader runs one independent Nano VM instance per GPU invocation, decoding the same opcode map NanoVM.kt documents. Three of its four storage buffers hold bytes and are cut into non-overlapping per-instance slices; the fourth, at binding 0, holds one NanoMapping per dispatch id and says where that instance's three slices begin. Nothing in the tree dispatches it. SOURCE reads the GLSL off the classpath and the helpers below build the buffers, so what is missing is the Vulkan plumbing: pipeline, descriptor sets, submit and readback.

Buffer bindings

Byte-packing

All three byte buffers pack 4 bytes per uint: the byte at logical offset B lives in word B >> 2 at bit-shift (B & 3) << 3. A host has to honour that when uploading code and RAM and when reading OUT results back, which is what packBytes, packPrograms and unpackBytes are for.

Push constants (20 bytes, std430)

Mapping buffer layout (NanoMapping struct, 20 bytes each)

Minimal dispatch example (pseudocode)

val N         = 1024               // number of Nano instances
val ramSize   = 32                 // bytes per VM
val codeSize  = 64                 // bytes per program
val outSize   = 64                 // max OUT bytes per VM

val mappings  = NanoComputeShader.buildMappings(N, ramSize, codeSize, outSize)
// upload mappings, ramBytes (N*ramSize), codeBytes (N*codeSize), outBytes (N*outSize)
// bind as buffers 0..3
// dispatch ceil(N / 64) groups of 64

object NanoComputeShader

The GLSL source, the mapping-buffer builder and the byte-packing helpers, and nothing that talks to a device.

object NanoComputeShader {

val SOURCE: String by lazy

Source of the nano_compute.glsl compute shader.

    val SOURCE: String by lazy {
        NanoComputeShader::class.java.getResourceAsStream("/nano_compute.glsl")?.bufferedReader()?.readText()
            ?: error("nano_compute.glsl not found on classpath")
    }

Mapping buffer

data class NanoMapping

Kotlin mirror of the NanoMapping GPU struct (20 bytes, std430).

@param ramBase byte offset into the RAM buffer for this instance @param codeBase byte offset into the code buffer for this instance @param outputBase byte offset into the OUT buffer for this instance @param maxSteps per-instance step cap; 0 means use the push-constant default @param outputCapacity maximum OUT bytes; zero disables output, overflow halts the VM

    data class NanoMapping(
        val ramBase: Int,
        val codeBase: Int,
        val outputBase: Int,
        val maxSteps: Int = 0,
        val outputCapacity: Int = 0,
    )

fun buildMappings

Build a flat ByteBuffer containing one NanoMapping per instance, with each instance allocated contiguous, non-overlapping slices.

Slice layout (zero-based): - RAM : i * ramSize … i * ramSize + ramSize - 1 - code : i * codeSize … i * codeSize + codeSize - 1 - OUT : i * outSize … i * outSize + outSize - 1

@param numInstances number of VM instances @param ramSize RAM bytes per instance (must be power of 2, typically 32) @param codeSize code bytes per instance (typically 64) @param outSize OUT-stream bytes allocated per instance @param maxSteps per-instance step cap (0 → use push-constant default)

    fun buildMappings(
        numInstances: Int,
        ramSize: Int = NanoVM.RAM_SIZE,
        codeSize: Int = 64,
        outSize: Int = 64,
        maxSteps: Int = 0,
    ): ByteBuffer {
        require(numInstances > 0) { "numInstances must be > 0" }
        require(ramSize > 0 && (ramSize and (ramSize - 1)) == 0) {
            "ramSize must be a power of 2 (got $ramSize)"
        }
        require(codeSize > 0) { "codeSize must be > 0" }
        require(outSize >= 0) { "outSize must be >= 0" }

        require(maxSteps >= 0) { "maxSteps must be >= 0" }
        // Check offsets before allocation; signed Int offsets cannot represent larger slices.
        Math.multiplyExact(numInstances, ramSize)
        Math.multiplyExact(numInstances, codeSize)
        Math.multiplyExact(numInstances, outSize)

        // Each NanoMapping = 5 × uint32 = 20 bytes
        val buf = ByteBuffer.allocate(Math.multiplyExact(numInstances, 20)).order(ByteOrder.LITTLE_ENDIAN)
        for (i in 0 until numInstances) {
            buf.putInt(i * ramSize) // ramBase
            buf.putInt(i * codeSize) // codeBase
            buf.putInt(i * outSize) // outputBase
            buf.putInt(maxSteps) // maxSteps
            buf.putInt(outSize) // outputCapacity
        }
        buf.flip()
        return buf
    }

Byte-packing helpers

fun packBytes(bytes: ByteArray): ByteBuffer

Pack bytes into a uint[]-compatible ByteBuffer using the shader's 4-bytes-per-uint little-endian layout.

The returned buffer is ready for upload to a Vulkan SSBO. Its capacity is ceil(bytes.size / 4) * 4 bytes.

    fun packBytes(bytes: ByteArray): ByteBuffer {
        val wordCount = (bytes.size + 3) / 4
        val buf = ByteBuffer.allocate(wordCount * 4).order(ByteOrder.LITTLE_ENDIAN)
        buf.put(bytes)
        // zero-pad to word boundary (ByteBuffer.allocate already zero-fills)
        buf.flip()
        return buf
    }

fun unpackBytes(buf: ByteBuffer, baseOffset: Int, length: Int): ByteArray

Unpack length bytes from a uint[]-packed ByteBuffer starting at byte offset baseOffset.

Use this to read back RAM or OUT results after a dispatch.

    fun unpackBytes(buf: ByteBuffer, baseOffset: Int, length: Int): ByteArray {
        require(baseOffset >= 0 && baseOffset + length <= buf.limit()) {
            "range [$baseOffset, ${baseOffset + length}) out of buffer bounds (limit=${buf.limit()})"
        }
        val out = ByteArray(length)
        val view = buf.duplicate().order(ByteOrder.LITTLE_ENDIAN)
        view.position(baseOffset)
        view.get(out)
        return out
    }

fun packPrograms

Pack programs (each up to codeSize bytes) into a single flat ByteBuffer ready for binding as the code SSBO.

Programs shorter than codeSize are zero-padded (byte 0x00 = NOP).

    fun packPrograms(
        programs: List<ByteArray>,
        codeSize: Int = 64,
    ): ByteBuffer {
        val wordCount = (programs.size * codeSize + 3) / 4
        val buf = ByteBuffer.allocate(wordCount * 4).order(ByteOrder.LITTLE_ENDIAN)
        for (prog in programs) {
            val slice = ByteArray(codeSize)
            prog.copyInto(slice, endIndex = minOf(prog.size, codeSize))
            buf.put(slice)
        }
        buf.flip()
        return buf
    }

Assemble convenience

fun assemble(source: String): ByteArray = nanoAssemble(source)

Assemble source with NanoAssembler and return the raw bytes, or throw NanoAssemblyException on error. Useful for building code buffers inline.

Delegates to the top-level nanoAssemble() function in NanoInterface.kt (same package).