Support for storage textures on WebGPU platform (#5760)

* Support for storage textures on WebGPU platform

* lint

---------

Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>

src/index.js

@@ -57,6 +57,7 @@ export * from './platform/audio/constants.js';
 // PLATFORM / GRAPHICS
 export * from './platform/graphics/constants.js';
 export { createGraphicsDevice } from './platform/graphics/graphics-device-create.js';
+export { BindGroupFormat, BindBufferFormat, BindTextureFormat, BindStorageTextureFormat } from './platform/graphics/bind-group-format.js';
 export { BlendState } from './platform/graphics/blend-state.js';
 export { Compute } from './platform/graphics/compute.js';
 export { DepthState } from './platform/graphics/depth-state.js';

src/platform/graphics/bind-group-format.js

@@ -3,7 +3,7 @@ import { Debug, DebugHelper } from '../../core/debug.js';
 
 import {
     TEXTUREDIMENSION_2D, TEXTUREDIMENSION_CUBE, TEXTUREDIMENSION_3D,
-    SAMPLETYPE_FLOAT
+    SAMPLETYPE_FLOAT, PIXELFORMAT_RGBA8
 } from './constants.js';
 
 let id = 0;
@@ -49,24 +49,54 @@ class BindTextureFormat {
     }
 }
 
+/**
+ * @ignore
+ */
+class BindStorageTextureFormat {
+    /** @type {import('./scope-id.js').ScopeId} */
+    scopeId;
+
+    constructor(name, format = PIXELFORMAT_RGBA8, textureDimension = TEXTUREDIMENSION_2D) {
+        /** @type {string} */
+        this.name = name;
+
+        // PIXELFORMAT_***
+        this.format = format;
+
+        // TEXTUREDIMENSION_***
+        this.textureDimension = textureDimension;
+    }
+}
+
 /**
  * @ignore
  */
 class BindGroupFormat {
+    compute = false;
+
     /**
      * @param {import('./graphics-device.js').GraphicsDevice} graphicsDevice - The graphics device
      * used to manage this vertex format.
      * @param {BindBufferFormat[]} [bufferFormats] - An array of bind buffer formats (uniform
      * buffers). Defaults to an empty array.
      * @param {BindTextureFormat[]} [textureFormats] - An array of bind texture formats (textures).
      * Defaults to an empty array.
+     * @param {BindStorageTextureFormat[]} [storageTextureFormats] - An array of bind storage texture
+     * formats (storage textures), used by the compute shader. Defaults to an empty array.
+     * @param {object} [options] - Object for passing optional arguments.
+     * @param {boolean} [options.compute] - If true, this bind group format is used by the compute
+     * shader.
      */
-    constructor(graphicsDevice, bufferFormats = [], textureFormats = []) {
+    constructor(graphicsDevice, bufferFormats = [], textureFormats = [], storageTextureFormats = [], options = {}) {
         this.id = id++;
         DebugHelper.setName(this, `BindGroupFormat_${this.id}`);
 
+        this.compute = options.compute ?? false;
+        Debug.assert(this.compute || storageTextureFormats.length === 0, "Storage textures can be specified only for compute");
+
         /** @type {import('./graphics-device.js').GraphicsDevice} */
         this.device = graphicsDevice;
+        const scope = graphicsDevice.scope;
 
         /** @type {BindBufferFormat[]} */
         this.bufferFormats = bufferFormats;
@@ -79,8 +109,6 @@ class BindGroupFormat {
         /** @type {BindTextureFormat[]} */
         this.textureFormats = textureFormats;
 
-        const scope = graphicsDevice.scope;
-
         // maps a texture format name to a slot index
         /** @type {Map<string, number>} */
         this.textureFormatsMap = new Map();
@@ -91,6 +119,19 @@ class BindGroupFormat {
             tf.scopeId = scope.resolve(tf.name);
         });
 
+        /** @type {BindStorageTextureFormat[]} */
+        this.storageTextureFormats = storageTextureFormats;
+
+        // maps a storage texture format name to a slot index
+        /** @type {Map<string, number>} */
+        this.storageTextureFormatsMap = new Map();
+        storageTextureFormats.forEach((tf, i) => {
+            this.storageTextureFormatsMap.set(tf.name, i);
+
+            // resolve scope id
+            tf.scopeId = scope.resolve(tf.name);
+        });
+
         this.impl = graphicsDevice.createBindGroupFormatImpl(this);
 
         Debug.trace(TRACEID_BINDGROUPFORMAT_ALLOC, `Alloc: Id ${this.id}`, this);
@@ -107,7 +148,7 @@ class BindGroupFormat {
      * Returns format of texture with specified name.
      *
      * @param {string} name - The name of the texture slot.
-     * @returns {BindTextureFormat} - The format.
+     * @returns {BindTextureFormat|null} - The format.
      */
     getTexture(name) {
         const index = this.textureFormatsMap.get(name);
@@ -118,6 +159,21 @@ class BindGroupFormat {
         return null;
     }
 
+    /**
+     * Returns format of storage texture with specified name.
+     *
+     * @param {string} name - The name of the texture slot.
+     * @returns {BindStorageTextureFormat|null} - The format.
+     */
+    getStorageTexture(name) {
+        const index = this.storageTextureFormatsMap.get(name);
+        if (index !== undefined) {
+            return this.storageTextureFormats[index];
+        }
+
+        return null;
+    }
+
     getShaderDeclarationTextures(bindGroup) {
         let code = '';
         let bindIndex = this.bufferFormats.length;
@@ -138,4 +194,4 @@ class BindGroupFormat {
     }
 }
 
-export { BindBufferFormat, BindTextureFormat, BindGroupFormat };
+export { BindBufferFormat, BindTextureFormat, BindGroupFormat, BindStorageTextureFormat };

src/platform/graphics/bind-group.js

@@ -49,6 +49,7 @@ class BindGroup {
         this.impl = graphicsDevice.createBindGroupImpl(this);
 
         this.textures = [];
+        this.storageTextures = [];
         this.uniformBuffers = [];
 
         /** @type {import('./uniform-buffer.js').UniformBuffer} */
@@ -104,21 +105,46 @@ class BindGroup {
         }
     }
 
+    /**
+     * Assign a storage texture to a named slot.
+     *
+     * @param {string} name - The name of the texture slot.
+     * @param {import('./texture.js').Texture} texture - Texture to assign to the slot.
+     */
+    setStorageTexture(name, texture) {
+        const index = this.format.storageTextureFormatsMap.get(name);
+        Debug.assert(index !== undefined, `Setting a storage texture [${name}] on a bind group with id: ${this.id} which does not contain in, while rendering [${DebugGraphics.toString()}]`, this);
+        if (this.storageTextures[index] !== texture) {
+            this.storageTextures[index] = texture;
+            this.dirty = true;
+        } else if (this.renderVersionUpdated < texture.renderVersionDirty) {
+            // if the texture properties have changed
+            this.dirty = true;
+        }
+    }
+
     /**
      * Applies any changes made to the bind group's properties.
      */
     update() {
 
         // TODO: implement faster version of this, which does not call SetTexture, which does a map lookup
+        const { textureFormats, storageTextureFormats } = this.format;
 
-        const textureFormats = this.format.textureFormats;
         for (let i = 0; i < textureFormats.length; i++) {
             const textureFormat = textureFormats[i];
             const value = textureFormat.scopeId.value;
             Debug.assert(value, `Value was not set when assigning texture slot [${textureFormat.name}] to a bind group, while rendering [${DebugGraphics.toString()}]`, this);
             this.setTexture(textureFormat.name, value);
         }
 
+        for (let i = 0; i < storageTextureFormats.length; i++) {
+            const storageTextureFormat = storageTextureFormats[i];
+            const value = storageTextureFormat.scopeId.value;
+            Debug.assert(value, `Value was not set when assigning storage texture slot [${storageTextureFormat.name}] to a bind group, while rendering [${DebugGraphics.toString()}]`, this);
+            this.setStorageTexture(storageTextureFormat.name, value);
+        }
+
         // update uniform buffer offsets
         this.uniformBufferOffsets.length = this.uniformBuffers.length;
         for (let i = 0; i < this.uniformBuffers.length; i++) {

src/platform/graphics/texture.js

@@ -58,6 +58,9 @@ class Texture {
      */
     renderVersionDirty = 0;
 
+    /** @protected */
+    _storage = false;
+
     /**
      * Create a new Texture instance.
      *
@@ -157,7 +160,10 @@ class Texture {
      * - {@link FUNC_NOTEQUAL}
      *
      * Defaults to {@link FUNC_LESS}.
-     * @param {Uint8Array[]|HTMLCanvasElement[]|HTMLImageElement[]|HTMLVideoElement[]} [options.levels] - Array of Uint8Array or other supported browser interface.
+     * @param {Uint8Array[]|HTMLCanvasElement[]|HTMLImageElement[]|HTMLVideoElement[]} [options.levels]
+     * - Array of Uint8Array or other supported browser interface.
+     * @param {boolean} [options.storage] - Defines if texture can be used as a storage texture by
+     * a compute shader. Defaults to false.
      * @example
      * // Create a 8x8x24-bit texture
      * const texture = new pc.Texture(graphicsDevice, {
@@ -201,6 +207,7 @@ class Texture {
             this._depth = 1;
         }
 
+        this._storage = options.storage ?? false;
         this._cubemap = options.cubemap ?? false;
         this.fixCubemapSeams = options.fixCubemapSeams ?? false;
         this._flipY = options.flipY ?? false;
@@ -546,6 +553,15 @@ class Texture {
         return this._mipmaps;
     }
 
+    /**
+     * Defines if texture can be used as a storage texture by a compute shader.
+     *
+     * @type {boolean}
+     */
+    get storage() {
+        return this._storage;
+    }
+
     /**
      * The width of the texture in pixels.
      *

src/platform/graphics/webgpu/constants.js

@@ -0,0 +1,44 @@
+import {
+    PIXELFORMAT_A8, PIXELFORMAT_L8, PIXELFORMAT_LA8, PIXELFORMAT_RGB565, PIXELFORMAT_RGBA5551, PIXELFORMAT_RGBA4,
+    PIXELFORMAT_RGB8, PIXELFORMAT_RGBA8, PIXELFORMAT_DXT1, PIXELFORMAT_DXT3, PIXELFORMAT_DXT5,
+    PIXELFORMAT_RGB16F, PIXELFORMAT_RGBA16F, PIXELFORMAT_RGB32F, PIXELFORMAT_RGBA32F, PIXELFORMAT_R32F, PIXELFORMAT_DEPTH,
+    PIXELFORMAT_DEPTHSTENCIL, PIXELFORMAT_111110F, PIXELFORMAT_SRGB, PIXELFORMAT_SRGBA, PIXELFORMAT_ETC1,
+    PIXELFORMAT_ETC2_RGB, PIXELFORMAT_ETC2_RGBA, PIXELFORMAT_PVRTC_2BPP_RGB_1, PIXELFORMAT_PVRTC_2BPP_RGBA_1,
+    PIXELFORMAT_PVRTC_4BPP_RGB_1, PIXELFORMAT_PVRTC_4BPP_RGBA_1, PIXELFORMAT_ASTC_4x4, PIXELFORMAT_ATC_RGB,
+    PIXELFORMAT_ATC_RGBA, PIXELFORMAT_BGRA8
+} from '../constants.js';
+
+// map of PIXELFORMAT_*** to GPUTextureFormat
+export const gpuTextureFormats = [];
+gpuTextureFormats[PIXELFORMAT_A8] = '';
+gpuTextureFormats[PIXELFORMAT_L8] = 'r8unorm';
+gpuTextureFormats[PIXELFORMAT_LA8] = 'rg8unorm';
+gpuTextureFormats[PIXELFORMAT_RGB565] = '';
+gpuTextureFormats[PIXELFORMAT_RGBA5551] = '';
+gpuTextureFormats[PIXELFORMAT_RGBA4] = '';
+gpuTextureFormats[PIXELFORMAT_RGB8] = 'rgba8unorm';
+gpuTextureFormats[PIXELFORMAT_RGBA8] = 'rgba8unorm';
+gpuTextureFormats[PIXELFORMAT_DXT1] = 'bc1-rgba-unorm';
+gpuTextureFormats[PIXELFORMAT_DXT3] = 'bc2-rgba-unorm';
+gpuTextureFormats[PIXELFORMAT_DXT5] = 'bc3-rgba-unorm';
+gpuTextureFormats[PIXELFORMAT_RGB16F] = '';
+gpuTextureFormats[PIXELFORMAT_RGBA16F] = 'rgba16float';
+gpuTextureFormats[PIXELFORMAT_RGB32F] = '';
+gpuTextureFormats[PIXELFORMAT_RGBA32F] = 'rgba32float';
+gpuTextureFormats[PIXELFORMAT_R32F] = 'r32float';
+gpuTextureFormats[PIXELFORMAT_DEPTH] = 'depth32float';
+gpuTextureFormats[PIXELFORMAT_DEPTHSTENCIL] = 'depth24plus-stencil8';
+gpuTextureFormats[PIXELFORMAT_111110F] = 'rg11b10ufloat';
+gpuTextureFormats[PIXELFORMAT_SRGB] = '';
+gpuTextureFormats[PIXELFORMAT_SRGBA] = '';
+gpuTextureFormats[PIXELFORMAT_ETC1] = '';
+gpuTextureFormats[PIXELFORMAT_ETC2_RGB] = 'etc2-rgb8unorm';
+gpuTextureFormats[PIXELFORMAT_ETC2_RGBA] = 'etc2-rgba8unorm';
+gpuTextureFormats[PIXELFORMAT_PVRTC_2BPP_RGB_1] = '';
+gpuTextureFormats[PIXELFORMAT_PVRTC_2BPP_RGBA_1] = '';
+gpuTextureFormats[PIXELFORMAT_PVRTC_4BPP_RGB_1] = '';
+gpuTextureFormats[PIXELFORMAT_PVRTC_4BPP_RGBA_1] = '';
+gpuTextureFormats[PIXELFORMAT_ASTC_4x4] = 'astc-4x4-unorm';
+gpuTextureFormats[PIXELFORMAT_ATC_RGB] = '';
+gpuTextureFormats[PIXELFORMAT_ATC_RGBA] = '';
+gpuTextureFormats[PIXELFORMAT_BGRA8] = 'bgra8unorm';

src/platform/graphics/webgpu/webgpu-bind-group-format.js

@@ -3,6 +3,7 @@ import { StringIds } from '../../../core/string-ids.js';
 import { SAMPLETYPE_FLOAT, SAMPLETYPE_UNFILTERABLE_FLOAT, SAMPLETYPE_DEPTH } from '../constants.js';
 
 import { WebgpuUtils } from './webgpu-utils.js';
+import { gpuTextureFormats } from './constants.js';
 
 const samplerTypes = [];
 samplerTypes[SAMPLETYPE_FLOAT] = 'filtering';
@@ -77,6 +78,7 @@ class WebgpuBindGroupFormat {
      * @returns {any} Returns the bind group descriptor.
      */
     createDescriptor(bindGroupFormat) {
+
         // all WebGPU bindings:
         // - buffer: GPUBufferBindingLayout, resource type is GPUBufferBinding
         // - sampler: GPUSamplerBindingLayout, resource type is GPUSampler
@@ -87,8 +89,9 @@ class WebgpuBindGroupFormat {
 
         // generate unique key
         let key = '';
-
         let index = 0;
+
+        // buffers
         bindGroupFormat.bufferFormats.forEach((bufferFormat) => {
 
             const visibility = WebgpuUtils.shaderStage(bufferFormat.visibility);
@@ -112,6 +115,7 @@ class WebgpuBindGroupFormat {
             });
         });
 
+        // textures
         bindGroupFormat.textureFormats.forEach((textureFormat) => {
 
             const visibility = WebgpuUtils.shaderStage(textureFormat.visibility);
@@ -126,6 +130,7 @@ class WebgpuBindGroupFormat {
 
             key += `#${index}T:${visibility}-${gpuSampleType}-${viewDimension}-${multisampled}`;
 
+            // texture
             entries.push({
                 binding: index++,
                 visibility: visibility,
@@ -160,6 +165,31 @@ class WebgpuBindGroupFormat {
             });
         });
 
+        // storage textures
+        bindGroupFormat.storageTextureFormats.forEach((textureFormat) => {
+
+            const { format, textureDimension } = textureFormat;
+            key += `#${index}ST:${format}-${textureDimension}`;
+
+            // storage texture
+            entries.push({
+                binding: index++,
+                visibility: GPUShaderStage.COMPUTE,
+                storageTexture: {
+
+                    // The access mode for this binding, indicating readability and writability.
+                    access: 'write-only', // only single option currently, more in the future
+
+                    // The required format of texture views bound to this binding.
+                    format: gpuTextureFormats[format],
+
+                    // Indicates the required dimension for texture views bound to this binding.
+                    // "1d", "2d", "2d-array", "cube", "cube-array", "3d"
+                    viewDimension: textureDimension
+                }
+            });
+        });
+
         /** @type {GPUBindGroupLayoutDescriptor} */
         const descr = {
             entries: entries

src/platform/graphics/webgpu/webgpu-bind-group.js

@@ -109,6 +109,25 @@ class WebgpuBindGroup {
             });
         });
 
+        // storage textures
+        bindGroup.storageTextures.forEach((tex, textureIndex) => {
+
+            /** @type {import('./webgpu-texture.js').WebgpuTexture} */
+            const wgpuTexture = tex.impl;
+
+            // texture
+            const view = wgpuTexture.getView(device);
+            Debug.assert(view, 'NULL texture view cannot be used by the bind group');
+            Debug.call(() => {
+                this.debugFormat += `${index}: ${bindGroup.format.storageTextureFormats[textureIndex].name}\n`;
+            });
+
+            entries.push({
+                binding: index++,
+                resource: view
+            });
+        });
+
         const descr = {
             layout: bindGroup.format.impl.bindGroupLayout,
             entries: entries