glre 0.36.0 → 0.37.0

package/src/webgl.ts

@@ -1,6 +1,15 @@
 import { nested as cached } from 'reev'
 import { loadingImage } from './utils/helpers'
-import { createAttrib, createProgram, createStorage, createTexture, createUniform } from './utils/program'
+import {
+        cleanStorage,
+        createAttachment,
+        createAttrib,
+        createProgram,
+        createStorage,
+        createTexture,
+        createUniform,
+} from './utils/program'
+import { compute, fragment, vertex } from './node'
 import type { GL, WebGLState } from './types'
 
 const vert = /* cpp */ `
@@ -11,100 +20,103 @@ void main() {
         gl_Position = vec4(x, y, 0.0, 1.0);
 }`.trim()
 
-export const webgl = async (gl: GL) => {
-        const c = gl.el!.getContext('webgl2')!
+const computeProgram = (gl: GL, c: WebGL2RenderingContext) => {
+        if (!gl.cs) return null // ignore if no compute shader
         c.getExtension('EXT_color_buffer_float')
-        const pg1 = createProgram(c, gl.vs, gl.fs, gl)!
-        const pg2 = createProgram(c, vert, gl.cs, gl)!
-        c.useProgram(pg1)
 
         let activeUnit = 0 // for texture units
         let currentNum = 0 // for storage buffers
 
-        const attribs = cached((key) => c.getAttribLocation(pg1, key))
-        const uniforms1 = cached((key) => c.getUniformLocation(pg1, key))
-        const uniforms2 = cached((key) => c.getUniformLocation(pg2, key))
-        const textures = cached(() => activeUnit++)
-        const storages = cached(() => {
-                const unit = activeUnit++
-                const a = { texture: c.createTexture(), buffer: c.createFramebuffer() }
-                const b = { texture: c.createTexture(), buffer: c.createFramebuffer() }
-                return { a, b, unit, width: 0, height: 0 }
+        const units = cached(() => activeUnit++)
+        const config = { isWebGL: true, gl, units }
+
+        const pg = createProgram(c, compute(gl.cs, config), vert, gl)!
+        const size = Math.ceil(Math.sqrt(gl.particles))
+
+        const uniforms = cached((key) => c.getUniformLocation(pg, key)!)
+        const storages = cached((key) => {
+                const array = new Float32Array(size * size * 4) // RGBA texture data
+                const ping = { texture: c.createTexture(), buffer: c.createFramebuffer() }
+                const pong = { texture: c.createTexture(), buffer: c.createFramebuffer() }
+                return { ping, pong, array, loc: uniforms(key), unit: units(key) }
         })
 
+        const _uniform = (key: string, value: number | number[]) => {
+                c.useProgram(pg)
+                createUniform(c, uniforms(key), value)
+        }
+
+        const _storage = (key: string, value: number[]) => {
+                const { ping, pong, unit, array } = storages(key)
+                createStorage(c, value, size, ping, pong, unit, array)
+        }
+
+        const clean = () => {
+                c.deleteProgram(pg)
+                cleanStorage(c, storages.map.values())
+        }
+
+        const render = () => {
+                c.useProgram(pg)
+                const attachments = storages.map.values().map(({ ping, pong, loc, unit }, index) => {
+                        const [i, o] = currentNum % 2 ? [ping, pong] : [pong, ping]
+                        return createAttachment(c, i, o, loc, unit, index)
+                })
+                c.drawBuffers(attachments)
+                c.drawArrays(c.TRIANGLES, 0, 3)
+                c.bindFramebuffer(c.FRAMEBUFFER, null)
+                currentNum++
+        }
+
+        return { render, clean, _uniform, _storage, storages }
+}
+
+export const webgl = async (gl: GL) => {
+        const config = { isWebGL: true, gl }
+        const c = gl.el!.getContext('webgl2')!
+        const cp = computeProgram(gl, c)
+        const pg = createProgram(c, fragment(gl.fs, config), vertex(gl.vs, config), gl)!
+        c.useProgram(pg)
+
+        let activeUnit = 0 // for texture units
+
+        const units = cached(() => activeUnit++)
+        const attribs = cached((key) => c.getAttribLocation(pg, key))
+        const uniforms = cached((key) => c.getUniformLocation(pg, key))
+
         const _attribute = (key = '', value: number[], iboValue: number[]) => {
                 const loc = attribs(key, true)
                 createAttrib(c, loc, gl.count, value, iboValue)
         }
 
         const _uniform = (key: string, value: number | number[]) => {
-                createUniform(c, uniforms1(key)!, value)
-                if (!pg2) return
-                c.useProgram(pg2)
-                createUniform(c, uniforms2(key)!, value)
-                c.useProgram(pg1)
+                c.useProgram(pg)
+                createUniform(c, uniforms(key)!, value)
+                cp?._uniform(key, value)
         }
 
         const _texture = (key: string, src: string) => {
+                c.useProgram(pg)
                 loadingImage(gl, src, (source) => {
-                        const loc = uniforms1(key)
-                        const unit = textures(key)
-                        createTexture(c, source, loc, unit)
+                        createTexture(c, source, uniforms(key), units(key))
                 })
         }
 
-        const _storage = (key: string, value: number[] | Float32Array) => {
-                const array = value instanceof Float32Array ? value : new Float32Array(value)
-                const storage = storages(key)
-                const size = Math.ceil(Math.sqrt(array.length))
-                storage.width = size
-                storage.height = size
-                createStorage(c, size, storage, array)
-                c.uniform1i(uniforms1(key), storage.unit)
-        }
-
         const clean = () => {
-                c.deleteProgram(pg1)
-                if (pg2) c.deleteProgram(pg2)
-                for (const { a, b } of storages.map.values()) {
-                        c.deleteTexture(a.texture)
-                        c.deleteTexture(b.texture)
-                        c.deleteFramebuffer(a.buffer)
-                        c.deleteFramebuffer(b.buffer)
-                }
+                cp?.clean()
+                c.deleteProgram(pg)
                 c.getExtension('WEBGL_lose_context')?.loseContext()
         }
 
-        const _compute = () => {
-                c.useProgram(pg2)
-                for (const [, storage] of storages.map) {
-                        const output = currentNum % 2 ? storage.b : storage.a
-                        c.bindFramebuffer(c.FRAMEBUFFER, output.buffer)
-                        c.framebufferTexture2D(c.FRAMEBUFFER, c.COLOR_ATTACHMENT0, c.TEXTURE_2D, output.texture, 0)
-                        c.viewport(0, 0, storage.width, storage.height)
-                        c.drawArrays(c.TRIANGLES, 0, 6)
-                        c.bindFramebuffer(c.FRAMEBUFFER, null)
-                }
-                currentNum++
-                c.useProgram(pg1)
-        }
-
         const render = () => {
-                if (pg2) _compute()
-                c.bindFramebuffer(c.FRAMEBUFFER, null)
-                c.clear(c.COLOR_BUFFER_BIT)
+                cp?.render()
+                c.useProgram(pg)
                 c.viewport(0, 0, ...gl.size)
-                for (const [key, { unit, a, b }] of storages.map) {
-                        const loc = uniforms1(key)
-                        const output = currentNum % 2 ? a : b
-                        c.activeTexture(c.TEXTURE0 + unit)
-                        c.bindTexture(c.TEXTURE_2D, output.texture)
-                        c.uniform1i(loc, unit)
-                }
                 c.drawArrays(c.TRIANGLES, 0, gl.count)
+                c.bindFramebuffer(c.FRAMEBUFFER, null)
         }
 
-        const webgl: WebGLState = { context: c, program: pg1 }
+        const webgl: WebGLState = { context: c, program: pg, storages: cp?.storages }
 
-        return { webgl, render, clean, _attribute, _uniform, _texture, _storage }
+        return { webgl, render, clean, _attribute, _uniform, _texture, _storage: cp?._storage }
 }

package/src/webgpu.ts

@@ -2,8 +2,8 @@ import { nested as cached } from 'reev'
 import { is, loadingImage } from './utils/helpers'
 import {
         createArrayBuffer,
-        createBindings,
         createBindGroup,
+        createBindings,
         createComputePipeline,
         createDepthTexture,
         createDescriptor,
@@ -15,15 +15,53 @@ import {
 import type { GL, WebGPUState } from './types'
 import { compute, fragment, vertex } from './node'
 
+const WORKING_GROUP_SIZE = 32
+
+const computeProgram = (gl: GL, device: GPUDevice, bindings: any) => {
+        let flush = (_pass: GPUComputePassEncoder) => {}
+
+        const storages = cached((_key, value: number[] | Float32Array) => {
+                const { array, buffer } = createArrayBuffer(device, value, 'storage')
+                const { binding, group } = bindings.storage()
+                return { array, buffer, binding, group }
+        })
+
+        const _storage = (key: string, value: number[] | Float32Array) => {
+                const { array, buffer } = storages(key, value)
+                device.queue.writeBuffer(buffer, 0, array as any)
+        }
+
+        const update = (bindGroups: GPUBindGroup[], bindGroupLayouts: GPUBindGroupLayout[], comp: string) => {
+                const pipeline = createComputePipeline(device, bindGroupLayouts, comp!)
+                flush = (pass) => {
+                        pass.setPipeline(pipeline)
+                        bindGroups.forEach((v, i) => pass.setBindGroup(i, v))
+                        const workgroupCount = Math.ceil(gl.particles / WORKING_GROUP_SIZE)
+                        pass.dispatchWorkgroups(workgroupCount, 1, 1)
+                        pass.end()
+                }
+        }
+
+        const render = (pass: GPUComputePassEncoder) => {
+                flush(pass)
+        }
+
+        const clean = () => {
+                for (const { buffer } of storages.map.values()) buffer.destroy()
+        }
+
+        return { storages, _storage, update, render, clean }
+}
+
 export const webgpu = async (gl: GL) => {
         const context = gl.el!.getContext('webgpu') as GPUCanvasContext
         const { device, format } = await createDevice(context, gl.error)
         const bindings = createBindings()
+        const cp = computeProgram(gl, device, bindings)
         let frag: string
-        let vert: string
         let comp: string
+        let vert: string
         let flush = (_pass: GPURenderPassEncoder) => {}
-        let computeFlush = (_pass: GPUComputePassEncoder) => {}
         let needsUpdate = true
         let depthTexture: GPUTexture
 
@@ -35,13 +73,6 @@ export const webgpu = async (gl: GL) => {
                 return { array, buffer, location, stride }
         })
 
-        const storages = cached((_key, value: number[] | Float32Array) => {
-                needsUpdate = true
-                const { array, buffer } = createArrayBuffer(device, value, 'storage')
-                const { binding, group } = bindings.storage()
-                return { array, buffer, binding, group }
-        })
-
         const uniforms = cached((_key, value: number[]) => {
                 needsUpdate = true
                 const { binding, group } = bindings.uniform()
@@ -61,14 +92,10 @@ export const webgpu = async (gl: GL) => {
                 device.queue.writeBuffer(buffer, 0, array as any)
         }
 
-        const _storage = (key: string, value: number[] | Float32Array) => {
-                const { array, buffer } = storages(key, value)
-                device.queue.writeBuffer(buffer, 0, array as any)
-        }
-
         const _uniform = (key: string, value: number | number[]) => {
                 if (is.num(value)) value = [value]
                 const { array, buffer } = uniforms(key, value)
+                array.set(value) // needs to set leatest value
                 device.queue.writeBuffer(buffer, 0, array as any)
         }
 
@@ -86,7 +113,7 @@ export const webgpu = async (gl: GL) => {
                         device,
                         uniforms.map.values(),
                         textures.map.values(),
-                        storages.map.values()
+                        cp.storages.map.values()
                 )
                 const pipeline = createPipeline(device, format, bufferLayouts, bindGroupLayouts, vert, frag)
                 flush = (pass) => {
@@ -96,33 +123,21 @@ export const webgpu = async (gl: GL) => {
                         pass.draw(gl.count, 1, 0, 0)
                         pass.end()
                 }
-                if (comp) {
-                        const computePipeline = createComputePipeline(device, bindGroupLayouts, comp)
-                        computeFlush = (pass) => {
-                                pass.setPipeline(computePipeline)
-                                bindGroups.forEach((v, i) => pass.setBindGroup(i, v))
-                                let maxElements = 1
-                                for (const { array } of storages.map.values())
-                                        maxElements = Math.max(maxElements, array.length)
-                                const workgroupCount = Math.ceil(maxElements / 64)
-                                pass.dispatchWorkgroups(workgroupCount)
-                                pass.end()
-                        }
-                }
+                if (gl.cs) cp.update(bindGroups, bindGroupLayouts, comp)
         }
 
         const render = () => {
                 if (!frag || !vert) {
                         const config = { isWebGL: false, gl }
                         frag = fragment(gl.fs, config) // needs to be before vertex
-                        comp = compute(gl.cs, config)
                         vert = vertex(gl.vs, config)
+                        comp = compute(gl.cs, config)
                 }
                 if (gl.loading) return // MEMO: loading after build node
                 if (needsUpdate) update()
                 needsUpdate = false
                 const encoder = device.createCommandEncoder()
-                if (comp) computeFlush(encoder.beginComputePass())
+                if (gl.cs) cp.render(encoder.beginComputePass())
                 flush(encoder.beginRenderPass(createDescriptor(context, depthTexture)))
                 device.queue.submit([encoder.finish()])
         }
@@ -139,19 +154,12 @@ export const webgpu = async (gl: GL) => {
                 for (const { texture } of textures.map.values()) texture.destroy()
                 for (const { buffer } of uniforms.map.values()) buffer.destroy()
                 for (const { buffer } of attribs.map.values()) buffer.destroy()
-                for (const { buffer } of storages.map.values()) buffer.destroy()
+                cp.clean()
         }
 
         resize()
 
-        return {
-                webgpu: { device, uniforms, textures, attribs, storages } as WebGPUState,
-                render,
-                resize,
-                clean,
-                _attribute,
-                _uniform,
-                _texture,
-                _storage,
-        }
+        const webgpu = { device, uniforms, textures, attribs, storages: cp.storages } as WebGPUState
+
+        return { webgpu, render, resize, clean, _attribute, _uniform, _texture, _storage: cp._storage }
 }