Work around naga/wgpu WGSL instance_index -> GLSL gl_InstanceID bug on WebGL2 (#9383)

naga and wgpu should polyfill WGSL instance_index functionality where it
is not available in GLSL. Until that is done, we can work around it in
bevy using a push constant which is converted to a uniform by naga and
wgpu.

# Objective

- Fixes #9375 

## Solution

- Use a push constant to pass in the base instance to the shader on
WebGL2 so that base instance + gl_InstanceID is used to correctly
represent the instance index.

## TODO

- [ ] Benchmark vs per-object dynamic offset MeshUniform as this will
now push a uniform value per-draw as well as update the dynamic offset
per-batch.
- [x] Test on DX12 AMD/NVIDIA to check that this PR does not regress any
problems that were observed there. (@Elabajaba @robtfm were testing that
last time - help appreciated. <3 )

---

## Changelog

- Added: `bevy_render::instance_index` shader import which includes a
workaround for the lack of a WGSL `instance_index` polyfill for WebGL2
in naga and wgpu for the time being. It uses a push_constant which gets
converted to a plain uniform by naga and wgpu.

## Migration Guide

Shader code before:

```
struct Vertex {
    @builtin(instance_index) instance_index: u32,
...
}

@vertex
fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
...

    var model = mesh[vertex_no_morph.instance_index].model;
```

After:

```
#import bevy_render::instance_index

struct Vertex {
    @builtin(instance_index) instance_index: u32,
...
}

@vertex
fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
...

    var model = mesh[bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)].model;
```

assets/shaders/custom_vertex_attribute.wgsl

@@ -1,5 +1,6 @@
 #import bevy_pbr::mesh_bindings   mesh
 #import bevy_pbr::mesh_functions  mesh_position_local_to_clip
+#import bevy_render::instance_index
 
 struct CustomMaterial {
     color: vec4<f32>,
@@ -22,7 +23,7 @@ struct VertexOutput {
 fn vertex(vertex: Vertex) -> VertexOutput {
     var out: VertexOutput;
     out.clip_position = mesh_position_local_to_clip(
-        mesh[vertex.instance_index].model,
+        mesh[bevy_render::instance_index::get_instance_index(vertex.instance_index)].model,
         vec4<f32>(vertex.position, 1.0),
     );
     out.blend_color = vertex.blend_color;

crates/bevy_pbr/src/prepass/mod.rs

@@ -31,10 +31,10 @@ use bevy_render::{
         BindGroupLayoutEntry, BindingResource, BindingType, BlendState, BufferBindingType,
         ColorTargetState, ColorWrites, CompareFunction, DepthBiasState, DepthStencilState,
         DynamicUniformBuffer, FragmentState, FrontFace, GpuArrayBufferIndex, MultisampleState,
-        PipelineCache, PolygonMode, PrimitiveState, RenderPipelineDescriptor, Shader, ShaderRef,
-        ShaderStages, ShaderType, SpecializedMeshPipeline, SpecializedMeshPipelineError,
-        SpecializedMeshPipelines, StencilFaceState, StencilState, TextureSampleType,
-        TextureViewDimension, VertexState,
+        PipelineCache, PolygonMode, PrimitiveState, PushConstantRange, RenderPipelineDescriptor,
+        Shader, ShaderRef, ShaderStages, ShaderType, SpecializedMeshPipeline,
+        SpecializedMeshPipelineError, SpecializedMeshPipelines, StencilFaceState, StencilState,
+        TextureSampleType, TextureViewDimension, VertexState,
     },
     renderer::{RenderDevice, RenderQueue},
     texture::{FallbackImagesDepth, FallbackImagesMsaa},
@@ -487,6 +487,14 @@ where
             PREPASS_SHADER_HANDLE.typed::<Shader>()
         };
 
+        let mut push_constant_ranges = Vec::with_capacity(1);
+        if cfg!(all(feature = "webgl", target_arch = "wasm32")) {
+            push_constant_ranges.push(PushConstantRange {
+                stages: ShaderStages::VERTEX,
+                range: 0..4,
+            });
+        }
+
         let mut descriptor = RenderPipelineDescriptor {
             vertex: VertexState {
                 shader: vert_shader_handle,
@@ -526,7 +534,7 @@ where
                 mask: !0,
                 alpha_to_coverage_enabled: false,
             },
-            push_constant_ranges: Vec::new(),
+            push_constant_ranges,
             label: Some("prepass_pipeline".into()),
         };
 

crates/bevy_pbr/src/prepass/prepass.wgsl

@@ -3,6 +3,7 @@
 #import bevy_pbr::skinning
 #import bevy_pbr::morph
 #import bevy_pbr::mesh_bindings mesh
+#import bevy_render::instance_index
 
 // Most of these attributes are not used in the default prepass fragment shader, but they are still needed so we can
 // pass them to custom prepass shaders like pbr_prepass.wgsl.
@@ -91,7 +92,7 @@ fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
 #else // SKINNED
     // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
     // See https://github.com/gfx-rs/naga/issues/2416
-    var model = mesh[vertex_no_morph.instance_index].model;
+    var model = mesh[bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)].model;
 #endif // SKINNED
 
     out.clip_position = bevy_pbr::mesh_functions::mesh_position_local_to_clip(model, vec4(vertex.position, 1.0));
@@ -112,7 +113,7 @@ fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
         vertex.normal,
         // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
         // See https://github.com/gfx-rs/naga/issues/2416
-        vertex_no_morph.instance_index
+        bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)
     );
 #endif // SKINNED
 
