From 5b885c96ee76a317864367524265e656c46d458e Mon Sep 17 00:00:00 2001
From: Martin Valigursky <mvaligursky@snapchat.com>
Date: Wed, 23 Apr 2025 16:22:20 +0100
Subject: [PATCH] WGSL support for batcing

---
 .../webgpu/webgpu-shader-processor-wgsl.js    |  2 +-
 .../shader-lib/chunks-wgsl/chunks-wgsl.js     |  4 +--
 .../chunks-wgsl/common/vert/skinBatch.js      | 26 +++++++++++++++++++
 3 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100644 src/scene/shader-lib/chunks-wgsl/common/vert/skinBatch.js

diff --git a/src/platform/graphics/webgpu/webgpu-shader-processor-wgsl.js b/src/platform/graphics/webgpu/webgpu-shader-processor-wgsl.js
index 9fdf8ddb613..3fe42dc4c22 100644
--- a/src/platform/graphics/webgpu/webgpu-shader-processor-wgsl.js
+++ b/src/platform/graphics/webgpu/webgpu-shader-processor-wgsl.js
@@ -196,7 +196,7 @@ class UniformLine {
 //     var<storage, read_write> storageBuffer : Buffer;
 //     var storageTexture : texture_storage_2d<rgba8unorm, write>;
 //     var videoTexture : texture_external;
-// eslint-disable-next-line
+
 const TEXTURE_REGEX = /^\s*var\s+(\w+)\s*:\s*(texture_\w+)(?:<(\w+)>)?;\s*$/;
 // eslint-disable-next-line
 const STORAGE_TEXTURE_REGEX = /^\s*var\s+([\w\d_]+)\s*:\s*(texture_storage_2d|texture_storage_2d_array)<([\w\d_]+),\s*(\w+)>\s*;\s*$/;
diff --git a/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js b/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js
index c4bddb6b703..d947d7f7fbc 100644
--- a/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js
+++ b/src/scene/shader-lib/chunks-wgsl/chunks-wgsl.js
@@ -165,7 +165,7 @@ import shadowPCF3PS from './lit/frag/lighting/shadowPCF3.js';
 import shadowPCF5PS from './lit/frag/lighting/shadowPCF5.js';
 // import shadowPCSSPS from './lit/frag/lighting/shadowPCSS.js';
 // import shadowSoftPS from './lit/frag/lighting/shadowSoft.js';
-// import skinBatchVS from './common/vert/skinBatch.js';
+import skinBatchVS from './common/vert/skinBatch.js';
 import skinVS from './common/vert/skin.js';
 import skyboxPS from './skybox/frag/skybox.js';
 import skyboxVS from './skybox/vert/skybox.js';
@@ -376,7 +376,7 @@ const shaderChunksWGSL = {
     shadowPCF5PS,
     // shadowPCSSPS,
     // shadowSoftPS,
-    // skinBatchVS,
+    skinBatchVS,
     skinVS,
     skyboxPS,
     skyboxVS,
diff --git a/src/scene/shader-lib/chunks-wgsl/common/vert/skinBatch.js b/src/scene/shader-lib/chunks-wgsl/common/vert/skinBatch.js
new file mode 100644
index 00000000000..a3f6258ba0b
--- /dev/null
+++ b/src/scene/shader-lib/chunks-wgsl/common/vert/skinBatch.js
@@ -0,0 +1,26 @@
+export default /* wgsl */`
+attribute vertex_boneIndices: f32;
+
+var texture_poseMap: texture_2d<f32>;
+
+fn getBoneMatrix(indexFloat: f32) -> mat4x4f {
+
+    let width = i32(textureDimensions(texture_poseMap).x);
+    let index: i32 = i32(indexFloat + 0.5) * 3;
+    let iy: i32 = index / width;
+    let ix: i32 = index % width;
+
+    // read elements of 4x3 matrix
+    let v1: vec4f = textureLoad(texture_poseMap, vec2i(ix + 0, iy), 0);
+    let v2: vec4f = textureLoad(texture_poseMap, vec2i(ix + 1, iy), 0);
+    let v3: vec4f = textureLoad(texture_poseMap, vec2i(ix + 2, iy), 0);
+
+    // transpose to 4x4 matrix
+    return mat4x4f(
+        v1.x, v2.x, v3.x, 0,
+        v1.y, v2.y, v3.y, 0,
+        v1.z, v2.z, v3.z, 0,
+        v1.w, v2.w, v3.w, 1.0
+    );
+}
+`;