Skip to content

SplatShader.kt

src/jvmMain/kotlin/afield/graphics/scene/fabricated/SplatShader.kt · 352 lines

@file:Suppress("MemberVisibilityCanBePrivate")

package afield.graphics.scene.fabricated

import afield.document.context.fromContext
import afield.document.graph.Box
import afield.graphics.render.ComputeMesh
import afield.graphics.render.ComputeModulePointer
import afield.graphics.vulkan.Renderer
import afield.graphics.vulkan.Target
import afield.graphics.vulkan.buffers.FrameBufferGroup
import afield.graphics.vulkan.buffers.defaultBuffer
import afield.graphics.vulkan.targets.PointAccumulation
import afield.graphics.vulkan.targets.PointBlend
import afield.graphics.vulkan.targets.PointToneMap
import afield.graphics.render.Module
import afield.graphics.render.pointer
import afield.math.linalg.Matrix4f
import afield.math.linalg.Rectanglef
import afield.math.linalg.Vector3f
import afield.ui.output.controlGroup
import afield.ui.output.displayOpen

The splat-shader box: a cloud of Gaussian splats uploaded once, a few lines of GLSL that rewrite every splat each frame, and Field's Vulkan splat renderer drawing the result.

SplatShader does the work. It uploads the parametric splats, binds them to the box's compute module as splatsIn and the renderer's cloud as splatsOut, dispatches the insert once per splat before the sort runs, and installs the cloud in an accumulation target, so a trail is a setting rather than a second code path. SplatShaderSettings is everything a viewer can change while it runs, from the camera and the turn rate to the opacity and gamma the splats are blended with, and showControls puts the demo's sliders in the box's controls output.

The clouds come from SplatSet and ProceduralSplats; the sort, the packed covariance form and the splat shaders themselves are GaussianSplatRenderer's.

class SplatShader

A Gaussian-splat cloud whose splats are rewritten every frame by a few lines of GLSL in a box: the madp splat-shader demo's "insert", on Field's Vulkan splat renderer.

The box's compute tab is the insert between two includes:

#version 450
#include "splat/insert"

float r = length(pos.xz);
pos.y += sin(r * 18.0 - t * 3.5) * 0.045;
rgba.rgb = mix(rgba.rgb, vec3(0.08, 0.85, 1.0), 0.5);

#include "splat/end"

and its code hands that module and a cloud to this class:

var module = computeModuleHere {}
SplatShader(module, ProceduralSplats.torusBraid()) {
    autoRotate = 0.32
}

Between the includes, pos, scales, rgba and quat are the splat's centre, radii, colour and rotation, yours to change; index (which splat) and t (seconds since the box ran) are there to read, with pi and the helpers sat, rot2, hash31 and noise31. splat/end clamps rgba to 0..1 and every scale to at least 1e-4, as the demo does, then builds the splat's covariance.

Every frame, the insert runs once per splat on the GPU, writing the cloud the renderer draws; then the renderer sorts those splats back to front and draws them. So a splat the insert moves is sorted where it went, which is what keeps a dissolve or a gravity well compositing in the right order. (Edit the splat vertex shader instead, as gaussianSplatModuleHere lets you, and the sort still sees the splats where they were uploaded.)

The cloud is framed as the demo frames it: centred on its median, scaled so nine splats in ten lie within 1.08 of the origin (SplatShaderSettings.size times that), turned about +y at SplatShaderSettings.autoRotate, and seen from (0, 0.38, 2.65) with a 54° lens.

The splats are drawn into a PointAccumulation that keeps 1 - fade of the previous frame, which is the demo's trail mode: with SplatShaderSettings.trail on, each frame is laid over a fading copy of the last, so moving splats leave streaks. With it off the accumulation is cleared every frame and the picture is the plain cloud.

Every SplatShader installs under the same name unless SplatShaderSettings.name says otherwise, so launching one replaces the one showing.

