glre 0.36.0 → 0.38.0

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 }
 }