@multiplekex/shallot 0.2.5 → 0.3.0

package/src/standard/raster/shadow.ts

@@ -0,0 +1,408 @@
+import type { ComputeNode, ExecutionContext } from "../compute";
+import {
+    perspective,
+    multiply,
+    invert,
+    lookAtMatrix,
+    orthographicBounds,
+    extractFrustumCorners,
+    extractFrustumPlanes,
+    invertMatrix,
+} from "../../core/math";
+import { SHADOW_STRUCT_WGSL } from "../render/surface/structs";
+import { VERTEX_BUFFERS, drawSlots } from "./forward";
+import type { Batch } from "./batch";
+
+const SHADOW_CASCADE_COUNT = 4;
+const SHADOW_ATLAS_SIZE = 2048;
+const SHADOW_CASCADE_SIZE = SHADOW_ATLAS_SIZE / 2;
+const SHADOW_BUFFER_SIZE = 288;
+
+export function computeCascadeSplits(
+    near: number,
+    far: number,
+    cascadeCount: number,
+    lambda = 0.5
+): Float32Array {
+    const splits = new Float32Array(cascadeCount);
+    const ratio = far / near;
+
+    for (let i = 0; i < cascadeCount; i++) {
+        const p = (i + 1) / cascadeCount;
+        const log = near * Math.pow(ratio, p);
+        const uniform = near + (far - near) * p;
+        splits[i] = lambda * log + (1 - lambda) * uniform;
+    }
+
+    return splits;
+}
+
+interface CascadeData {
+    viewProj: Float32Array;
+    frustumPlanes: Float32Array;
+    texelSize: number;
+}
+
+export function computeCascadeMatrix(
+    cameraWorld: Float32Array,
+    fov: number,
+    aspect: number,
+    nearSplit: number,
+    farSplit: number,
+    lightDir: [number, number, number],
+    shadowMapSize: number
+): CascadeData {
+    const proj = perspective(fov, aspect, nearSplit, farSplit);
+    const view = invert(cameraWorld);
+    const viewProj = multiply(proj, view);
+    const invViewProj = invertMatrix(viewProj);
+
+    const corners = extractFrustumCorners(invViewProj, 0, 1);
+
+    let centerX = 0,
+        centerY = 0,
+        centerZ = 0;
+    for (let i = 0; i < 8; i++) {
+        centerX += corners[i * 3];
+        centerY += corners[i * 3 + 1];
+        centerZ += corners[i * 3 + 2];
+    }
+    centerX /= 8;
+    centerY /= 8;
+    centerZ /= 8;
+
+    const [lightDirX, lightDirY, lightDirZ] = lightDir;
+    const len = Math.sqrt(lightDirX * lightDirX + lightDirY * lightDirY + lightDirZ * lightDirZ);
+    const normLightX = lightDirX / len;
+    const normLightY = lightDirY / len;
+    const normLightZ = lightDirZ / len;
+
+    let maxRadius = 0;
+    for (let i = 0; i < 8; i++) {
+        const dx = corners[i * 3] - centerX;
+        const dy = corners[i * 3 + 1] - centerY;
+        const dz = corners[i * 3 + 2] - centerZ;
+        const dist = Math.sqrt(dx * dx + dy * dy + dz * dz);
+        maxRadius = Math.max(maxRadius, dist);
+    }
+
+    const shadowDistance = maxRadius * 2;
+    const eyeX = centerX - normLightX * shadowDistance;
+    const eyeY = centerY - normLightY * shadowDistance;
+    const eyeZ = centerZ - normLightZ * shadowDistance;
+
+    const lightView = lookAtMatrix(eyeX, eyeY, eyeZ, centerX, centerY, centerZ);
+
+    let minX = Infinity,
+        maxX = -Infinity;
+    let minY = Infinity,
+        maxY = -Infinity;
+    let minZ = Infinity,
+        maxZ = -Infinity;
+
+    for (let i = 0; i < 8; i++) {
+        const wx = corners[i * 3];
+        const wy = corners[i * 3 + 1];
+        const wz = corners[i * 3 + 2];
+
+        const lx = lightView[0] * wx + lightView[4] * wy + lightView[8] * wz + lightView[12];
+        const ly = lightView[1] * wx + lightView[5] * wy + lightView[9] * wz + lightView[13];
+        const lz = lightView[2] * wx + lightView[6] * wy + lightView[10] * wz + lightView[14];
+
+        minX = Math.min(minX, lx);
+        maxX = Math.max(maxX, lx);
+        minY = Math.min(minY, ly);
+        maxY = Math.max(maxY, ly);
+        minZ = Math.min(minZ, lz);
+        maxZ = Math.max(maxZ, lz);
+    }
+
+    const extendBack = shadowDistance;
+    minZ -= extendBack;
+    maxZ += maxRadius;
+
+    const texelSize = (maxX - minX) / shadowMapSize;
+    minX = Math.floor(minX / texelSize) * texelSize;
+    maxX = Math.ceil(maxX / texelSize) * texelSize;
+    minY = Math.floor(minY / texelSize) * texelSize;
+    maxY = Math.ceil(maxY / texelSize) * texelSize;
+
+    const lightProj = orthographicBounds(minX, maxX, minY, maxY, -maxZ, -minZ);
+    const cascadeViewProj = multiply(lightProj, lightView);
+    const frustumPlanes = extractFrustumPlanes(cascadeViewProj);
+
+    return { viewProj: cascadeViewProj, frustumPlanes, texelSize };
+}
+
+const shadowBuffer = new ArrayBuffer(SHADOW_BUFFER_SIZE);
+const shadowF32 = new Float32Array(shadowBuffer);
+
+function uploadShadowData(
+    device: GPUDevice,
+    buffer: GPUBuffer,
+    cascades: CascadeData[],
+    splits: Float32Array
+): void {
+    for (let i = 0; i < SHADOW_CASCADE_COUNT; i++) {
+        shadowF32.set(cascades[i].viewProj, i * 16);
+    }
+    shadowF32[64] = splits[0];
+    shadowF32[65] = splits[1];
+    shadowF32[66] = splits[2];
+    shadowF32[67] = splits[3];
+
+    shadowF32[68] = cascades[0].texelSize;
+    shadowF32[69] = cascades[1].texelSize;
+    shadowF32[70] = cascades[2].texelSize;
+    shadowF32[71] = cascades[3].texelSize;
+
+    device.queue.writeBuffer(buffer, 0, shadowBuffer);
+}
+
+export function createShadowBuffer(device: GPUDevice): GPUBuffer {
+    return device.createBuffer({
+        label: "shadow",
+        size: SHADOW_BUFFER_SIZE,
+        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
+}
+
+export function createShadowAtlas(device: GPUDevice): GPUTexture {
+    return device.createTexture({
+        label: "shadow-atlas",
+        size: [SHADOW_ATLAS_SIZE, SHADOW_ATLAS_SIZE, 1],
+        format: "depth32float",
+        usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.TEXTURE_BINDING,
+    });
+}
+
+interface ShadowUploadConfig {
+    shadow: GPUBuffer;
+    getCameraData: () => {
+        world: Float32Array;
+        fov: number;
+        near: number;
+        far: number;
+        width: number;
+        height: number;
+    } | null;
+    getLightDir: () => [number, number, number];
+    shadowsEnabled: () => boolean;
+    getShadowDistance: () => number;
+}
+
+export function createShadowUploadNode(config: ShadowUploadConfig): ComputeNode {
+    let cascades: CascadeData[] = [];
+    let splits: Float32Array = new Float32Array(SHADOW_CASCADE_COUNT);
+
+    return {
+        id: "shadow-cascade-upload",
+        inputs: [{ id: "data", access: "read" }],
+        outputs: [{ id: "shadow-cascades", access: "write" }],
+
+        execute(ctx: ExecutionContext) {
+            if (!config.shadowsEnabled()) return;
+
+            const camera = config.getCameraData();
+            if (!camera) return;
+
+            const { world, fov, near, far, width, height } = camera;
+            const aspect = width / height;
+            const lightDir = config.getLightDir();
+
+            const shadowDistance = config.getShadowDistance();
+            const effectiveFar = Math.min(far, shadowDistance);
+
+            splits = computeCascadeSplits(near, effectiveFar, SHADOW_CASCADE_COUNT);
+            cascades = [];
+
+            let prevSplit = near;
+            for (let i = 0; i < SHADOW_CASCADE_COUNT; i++) {
+                const cascade = computeCascadeMatrix(
+                    world,
+                    fov,
+                    aspect,
+                    prevSplit,
+                    splits[i],
+                    lightDir,
+                    SHADOW_CASCADE_SIZE
+                );
+                cascades.push(cascade);
+                prevSplit = splits[i];
+            }
+
+            uploadShadowData(ctx.device, config.shadow, cascades, splits);
+        },
+    };
+}
+
+interface ShadowForwardGPU {
+    pipeline: GPURenderPipeline;
+    cascadeIndexBuffers: GPUBuffer[];
+    cascadeBindGroups: GPUBindGroup[];
+}
+
+const shadowDepthShader = /* wgsl */ `
+struct VertexInput {
+    @location(0) position: vec3<f32>,
+    @location(1) normal: vec3<f32>,
+    @builtin(instance_index) instance: u32,
+}
+
+struct VertexOutput {
+    @builtin(position) position: vec4<f32>,
+}
+
+${SHADOW_STRUCT_WGSL}
+
+@group(0) @binding(0) var<uniform> shadow: Shadow;
+@group(0) @binding(1) var<storage, read> entityIds: array<u32>;
+@group(0) @binding(2) var<storage, read> matrices: array<mat4x4<f32>>;
+@group(0) @binding(3) var<storage, read> sizes: array<vec4<f32>>;
+@group(0) @binding(4) var<uniform> cascadeIndex: u32;
+
+fn getCascadeViewProj(cascade: u32) -> mat4x4<f32> {
+    switch cascade {
+        case 0u: { return shadow.cascade0ViewProj; }
+        case 1u: { return shadow.cascade1ViewProj; }
+        case 2u: { return shadow.cascade2ViewProj; }
+        default: { return shadow.cascade3ViewProj; }
+    }
+}
+
+@vertex
+fn vs(input: VertexInput) -> VertexOutput {
+    let eid = entityIds[input.instance];
+    let world = matrices[eid];
+    let scaledPos = input.position * sizes[eid].xyz;
+    let worldPos = (world * vec4<f32>(scaledPos, 1.0)).xyz;
+    let viewProj = getCascadeViewProj(cascadeIndex);
+
+    var output: VertexOutput;
+    output.position = viewProj * vec4<f32>(worldPos, 1.0);
+    return output;
+}
+
+@fragment
+fn fs() {}
+`;
+
+async function prepareShadowForwardGPU(
+    device: GPUDevice,
+    config: ShadowForwardConfig
+): Promise<ShadowForwardGPU> {
+    const module = device.createShaderModule({ code: shadowDepthShader });
+    const pipeline = await device.createRenderPipelineAsync({
+        layout: "auto",
+        vertex: { module, entryPoint: "vs", buffers: VERTEX_BUFFERS },
+        fragment: { module, entryPoint: "fs", targets: [] },
+        depthStencil: {
+            format: "depth32float",
+            depthWriteEnabled: true,
+            depthCompare: "less",
+        },
+        primitive: { topology: "triangle-list", cullMode: "back" },
+    });
+
+    const cascadeIndexBuffers: GPUBuffer[] = [];
+    for (let i = 0; i < 4; i++) {
+        const buf = device.createBuffer({
+            label: `cascade-index-${i}`,
+            size: 4,
+            usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+        });
+        device.queue.writeBuffer(buf, 0, new Uint32Array([i]));
+        cascadeIndexBuffers.push(buf);
+    }
+
+    const layout = pipeline.getBindGroupLayout(0);
+    const cascadeBindGroups = cascadeIndexBuffers.map((indexBuf) =>
+        device.createBindGroup({
+            layout,
+            entries: [
+                { binding: 0, resource: { buffer: config.shadow } },
+                { binding: 1, resource: { buffer: config.batching.entityIds } },
+                { binding: 2, resource: { buffer: config.matrices } },
+                { binding: 3, resource: { buffer: config.sizes } },
+                { binding: 4, resource: { buffer: indexBuf } },
+            ],
+        })
+    );
+
+    return { pipeline, cascadeIndexBuffers, cascadeBindGroups };
+}
+
+interface ShadowForwardConfig {
+    shadow: GPUBuffer;
+    matrices: GPUBuffer;
+    sizes: GPUBuffer;
+    batching: Batch;
+    atlas: GPUTexture;
+    shadowsEnabled: () => boolean;
+}
+
+export function createShadowForwardNode(config: ShadowForwardConfig): ComputeNode[] {
+    let gpu: ShadowForwardGPU | null = null;
+
+    const nodes: ComputeNode[] = [];
+
+    for (let cascade = 0; cascade < 4; cascade++) {
+        const inputId = cascade === 0 ? "shadow-cascades" : `shadow-cascade-${cascade - 1}`;
+
+        nodes.push({
+            id: `shadow-render-${cascade}`,
+            inputs: [
+                { id: inputId, access: "read" },
+                { id: "batched", access: "read" },
+            ],
+            outputs: [{ id: `shadow-cascade-${cascade}`, access: "write" }],
+
+            async prepare(device: GPUDevice) {
+                if (!gpu) {
+                    gpu = await prepareShadowForwardGPU(device, config);
+                }
+            },
+
+            execute(ctx: ExecutionContext) {
+                if (!gpu || !config.shadowsEnabled()) return;
+
+                const buffers = config.batching.buffers;
+                const offsetX = (cascade % 2) * SHADOW_CASCADE_SIZE;
+                const offsetY = Math.floor(cascade / 2) * SHADOW_CASCADE_SIZE;
+
+                const pass = ctx.encoder.beginRenderPass({
+                    colorAttachments: [],
+                    depthStencilAttachment: {
+                        view: config.atlas.createView(),
+                        depthClearValue: 1.0,
+                        depthLoadOp: cascade === 0 ? "clear" : "load",
+                        depthStoreOp: "store",
+                    },
+                });
+
+                pass.setViewport(offsetX, offsetY, SHADOW_CASCADE_SIZE, SHADOW_CASCADE_SIZE, 0, 1);
+                pass.setScissorRect(offsetX, offsetY, SHADOW_CASCADE_SIZE, SHADOW_CASCADE_SIZE);
+                pass.setPipeline(gpu.pipeline);
+                pass.setBindGroup(0, gpu.cascadeBindGroups[cascade]);
+
+                drawSlots(pass, config.batching.opaque, buffers, ctx.device);
+
+                pass.end();
+            },
+        });
+    }
+
+    const doneNode: ComputeNode = {
+        id: "shadow-done",
+        inputs: [
+            { id: "shadow-cascade-0", access: "read" },
+            { id: "shadow-cascade-1", access: "read" },
+            { id: "shadow-cascade-2", access: "read" },
+            { id: "shadow-cascade-3", access: "read" },
+        ],
+        outputs: [{ id: "shadow-atlas", access: "write" }],
+        execute() {},
+    };
+
+    return [...nodes, doneNode];
+}