class SplatShader(
    /** The compute module holding the insert, usually `computeModuleHere {}`. */
    val module: ComputeModulePointer,
    /** The splats the insert starts from each frame. */
    val splats: SplatSet,
    /** The renderer to draw on; by default the one this box is running in. */
    val renderer: Renderer = fromContext(),
    configure: SplatShaderSettings.() -> Unit = {}
) {

val settings = SplatShaderSettings().apply(configure)

The settings after configure ran; the framing, rotation and splat look can change live.

val source: FrameBufferGroup

The parametric splats, as uploaded: what the insert reads.

val cloud: GaussianSplatRenderer.Cloud

The renderer's cloud: what the insert writes, and what is sorted and drawn.

val trail: PointAccumulation

What the cloud is drawn into, and what leaves the trail.

val mesh: ComputeMesh

The dispatch in module that runs the insert, one invocation per splat.

val fit: SplatSet.Fit = splats.fit()

The SplatSet.fit the framing uses.

var t: Double = 0.0

Seconds since this was made: the insert's t. Set it to scrub.

var angle: Double = 0.0

The cloud's current turn about +y, in radians.

var playing: Boolean = true

When false, t and the rotation stand still and the insert sees the same t each frame.

    private var lastNanos = 0L

init

body — 82 lines, from line 127 
    init {
        require(splats.count > 0) { "a SplatShader needs at least one splat" }
        val count = splats.count
        // Nothing frees a GPU buffer or a cloud once made, so a re-run of the same size takes the
        // last one's back rather than allocating its own.
        trail = synchronized(kept) {
            trails.getOrPut(renderer) { HashMap() }.getOrPut(settings.name) {
                PointAccumulation(renderer) {
                    name = "${settings.name}-trail"
                    blend = PointBlend.ALPHA
                    toneMap = PointToneMap.LINEAR
                    exposure = 1.0
                    fade = 1.0
                }
            }
        }
        val shape = Shape(settings.name, count, settings.sortMode)
        val previous = synchronized(kept) { kept.getOrPut(renderer) { HashMap() }[shape] }
        if (previous != null) {
            source = previous.source
            cloud = previous.cloud
            cloud.beforeSort.clear()
        } else {
            source = defaultBuffer(name = "${settings.name}-source") {
                size = (count * 4).vec4s
                perFrame = false
            }.toBuffer(renderer)
            val (min, max) = splats.bounds()
            val config = GaussianSplatRenderer.Config(
                sortMode = settings.sortMode,
                depthKey = GaussianSplatRenderer.DepthKey.ViewDepth,
                // the insert can move splats under a still camera
                sortOnlyWhenViewChanges = false,
                maxScreenRadius = 256f
            )
            // the built-in splat shaders, with the trail's target format and the renderer's raster state
            val renderModule = Module(
                GaussianSplatRenderer.configureModule(
                    trail.pointParameters("${settings.name}-splats") {
                        vertex = GaussianSplatRenderer.splatVertexShader
                        fragment = GaussianSplatRenderer.splatFragmentShader
                    },
                    config = config,
                    name = "${settings.name}-splats-$count"
                )
            ).pointer()
            cloud = GaussianSplatRenderer.Cloud(
                renderer = renderer,
                camera = afield.graphics.vulkan.util.Camera(renderer),
                table = null,
                // Zeros: opacity 0 hides every splat until the insert has written it.
                packed = GaussianSplatRenderer.PackedSplats(FloatArray(count * SplatSet.STRIDE), count, min, max),
                config = config,
                sourceName = settings.name,
                capacity = count,
                suppliedRenderModule = renderModule
            )
            synchronized(kept) { kept.getOrPut(renderer) { HashMap() }[shape] = Resources(source, cloud) }
        }
        source.write { put(splats.floats.copyOf(count * SplatSet.STRIDE), padding = false) }

        mesh = ComputeMesh(renderer, module)
        mesh.push.bind("splatsIn") { source }
        mesh.push.bind("splatsOut") { requireNotNull(cloud.splatData.result) }
        mesh.push.bind("splatCount") { count }
        mesh.push.bind("t") { t.toFloat() }
        mesh.x = (count + 63) / 64
        mesh.y = 1
        mesh.z = 1
        module += "${settings.name}-insert" to mesh

        cloud.beforeSort["insert"] = Target { command, _ ->
            tick()
            module.render(command, Rectanglef())
        }
        cloud.bindModel { model() }
        applyLook()
        // one cloud in the trail at a time, under this name
        trail.add(settings.name, cloud.renderModule)
        trail.clear()
        trail.installInto(renderer.renderList)
    }

fun model(): Matrix4f

The model matrix this frame: the fit, the size, then the turn about +y.

    fun model(): Matrix4f {
        val base = if (settings.fit) fit.model(1.08f * settings.size.toFloat())
        else Matrix4f().scale(settings.size.toFloat())
        return Matrix4f().rotateY(angle.toFloat()).mul(base)
    }

fun applyLook()

Push settings' camera and splat look onto the cloud; called once a frame, so edits show.

    fun applyLook() {
        val camera = cloud.camera
        camera.at = Vector3f(settings.cameraAt)
        camera.target = Vector3f(settings.cameraTarget)
        camera.up = Vector3f(0f, 1f, 0f)
        camera.p.fov = Math.toRadians(settings.fovDegrees)
        camera.p.zNear = 0.01
        camera.p.zFar = 100.0
        trail.settings.fade = if (settings.trail) settings.fade.coerceIn(0.0, 1.0) else 1.0
        cloud.sizeScale = settings.splatSize
        cloud.alphaScale = settings.alphaScale
        cloud.colorGamma = settings.colorGamma
    }

    private fun tick() {
        val now = System.nanoTime()
        val dt = if (lastNanos == 0L) 0.0 else ((now - lastNanos) / 1e9).coerceIn(0.0, 0.25)
        lastNanos = now
        if (playing) {
            t += dt
            // the demo's slider adds speed × 0.01 radians per frame at 60 frames a second
            angle += settings.autoRotate * 0.6 * dt
        }
        applyLook()
    }

fun remove()

Take the cloud and its trail off the renderer. Launching another SplatShader with the same name does this too.

    fun remove() {
        trail.target.modules.remove(settings.name)
        renderer.renderList.preTargets.remove("${trail.settings.name}-accumulation")
        renderer.renderList.modules.remove("${trail.settings.name}-display")
    }

    private data class Shape(val name: String, val count: Int, val sortMode: GaussianSplatRenderer.SortMode)

    private class Resources(val source: FrameBufferGroup, val cloud: GaussianSplatRenderer.Cloud)

companion object

    companion object {

The buffers and cloud each shape was last built with, per renderer.

        private val kept = java.util.WeakHashMap<Renderer, HashMap<Shape, Resources>>()

One trail per name per renderer: its target is window-sized, and nothing frees it.

        private val trails = java.util.WeakHashMap<Renderer, HashMap<String, PointAccumulation>>()

        private val includes = setOf("insert", "end")

fun hasInclude(name: String): Boolean = name.removeSuffix(".glsl") in includes

True for the names #include "splat/<name>" resolves: insert and end.

fun include(name: String): String

The packaged source of #include "splat/<name>".

        fun include(name: String): String {
            val bare = name.removeSuffix(".glsl")
            require(bare in includes) { "Unknown splat shader include: $name (there are ${includes.joinToString()})" }
            return checkNotNull(SplatShader::class.java.getResourceAsStream("/afield/four/splat/$bare.glsl")) {
                "Missing packaged splat shader include: $bare"
            }.bufferedReader().use { it.readText() }
        }

fun computeSource(insert: String): String

A complete compute shader for insert — what a box's compute tab expands to — for building a SplatShader's module outside a box, as the GPU smoke does.

        fun computeSource(insert: String): String =
            "#version 450\n" + include("insert") + "\n" + insert + "\n" + include("end")
    }
}