@@ -122,7 +123,7 @@ fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
         vertex.tangent,
         // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
         // See https://github.com/gfx-rs/naga/issues/2416
-        vertex_no_morph.instance_index
+        bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)
     );
 #endif // VERTEX_TANGENTS
 #endif // NORMAL_PREPASS
@@ -132,7 +133,7 @@ fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
     // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
     // See https://github.com/gfx-rs/naga/issues/2416
     out.previous_world_position = bevy_pbr::mesh_functions::mesh_position_local_to_world(
-        mesh[vertex_no_morph.instance_index].previous_model,
+        mesh[bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)].previous_model,
         vec4<f32>(vertex.position, 1.0)
     );
 #endif // MOTION_VECTOR_PREPASS

crates/bevy_pbr/src/render/mesh.rs

@@ -891,6 +891,14 @@ impl SpecializedMeshPipeline for MeshPipeline {
             ));
         }
 
+        let mut push_constant_ranges = Vec::with_capacity(1);
+        if cfg!(all(feature = "webgl", target_arch = "wasm32")) {
+            push_constant_ranges.push(PushConstantRange {
+                stages: ShaderStages::VERTEX,
+                range: 0..4,
+            });
+        }
+
         Ok(RenderPipelineDescriptor {
             vertex: VertexState {
                 shader: MESH_SHADER_HANDLE.typed::<Shader>(),
@@ -909,7 +917,7 @@ impl SpecializedMeshPipeline for MeshPipeline {
                 })],
             }),
             layout: bind_group_layout,
