Reuse webgpu resources

cocos/game/splash-screen.ts

@@ -580,7 +580,6 @@ export class SplashScreen {
                     );
 
                     cmdBuff.bindPipelineState(bgPso);
-                    cmdBuff.bindDescriptorSet(SetIndex.GLOBAL, pipeline.descriptorSet);
                     cmdBuff.bindDescriptorSet(SetIndex.MATERIAL, bgPass.descriptorSet);
                     cmdBuff.bindInputAssembler(quadAssmebler);
                     cmdBuff.draw(quadAssmebler);
@@ -597,7 +596,6 @@ export class SplashScreen {
                     );
 
                     cmdBuff.bindPipelineState(logoPso);
-                    cmdBuff.bindDescriptorSet(SetIndex.GLOBAL, pipeline.descriptorSet);
                     cmdBuff.bindDescriptorSet(SetIndex.MATERIAL, logoPass.descriptorSet);
                     cmdBuff.bindInputAssembler(quadAssmebler);
                     cmdBuff.draw(quadAssmebler);
@@ -614,7 +612,6 @@ export class SplashScreen {
                     );
 
                     cmdBuff.bindPipelineState(watermarkPso);
-                    cmdBuff.bindDescriptorSet(SetIndex.GLOBAL, pipeline.descriptorSet);
                     cmdBuff.bindDescriptorSet(SetIndex.MATERIAL, wartermarkPass.descriptorSet);
                     cmdBuff.bindInputAssembler(quadAssmebler);
                     cmdBuff.draw(quadAssmebler);

cocos/gfx/base/define.ts

@@ -2255,6 +2255,8 @@ export const DESCRIPTOR_SAMPLER_TYPE = DescriptorType.SAMPLER_TEXTURE | Descript
 
 export const DESCRIPTOR_DYNAMIC_TYPE = DescriptorType.DYNAMIC_STORAGE_BUFFER | DescriptorType.DYNAMIC_UNIFORM_BUFFER;
 
+export const DESCRIPTOR_STORAGE_BUFFER_TYPE = DescriptorType.STORAGE_BUFFER | DescriptorType.DYNAMIC_STORAGE_BUFFER;
+
 export const DRAW_INFO_SIZE = 28;
 
 export type BufferSource = ArrayBuffer | IndirectBuffer;

cocos/gfx/webgpu/define.ts

@@ -22,6 +22,7 @@
  THE SOFTWARE.
 */
 
+import { BufferFlagBit, BufferInfo, BufferUsageBit, MemoryUsageBit } from '../base/define';
 import { DescriptorSet } from '../base/descriptor-set';
 import { DescriptorSetLayout } from '../base/descriptor-set-layout';
 import { WebGPUBuffer } from './webgpu-buffer';
@@ -58,16 +59,54 @@ export function hashCombineStr (str: string, currHash: number): number {
     return hashCombine(hash, currHash);
 }
 
