glre 0.39.0 → 0.40.0

package/src/utils/webgpu.ts

@@ -1,5 +1,5 @@
 import { nested as cached } from 'reev'
-import { is, loadingImage } from './helpers'
+import { is, loadingImage, getStride, WGSL_FS, WGSL_VS } from './helpers'
 import {
         createArrayBuffer,
         createBindGroup,
@@ -11,33 +11,10 @@ import {
         createPipeline,
         createTextureSampler,
         createVertexBuffers,
+        workgroupCount,
 } from './pipeline'
 import type { GL, WebGPUState } from '../types'
 
-const WORKING_GROUP_SIZE = 32
-
-const DEFAULT_VERTEX = /* rust */ `
-struct In { @builtin(vertex_index) vertex_index: u32 }
-struct Out { @builtin(position) position: vec4f }
-@vertex
-fn main(in: In) -> Out {
-  var out: Out;
-  var x = f32(in.vertex_index % 2) * 4.0 - 1.0;
-  var y = f32(in.vertex_index / 2) * 4.0 - 1.0;
-  out.position = vec4f(x, y, 0.0, 1.0);
-  return out;
-}
-`.trim()
-
-const DEFAULT_FRAGMENT = /* rust */ `
-struct Out { @builtin(position) position: vec4f }
-@group(0) @binding(0) var<uniform> iResolution: vec2f;
-@fragment
-fn main(out: Out) -> @location(0) vec4f {
-  return vec4f(fract((out.position.xy / iResolution)), 0.0, 1.0);
-}
-`
-
 const computeProgram = (gl: GL, device: GPUDevice, bindings: any) => {
         let flush = (_pass: GPUComputePassEncoder) => {}
 
@@ -57,8 +34,8 @@ const computeProgram = (gl: GL, device: GPUDevice, bindings: any) => {
                 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)
+                        const { x, y, z } = workgroupCount(gl.particleCount)
+                        pass.dispatchWorkgroups(x, y, z)
                         pass.end()
                 }
         }
@@ -86,12 +63,12 @@ export const webgpu = async (gl: GL) => {
         let needsUpdate = true
         let depthTexture: GPUTexture
 
-        const attribs = cached((_key, value: number[]) => {
+        const attribs = cached((_key, value: number[], isInstance = false) => {
                 needsUpdate = true
-                const stride = value.length / gl.count
+                const stride = getStride(value.length, isInstance ? gl.instanceCount : gl.count)
                 const { location } = bindings.attrib()
                 const { array, buffer } = createArrayBuffer(device, value, 'attrib')
-                return { array, buffer, location, stride }
+                return { array, buffer, location, stride, isInstance }
         })
 
         const uniforms = cached((_key, value: number[]) => {
@@ -110,13 +87,20 @@ export const webgpu = async (gl: GL) => {
 
         const _attribute = (key = '', value: number[]) => {
                 const { array, buffer } = attribs(key, value)
+                array.set(value)
+                device.queue.writeBuffer(buffer, 0, array as any)
+        }
+
+        const _instance = (key: string, value: number[]) => {
+                const { array, buffer } = attribs(key, value, true)
+                array.set(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
+                array.set(value)
                 device.queue.writeBuffer(buffer, 0, array as any)
         }
 
@@ -141,7 +125,7 @@ export const webgpu = async (gl: GL) => {
                         pass.setPipeline(pipeline)
                         bindGroups.forEach((v, i) => pass.setBindGroup(i, v))
                         vertexBuffers.forEach((v, i) => pass.setVertexBuffer(i, v))
-                        pass.draw(gl.count, 1, 0, 0)
+                        pass.draw(gl.count, gl.instanceCount, 0, 0)
                         pass.end()
                 }
                 if (gl.cs) cp.update(bindGroups, bindGroupLayouts, comp)
@@ -150,8 +134,8 @@ export const webgpu = async (gl: GL) => {
         const render = () => {
                 if (!frag || !vert) {
                         const config = { isWebGL: false, gl }
-                        frag = gl.fs ? (is.str(gl.fs) ? gl.fs : gl.fs.fragment(config)) : DEFAULT_FRAGMENT
-                        vert = gl.vs ? (is.str(gl.vs) ? gl.vs : gl.vs.vertex(config)) : DEFAULT_VERTEX
+                        frag = gl.fs ? (is.str(gl.fs) ? gl.fs : gl.fs.fragment(config)) : WGSL_FS
+                        vert = gl.vs ? (is.str(gl.vs) ? gl.vs : gl.vs.vertex(config)) : WGSL_VS
                         comp = gl.cs ? (is.str(gl.cs) ? gl.cs : gl.cs.compute(config)) : ''
                 }
                 if (gl.loading) return // MEMO: loading after build node
@@ -182,5 +166,5 @@ export const webgpu = async (gl: GL) => {
 
         const webgpu = { device, uniforms, textures, attribs, storages: cp.storages } as WebGPUState
 
-        return { webgpu, render, resize, clean, _attribute, _uniform, _texture, _storage: cp._storage }
+        return { webgpu, render, resize, clean, _attribute, _instance, _uniform, _texture, _storage: cp._storage }
 }