-            push_constant_ranges: Vec::new(),
+            push_constant_ranges,
             primitive: PrimitiveState {
                 front_face: FrontFace::Ccw,
                 cull_mode: Some(Face::Back),
@@ -1300,6 +1308,12 @@ impl<P: PhaseItem> RenderCommand<P> for DrawMesh {
     ) -> RenderCommandResult {
         if let Some(gpu_mesh) = meshes.into_inner().get(mesh_handle) {
             pass.set_vertex_buffer(0, gpu_mesh.vertex_buffer.slice(..));
+            #[cfg(all(feature = "webgl", target_arch = "wasm32"))]
+            pass.set_push_constants(
+                ShaderStages::VERTEX,
+                0,
+                &(batch_indices.index as i32).to_le_bytes(),
+            );
             match &gpu_mesh.buffer_info {
                 GpuBufferInfo::Indexed {
                     buffer,

crates/bevy_pbr/src/render/mesh.wgsl

@@ -3,6 +3,7 @@
 #import bevy_pbr::morph
 #import bevy_pbr::mesh_bindings       mesh
 #import bevy_pbr::mesh_vertex_output  MeshVertexOutput
+#import bevy_render::instance_index
 
 struct Vertex {
     @builtin(instance_index) instance_index: u32,
@@ -66,7 +67,7 @@ fn vertex(vertex_no_morph: Vertex) -> MeshVertexOutput {
 #else
     // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
     // See https://github.com/gfx-rs/naga/issues/2416
-    var model = mesh[vertex_no_morph.instance_index].model;
+    var model = mesh[bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)].model;
 #endif
 
 #ifdef VERTEX_NORMALS
@@ -77,7 +78,7 @@ fn vertex(vertex_no_morph: Vertex) -> MeshVertexOutput {
         vertex.normal,
         // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
         // See https://github.com/gfx-rs/naga/issues/2416
-        vertex_no_morph.instance_index
+        bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)
     );
 #endif
 #endif
@@ -97,7 +98,7 @@ fn vertex(vertex_no_morph: Vertex) -> MeshVertexOutput {
         vertex.tangent,
         // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
         // See https://github.com/gfx-rs/naga/issues/2416
-        vertex_no_morph.instance_index
+        bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index)
     );
 #endif
 
@@ -108,7 +109,7 @@ fn vertex(vertex_no_morph: Vertex) -> MeshVertexOutput {
 #ifdef VERTEX_OUTPUT_INSTANCE_INDEX
     // Use vertex_no_morph.instance_index instead of vertex.instance_index to work around a wgpu dx12 bug.
     // See https://github.com/gfx-rs/naga/issues/2416
-    out.instance_index = vertex_no_morph.instance_index;
+    out.instance_index = bevy_render::instance_index::get_instance_index(vertex_no_morph.instance_index);
 #endif
 
     return out;

crates/bevy_pbr/src/render/wireframe.wgsl

@@ -1,5 +1,6 @@
 #import bevy_pbr::mesh_bindings    mesh
 #import bevy_pbr::mesh_functions   mesh_position_local_to_clip
+#import bevy_render::instance_index
 
 #ifdef SKINNED
     #import bevy_pbr::skinning
@@ -23,7 +24,7 @@ fn vertex(vertex: Vertex) -> VertexOutput {
 #ifdef SKINNED
     let model = bevy_pbr::skinning::skin_model(vertex.joint_indexes, vertex.joint_weights);
 #else
-    let model = mesh[vertex.instance_index].model;
+    let model = mesh[bevy_render::instance_index::get_instance_index(vertex.instance_index)].model;
 #endif
 
     var out: VertexOutput;

crates/bevy_render/src/instance_index.wgsl

@@ -0,0 +1,17 @@
+#define_import_path bevy_render::instance_index
+
+#ifdef BASE_INSTANCE_WORKAROUND
+// naga and wgpu should polyfill WGSL instance_index functionality where it is
+// not available in GLSL. Until that is done, we can work around it in bevy
+// using a push constant which is converted to a uniform by naga and wgpu.
+// https://github.com/gfx-rs/wgpu/issues/1573
+var<push_constant> base_instance: i32;
+
+fn get_instance_index(instance_index: u32) -> u32 {
+    return u32(base_instance) + instance_index;
+}
+#else
+fn get_instance_index(instance_index: u32) -> u32 {
+    return instance_index;
+}
+#endif

crates/bevy_render/src/lib.rs

@@ -26,6 +26,7 @@ pub mod texture;
 pub mod view;
 
 use bevy_hierarchy::ValidParentCheckPlugin;
+use bevy_reflect::TypeUuid;
 pub use extract_param::Extract;
 
 pub mod prelude {
@@ -56,7 +57,7 @@ use crate::{
     view::{ViewPlugin, WindowRenderPlugin},
 };
 use bevy_app::{App, AppLabel, Plugin, SubApp};
-use bevy_asset::{AddAsset, AssetServer};
+use bevy_asset::{load_internal_asset, AddAsset, AssetServer, HandleUntyped};
 use bevy_ecs::{prelude::*, schedule::ScheduleLabel, system::SystemState};
 use bevy_utils::tracing::debug;
 use std::{
@@ -208,6 +209,9 @@ struct FutureRendererResources(
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, AppLabel)]
 pub struct RenderApp;
 
+pub const INSTANCE_INDEX_SHADER_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 10313207077636615845);
+
 impl Plugin for RenderPlugin {
     /// Initializes the renderer, sets up the [`RenderSet`](RenderSet) and creates the rendering sub-app.
     fn build(&self, app: &mut App) {
@@ -354,6 +358,16 @@ impl Plugin for RenderPlugin {
     }
 
     fn finish(&self, app: &mut App) {
+        load_internal_asset!(
+            app,
+            INSTANCE_INDEX_SHADER_HANDLE,
+            "instance_index.wgsl",
+            Shader::from_wgsl_with_defs,
+            vec![
+                #[cfg(all(feature = "webgl", target_arch = "wasm32"))]
+                "BASE_INSTANCE_WORKAROUND".into()
+            ]
+        );
         if let Some(future_renderer_resources) =
             app.world.remove_resource::<FutureRendererResources>()
         {