+interface WebGPU {
+    glslang: any;
+    twgsl: any;
+}
+
+export const webGPU: WebGPU = {
+    glslang: undefined,
+    twgsl: undefined,
+};
+function overrideClass (wasm): void {
+    if ('compileGLSL' in wasm) {
+        webGPU.glslang = wasm;
+    } else if ('convertSpirV2WGSL' in wasm) {
+        webGPU.twgsl = wasm;
+    }
+}
+
+export function overrideWebGPUDefine (wasm): void {
+    overrideClass(wasm);
+}
+
 export class DefaultResources {
     // hash, targetResource
     buffersDescLayout: Map<number, WebGPUBuffer> = new Map<number, WebGPUBuffer>();
     texturesDescLayout: Map<number, WebGPUTexture> = new Map<number, WebGPUTexture>();
     samplersDescLayout: Map<number, WebGPUSampler> = new Map<number, WebGPUSampler>();
     buffer!: WebGPUBuffer;
+    storageBuffers: WebGPUBuffer[] = [];
     texture!: WebGPUTexture;
+    cubeTexture!: WebGPUTexture;
     sampler!: WebGPUSampler;
     setLayout!: DescriptorSetLayout;
     descSet!: DescriptorSet;
+    getStorageBuffer (idx: number): WebGPUBuffer {
+        if (this.storageBuffers[idx]) {
+            return this.storageBuffers[idx];
+        }
+        const bufferInfo = new BufferInfo(
+            BufferUsageBit.STORAGE,
+            MemoryUsageBit.DEVICE,
+            16,
+            16, // in bytes
+            BufferFlagBit.NONE,
+        );
+        const defaultBuff = WebGPUDeviceManager.instance.createBuffer(bufferInfo) as WebGPUBuffer;
+        this.storageBuffers[idx] = defaultBuff;
+        return defaultBuff;
+    }
 }
 
 export function isBound (binds: number[], compares: number[]): boolean {

cocos/gfx/webgpu/instantiated.ts

@@ -0,0 +1,125 @@
+/*
+ Copyright (c) 2023 Xiamen Yaji Software Co., Ltd.
+
+ https://www.cocos.com/
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights to
+ use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+ of the Software, and to permit persons to whom the Software is furnished to do so,
+ subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE.
+*/
+
+import { fetchBuffer, ensureWasmModuleReady, instantiateWasm, fetchUrl } from 'pal/wasm';
+import { NATIVE_CODE_BUNDLE_MODE } from 'internal:constants';
+import { error, sys } from '../../core';
+import { NativeCodeBundleMode } from '../../misc/webassembly-support';
+import { WebGPUWasm } from './webgpu-core';
+import { overrideWebGPUDefine } from './define';
+
+const PAGESIZE = 65536; // 64KiB
+
+// How many pages of the wasm memory
+// TODO: let this can be canfiguable by user.
+const PAGECOUNT = 32 * 16;
+
+// How mush memory size of the wasm memory
+const MEMORYSIZE = PAGESIZE * PAGECOUNT; // 32 MiB
+
+let wasmInstance: WebGPUWasm.instance = null!;
+const registerList: any[] = [];
+
+function initWasm (wasmFactory, wasmUrl: string): Promise<void> {
+    return new Promise<void>((resolve, reject) => {
+        const errorMessage = (err: any): string => `[WebGPU]: WebGPU wasm load failed: ${err}`;
+        // eslint-disable-next-line @typescript-eslint/no-floating-promises
+        fetchUrl(wasmUrl).then((currUrl) => {
+            wasmFactory(currUrl).then((Instance: any) => {
+                wasmInstance = Instance;
+                registerList.forEach((cb) => {
+                    cb(wasmInstance);
+                });
+            }).then(resolve).catch((err: any) => reject(errorMessage(err)));
+        });
+    });
+}
+
+function initAsmJS (asmFactory, asmJsMemUrl: string): Promise<void> {
+    return new Promise<void>((resolve, reject) => {
+        fetchBuffer(asmJsMemUrl).then((arrayBuffer) => {
+            const wasmMemory: any = {};
+            wasmMemory.buffer = new ArrayBuffer(MEMORYSIZE);
+            const module = {
+                wasmMemory,
+                memoryInitializerRequest: {
+                    response: arrayBuffer,
+                    status: 200,
+                } as Partial<XMLHttpRequest>,
+            };
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+            return asmFactory(module).then((instance: any) => {
+                wasmInstance = instance;
+                registerList.forEach((cb) => {
+                    cb(wasmInstance);
+                });
+            });
+        }).then(resolve).catch(reject);
+    });
+}
+
+function shouldUseWasmModule (): boolean {
+    if (NATIVE_CODE_BUNDLE_MODE === (NativeCodeBundleMode.BOTH as number)) {
+        return sys.hasFeature(sys.Feature.WASM);
+    } else if (NATIVE_CODE_BUNDLE_MODE === (NativeCodeBundleMode.WASM as number)) {
+        return true;
+    } else {
+        return false;
+    }
+}
+
+export async function waitForWebGPUWasmInstantiation (): Promise<void> {
+    const errorReport = (msg: any): void => {
+        error('Error during WebGPU Wasm instantiation:', msg);
+    };
+    try {
+        await ensureWasmModuleReady();
+        if (shouldUseWasmModule()) {
+            const [
+                glslModule,
+                glslWasmModule,
+                twgslModule,
+                twgslWasmModule,
+            ] = await Promise.all([
+                import('external:emscripten/webgpu/glslang.js'),
+                import('external:emscripten/webgpu/glslang.wasm'),
+                import('external:emscripten/webgpu/twgsl.js'),
+                import('external:emscripten/webgpu/twgsl.wasm'),
+            ]);
+            const glslFactory = glslModule.default;
+            const glslWasmUrl = glslWasmModule.default;
+            const twgslFactory = twgslModule.default;
+            const twgslWasmUrl = twgslWasmModule.default;
+            await Promise.all([initWasm(glslFactory, glslWasmUrl), initWasm(twgslFactory, twgslWasmUrl)]);
+        } else {
+            throw new Error('Wasm module is not supported in this environment.');
+        }
+    } catch (error) {
+        errorReport(error);
+    }
+}
+
+registerList.push(overrideWebGPUDefine);
+
+export const WEBGPU_WASM = 1;

cocos/gfx/webgpu/webgpu-command-buffer.ts

@@ -38,7 +38,7 @@ import {
     Viewport,
     Filter,
     TextureBlit,
-    ShaderStageFlagBit,
+    DescriptorSetInfo,
 } from '../base/define';
 import { Framebuffer } from '../base/framebuffer';
 import { InputAssembler } from '../base/input-assembler';
@@ -70,10 +70,9 @@ import { INT_MAX } from '../../core/math/bits';
 import { GeneralBarrier } from '../base/states/general-barrier';
 import { TextureBarrier } from '../base/states/texture-barrier';
 import { BufferBarrier } from '../base/states/buffer-barrier';
-import { DescUpdateFrequency, WebGPUDeviceManager } from './define';
+import { WebGPUDeviceManager } from './define';
 import { WebGPUSwapchain } from './webgpu-swapchain';
 import { WebGPUPipelineLayout } from './webgpu-pipeline-layout';
-import { WebGPUDescriptorSetLayout } from './webgpu-descriptor-set-layout';
 import { error, errorID } from '../../core';
 
 export interface IWebGPUDepthBias {
@@ -101,6 +100,7 @@ export interface IWebGPUStencilCompareMask {
 interface CommandEncoder { commandEncoder: GPUCommandEncoder, renderPassEncoder: GPURenderPassEncoder }
 let currPipelineState: WebGPUPipelineState | null = null;
 const descriptorSets: WebGPUDescriptorSet[] = [];
+const groupSets: number[] = [0, 1, 2];
 const renderAreas: Rect[] = [];
 export class WebGPUCommandBuffer extends CommandBuffer {
     public pipelineBarrier (
@@ -305,7 +305,7 @@ export class WebGPUCommandBuffer extends CommandBuffer {
             descriptorSets[set] = descriptorSet as WebGPUDescriptorSet;
             this._isStateValid = true;
         }
-        if (dynamicOffsets) {
+        if (dynamicOffsets && dynamicOffsets.length) {
             const offsets = this._curDynamicOffsets[set];
             const dynamicOffsetSize = dynamicOffsets.length;
             for (let i = 0; i < dynamicOffsetSize; i++) offsets[i] = dynamicOffsets[i];
@@ -613,48 +613,23 @@ export class WebGPUCommandBuffer extends CommandBuffer {
         if (!this._curGPUPipelineState) {
             return;
         }
-        const gpuShader = this._curGPUPipelineState.gpuShader!;
-        const bindingMaps = gpuShader.bindings;
-        let vertBinds: number[][] = []; let fragBinds: number[][] = []; let vertAttrs;
-        for (const stage of gpuShader.gpuStages) {
-            if (stage.type === ShaderStageFlagBit.VERTEX) {
-                vertBinds = stage.bindings;
-                vertAttrs = stage.attrs;
-            } else if (stage.type === ShaderStageFlagBit.FRAGMENT) {
-                fragBinds = stage.bindings;
-            }
-        }
         const gpuPipelineLayout = this._curGPUPipelineState.gpuPipelineLayout as IWebGPUGPUPipelineLayout;
         const wgpuPipLayout = (currPipelineState?.pipelineLayout as WebGPUPipelineLayout);
-        let needFetchPipLayout = false;
-        const descSize = descriptorSets.length;
-        for (let i = 0; i < descSize; i++) {
-            descriptorSets[i].prepare(
-                i ? DescUpdateFrequency.NORMAL : DescUpdateFrequency.LOW,
-                bindingMaps.get(i)!,
-                vertBinds[i] || [],
-                fragBinds[i] || [],
-            );
-            const layout = descriptorSets[i].layout as WebGPUDescriptorSetLayout;
-            const currGrpLayout = layout.gpuDescriptorSetLayout!.bindGroupLayout;
-            const notEqualLayout = gpuPipelineLayout.gpuBindGroupLayouts[i] !== currGrpLayout;
-            if (layout.hasChanged || notEqualLayout) {
-                if (notEqualLayout) {
-                    wgpuPipLayout.changeSetLayout(i, layout);
-                }
-                layout.resetChanged();
-                needFetchPipLayout = true;
+        const device = WebGPUDeviceManager.instance;
+        for (let i = 0; i < groupSets.length; i++) {
+            const currSetIdx = groupSets[i];
+            const currDesc = descriptorSets[currSetIdx];
+            if (currDesc) {
+                currDesc.prepare();
+            } else {
+                const currLayout = wgpuPipLayout.setLayouts[currSetIdx];
+                const currLayoutInfo = new DescriptorSetInfo(currLayout);
+                const newDescSet = device.createDescriptorSet(currLayoutInfo) as WebGPUDescriptorSet;
+                descriptorSets[currSetIdx] = newDescSet;
+                newDescSet.prepare(true);
             }
         }
-
-        if (needFetchPipLayout || !wgpuPipLayout.gpuPipelineLayout!.nativePipelineLayout
-            || this._curGPUPipelineState.pipelineState!.layout !== gpuPipelineLayout.nativePipelineLayout) {
-            wgpuPipLayout.fetchPipelineLayout(false);
-            this._curWebGPUPipelineState?.updatePipelineLayout();
-            needFetchPipLayout = true;
-        }
-        this._curWebGPUPipelineState!.prepare(this._curGPUInputAssembler!, needFetchPipLayout);
-        const { dynamicOffsetIndices } = gpuPipelineLayout;
+        this._curWebGPUPipelineState!.prepare(this._curGPUInputAssembler!);
         // ----------------------------wgpu pipline state-----------------------------
         const wgpuPipeline = this._curGPUPipelineState.nativePipeline as GPURenderPipeline;
         const pplFunc = (passEncoder: GPURenderPassEncoder): void => {
@@ -668,20 +643,29 @@ export class WebGPUCommandBuffer extends CommandBuffer {
             };
             this._renderPassFuncQueue.push(stencilRefFunc);
         }
-        const currGPUDescSize = this._curGPUDescriptorSets.length;
+        const currGPUDescSize = groupSets.length;
         const wgpuBindGroups = new Array<GPUBindGroup>(currGPUDescSize);
         const wgpuDynOffsets = new Array<number[]>(currGPUDescSize);
         for (let i = 0; i < currGPUDescSize; i++) {
-            const curGpuDesc = this._curGPUDescriptorSets[i];
-            wgpuBindGroups[i] = curGpuDesc.bindGroup;
-            wgpuDynOffsets[i] = [...this._curDynamicOffsets[i]];
-            if (!descriptorSets[i].dynamicOffsetCount) {
-                wgpuDynOffsets[i] = [];
-            } else if (descriptorSets[i]  && descriptorSets[i].dynamicOffsetCount !== wgpuDynOffsets[i].length) {
-                wgpuDynOffsets[i].length = descriptorSets[i].dynamicOffsetCount;
-                for (let j = 0; j < descriptorSets[i].dynamicOffsetCount; j++) {
-                    if (!wgpuDynOffsets[i][j]) {
-                        wgpuDynOffsets[i][j] = 0;
+            const currSetIdx = groupSets[i];
+            const descObj = descriptorSets[currSetIdx];
+            const curGpuDesc = descObj.gpuDescriptorSet;
+            wgpuBindGroups[currSetIdx] = curGpuDesc.bindGroup;
+            wgpuDynOffsets[currSetIdx] = [...this._curDynamicOffsets[currSetIdx]];
+            if (!descObj.dynamicOffsetCount) {
+                wgpuDynOffsets[currSetIdx] = [];
+            } else if (descObj && descObj.dynamicOffsetCount !== wgpuDynOffsets[currSetIdx].length) {
+                wgpuDynOffsets[currSetIdx].length = descObj.dynamicOffsetCount;
+                for (let j = 0; j < descObj.dynamicOffsetCount; j++) {
+                    const currOffset = wgpuDynOffsets[currSetIdx][j];
+                    if (!currOffset) {
+                        wgpuDynOffsets[currSetIdx][j] = 0;
+                    } else {
+                        const currBind = descObj.dynamicOffsets[j];
+                        const bindObj = descObj.gpuDescriptorSet.gpuDescriptors[currBind];
+                        if (bindObj && bindObj.gpuBuffer && currOffset > bindObj.gpuBuffer.gpuBuffer!.size) {
+                            wgpuDynOffsets[currSetIdx][j] = 0;
+                        }
                     }
                 }
             }
@@ -690,8 +674,12 @@ export class WebGPUCommandBuffer extends CommandBuffer {
         const bgfunc = (passEncoder: GPURenderPassEncoder): void => {
             const gpuBindGroupSize = wgpuBindGroups.length;
             for (let i = 0; i < gpuBindGroupSize; i++) {
+                let currBindGroup = wgpuBindGroups[i];
+                if (!currBindGroup) {
+                    currBindGroup = (device.defaultResource.descSet as WebGPUDescriptorSet).gpuDescriptorSet.bindGroup;
+                }
                 // FIXME: this is a special sentence that 2 in 3 parameters I'm not certain.
-                passEncoder.setBindGroup(i, wgpuBindGroups[i], wgpuDynOffsets[i]);
+                passEncoder.setBindGroup(i, currBindGroup, wgpuDynOffsets[i]);
             }
         };
         this._renderPassFuncQueue.push(bgfunc);