class SplatShaderSettings

What a SplatShader can be told; see each property.

class SplatShaderSettings {

var name: String = "splat-shader"

The renderer module name. Two SplatShaders with one name replace each other.

var autoRotate: Double = 0.32

Turn about +y, in the demo's slider units: 0.32 is its default, 0 stands still.

var size: Double = 1.0

The demo's size slider: a multiplier on the fitted size.

var trail: Boolean = false

The demo's trail mode: lay each frame over a fading copy of the last, so moving splats leave streaks. fade is how much of the old picture goes each frame.

var fade: Double = 0.12

With trail on, the part of the previous frame that fades away each frame: the demo's 0.12.

var fit: Boolean = true

Frame the cloud by its median and 90th-percentile radius; false draws it as it is.

var splatSize: Float = (1.0 / kotlin.math.sqrt(2.2)).toFloat()

Multiplies every splat's screen size. The default, 1/√2.2, is what makes a splat look the size it does in the demo. Spark blends in display space, so a Gaussian's faint tails stay dark; Field blends in linear light and the window encodes the result to sRGB, which lifts those tails, and at the same radius the visible dot comes out about twice as wide. For a splat over black, encode(c · a') = c · a wants a' ≈ a^2.2, which is the same Gaussian narrowed by √2.2. Set 1 for the splats' true size.

var alphaScale: Float = 1f

Multiplies every splat's opacity.

var colorGamma: Float = 2.2f

See GaussianSplatRenderer.Config.colorGamma; 2.2 is right for Field's window.

var cameraAt: Vector3f = Vector3f(0f, 0.38f, 2.65f)

var cameraTarget: Vector3f = Vector3f(0f, 0f, 0f)

var fovDegrees: Double = 54.0

Vertical field of view.

var sortMode: GaussianSplatRenderer.SortMode = GaussianSplatRenderer.SortMode.ExactRadix

GaussianSplatRenderer.SortMode.ExactRadix by default: the keys are view-depth float bits.

fun Box.showControls(shader: SplatShader, channel: String = "controls")

The demo's viewer toolbar for shader, in this box's controls output: the auto-rotate speed, the size, a splat-size multiplier, play/pause for t and the turn, and a button that starts t again from zero. The panel opens, as the demo's toolbar is always in view.

fun Box.showControls(shader: SplatShader, channel: String = "controls") {
    val s = shader.settings
    controlGroup(channel) {
        slider("turn", min = 0.0, max = 1.5, step = 0.01, initialValue = s.autoRotate) { s.autoRotate = it }
        slider("size", min = 0.35, max = 2.5, step = 0.01, initialValue = s.size) { s.size = it }
        slider("splat size", min = 0.25, max = 4.0, step = 0.01, initialValue = s.splatSize.toDouble()) { s.splatSize = it.toFloat() }
        checkbox("play", initialValue = shader.playing) { shader.playing = it }
        checkbox("trail", initialValue = s.trail) { s.trail = it }
        slider("fade", min = 0.0, max = 1.0, step = 0.01, initialValue = s.fade) { s.fade = it }
        button("t = 0") { down -> if (down) shader.t = 0.0 }
    }
    displayOpen(channel)
}