Add bindless support back to ExtendedMaterial. (#18025)

PR #17898 disabled bindless support for `ExtendedMaterial`. This commit
adds it back. It also adds a new example, `extended_material_bindless`,
showing how to use it.

Author: pcwalton

Cargo.toml

@@ -4317,3 +4317,14 @@ name = "`no_std` Compatible Library"
 description = "Example library compatible with `std` and `no_std` targets"
 category = "Embedded"
 wasm = true
+
+[[example]]
+name = "extended_material_bindless"
+path = "examples/shader/extended_material_bindless.rs"
+doc-scrape-examples = true
+
+[package.metadata.example.extended_material_bindless]
+name = "Extended Bindless Material"
+description = "Demonstrates bindless `ExtendedMaterial`"
+category = "Shaders"
+wasm = false

assets/shaders/extended_material_bindless.wgsl

@@ -0,0 +1,107 @@
+// The shader that goes with `extended_material_bindless.rs`.
+//
+// This code demonstrates how to write shaders that are compatible with both
+// bindless and non-bindless mode. See the `#ifdef BINDLESS` blocks.
+
+#import bevy_pbr::{
+    forward_io::{FragmentOutput, VertexOutput},
+    mesh_bindings::mesh,
+    pbr_fragment::pbr_input_from_standard_material,
+    pbr_functions::{apply_pbr_lighting, main_pass_post_lighting_processing},
+}
+#import bevy_render::bindless::{bindless_samplers_filtering, bindless_textures_2d}
+
+#ifdef BINDLESS
+#import bevy_pbr::pbr_bindings::{material_array, material_indices}
+#else   // BINDLESS
+#import bevy_pbr::pbr_bindings::material
+#endif  // BINDLESS
+
+// Stores the indices of the bindless resources in the bindless resource arrays,
+// for the `ExampleBindlessExtension` fields.
+struct ExampleBindlessExtendedMaterialIndices {
+    // The index of the `ExampleBindlessExtendedMaterial` data in
+    // `example_extended_material`.
+    material: u32,
+    // The index of the texture we're going to modulate the base color with in
+    // the `bindless_textures_2d` array.
+    modulate_texture: u32,
+    // The index of the sampler we're going to sample the modulated texture with
+    // in the `bindless_samplers_filtering` array.
+    modulate_texture_sampler: u32,
+}
+
+// Plain data associated with this example material.
+struct ExampleBindlessExtendedMaterial {
+    // The color that we multiply the base color, base color texture, and
+    // modulated texture with.
+    modulate_color: vec4<f32>,
+}
+
+#ifdef BINDLESS
+
+// The indices of the bindless resources in the bindless resource arrays, for
+// the `ExampleBindlessExtension` fields.
+@group(2) @binding(100) var<storage> example_extended_material_indices:
+    array<ExampleBindlessExtendedMaterialIndices>;
+// An array that holds the `ExampleBindlessExtendedMaterial` plain old data,
+// indexed by `ExampleBindlessExtendedMaterialIndices.material`.
+@group(2) @binding(101) var<storage> example_extended_material:
+    array<ExampleBindlessExtendedMaterial>;
+
+#else   // BINDLESS
+
+// In non-bindless mode, we simply use a uniform for the plain old data.
+@group(2) @binding(50) var<uniform> example_extended_material: ExampleBindlessExtendedMaterial;
+@group(2) @binding(51) var modulate_texture: texture_2d<f32>;
+@group(2) @binding(52) var modulate_sampler: sampler;
+
+#endif  // BINDLESS
+
+@fragment
+fn fragment(
+    in: VertexOutput,
+    @builtin(front_facing) is_front: bool,
+) -> FragmentOutput {
+#ifdef BINDLESS
+    // Fetch the material slot. We'll use this in turn to fetch the bindless
+    // indices from `example_extended_material_indices`.
+    let slot = mesh[in.instance_index].material_and_lightmap_bind_group_slot & 0xffffu;
+#endif  // BINDLESS
+
+    // Generate a `PbrInput` struct from the `StandardMaterial` bindings.
+    var pbr_input = pbr_input_from_standard_material(in, is_front);
+
+    // Calculate the UV for the texture we're about to sample.
+#ifdef BINDLESS
+    let uv_transform = material_array[material_indices[slot].material].uv_transform;
+#else   // BINDLESS
+    let uv_transform = material.uv_transform;
+#endif  // BINDLESS
+    let uv = (uv_transform * vec3(in.uv, 1.0)).xy;
+
+    // Multiply the base color by the `modulate_texture` and `modulate_color`.
+#ifdef BINDLESS
+    // Notice how we fetch the texture, sampler, and plain extended material
+    // data from the appropriate arrays.
+    pbr_input.material.base_color *= textureSample(
+        bindless_textures_2d[example_extended_material_indices[slot].modulate_texture],
+        bindless_samplers_filtering[
+            example_extended_material_indices[slot].modulate_texture_sampler
+        ],
+        uv
+    ) * example_extended_material[example_extended_material_indices[slot].material].modulate_color;
+#else   // BINDLESS
+    pbr_input.material.base_color *= textureSample(modulate_texture, modulate_sampler, uv) *
+        example_extended_material.modulate_color;
+#endif  // BINDLESS
+
+    var out: FragmentOutput;
+    // Apply lighting.
+    out.color = apply_pbr_lighting(pbr_input);
+    // Apply in-shader post processing (fog, alpha-premultiply, and also
+    // tonemapping, debanding if the camera is non-HDR). Note this does not
+    // include fullscreen postprocessing effects like bloom.
+    out.color = main_pass_post_lighting_processing(pbr_input, out.color);
+    return out;
+}

crates/bevy_pbr/src/extended_material.rs

@@ -1,13 +1,16 @@
+use alloc::borrow::Cow;
+
 use bevy_asset::{Asset, Handle};
 use bevy_ecs::system::SystemParamItem;
+use bevy_platform_support::{collections::HashSet, hash::FixedHasher};
 use bevy_reflect::{impl_type_path, Reflect};
 use bevy_render::{
     alpha::AlphaMode,
     mesh::MeshVertexBufferLayoutRef,
     render_resource::{
-        AsBindGroup, AsBindGroupError, BindGroupLayout, BindlessDescriptor,
-        BindlessSlabResourceLimit, RenderPipelineDescriptor, Shader, ShaderRef,
-        SpecializedMeshPipelineError, UnpreparedBindGroup,
+        AsBindGroup, AsBindGroupError, BindGroupLayout, BindGroupLayoutEntry, BindlessDescriptor,
+        BindlessResourceType, BindlessSlabResourceLimit, RenderPipelineDescriptor, Shader,
+        ShaderRef, SpecializedMeshPipelineError, UnpreparedBindGroup,
     },
     renderer::RenderDevice,
 };
@@ -156,27 +159,53 @@ impl<B: Material, E: MaterialExtension> AsBindGroup for ExtendedMaterial<B, E> {
     type Param = (<B as AsBindGroup>::Param, <E as AsBindGroup>::Param);
 
     fn bindless_slot_count() -> Option<BindlessSlabResourceLimit> {
-        // For now, disable bindless in `ExtendedMaterial`.
-        if B::bindless_slot_count().is_some() && E::bindless_slot_count().is_some() {
-            panic!("Bindless extended materials are currently unsupported")
+        // We only enable bindless if both the base material and its extension
+        // are bindless. If we do enable bindless, we choose the smaller of the
+        // two slab size limits.
+        match (B::bindless_slot_count()?, E::bindless_slot_count()?) {
+            (BindlessSlabResourceLimit::Auto, BindlessSlabResourceLimit::Auto) => {
+                Some(BindlessSlabResourceLimit::Auto)
+            }
+            (BindlessSlabResourceLimit::Auto, BindlessSlabResourceLimit::Custom(limit))
+            | (BindlessSlabResourceLimit::Custom(limit), BindlessSlabResourceLimit::Auto) => {
+                Some(BindlessSlabResourceLimit::Custom(limit))
+            }
+            (
+                BindlessSlabResourceLimit::Custom(base_limit),
+                BindlessSlabResourceLimit::Custom(extended_limit),
+            ) => Some(BindlessSlabResourceLimit::Custom(
+                base_limit.min(extended_limit),
+            )),
         }
-        None
     }
 
     fn unprepared_bind_group(
         &self,
         layout: &BindGroupLayout,
         render_device: &RenderDevice,
         (base_param, extended_param): &mut SystemParamItem<'_, '_, Self::Param>,
-        _: bool,
+        mut force_non_bindless: bool,
     ) -> Result<UnpreparedBindGroup<Self::Data>, AsBindGroupError> {
+        force_non_bindless = force_non_bindless || Self::bindless_slot_count().is_none();
+
         // add together the bindings of the base material and the user material
         let UnpreparedBindGroup {
             mut bindings,
             data: base_data,
-        } = B::unprepared_bind_group(&self.base, layout, render_device, base_param, true)?;
-        let extended_bindgroup =
-            E::unprepared_bind_group(&self.extension, layout, render_device, extended_param, true)?;
+        } = B::unprepared_bind_group(
+            &self.base,
+            layout,
+            render_device,
+            base_param,
+            force_non_bindless,
+        )?;
+        let extended_bindgroup = E::unprepared_bind_group(
+            &self.extension,
+            layout,
+            render_device,
+            extended_param,
+            force_non_bindless,
+        )?;
 
         bindings.extend(extended_bindgroup.bindings.0);
 
@@ -188,23 +217,72 @@ impl<B: Material, E: MaterialExtension> AsBindGroup for ExtendedMaterial<B, E> {
 
     fn bind_group_layout_entries(
         render_device: &RenderDevice,
-        _: bool,
-    ) -> Vec<bevy_render::render_resource::BindGroupLayoutEntry>
+        mut force_non_bindless: bool,
+    ) -> Vec<BindGroupLayoutEntry>
     where
         Self: Sized,
     {
-        // add together the bindings of the standard material and the user material
-        let mut entries = B::bind_group_layout_entries(render_device, true);
-        entries.extend(E::bind_group_layout_entries(render_device, true));
+        force_non_bindless = force_non_bindless || Self::bindless_slot_count().is_none();
+
+        // Add together the bindings of the standard material and the user
+        // material, skipping duplicate bindings. Duplicate bindings will occur
+        // when bindless mode is on, because of the common bindless resource
+        // arrays, and we need to eliminate the duplicates or `wgpu` will
+        // complain.
+        let mut entries = vec![];
+        let mut seen_bindings = HashSet::<_>::with_hasher(FixedHasher);
+        for entry in B::bind_group_layout_entries(render_device, force_non_bindless)
+            .into_iter()
+            .chain(E::bind_group_layout_entries(render_device, force_non_bindless).into_iter())
+        {
+            if seen_bindings.insert(entry.binding) {
+                entries.push(entry);
+            }
+        }
         entries
     }
 
     fn bindless_descriptor() -> Option<BindlessDescriptor> {
-        if B::bindless_descriptor().is_some() && E::bindless_descriptor().is_some() {
-            panic!("Bindless extended materials are currently unsupported")
+        // We're going to combine the two bindless descriptors.
+        let base_bindless_descriptor = B::bindless_descriptor()?;
+        let extended_bindless_descriptor = E::bindless_descriptor()?;
+
+        // Combining the buffers and index tables is straightforward.
+
+        let mut buffers = base_bindless_descriptor.buffers.to_vec();
+        let mut index_tables = base_bindless_descriptor.index_tables.to_vec();
+
+        buffers.extend(extended_bindless_descriptor.buffers.iter().cloned());
+        index_tables.extend(extended_bindless_descriptor.index_tables.iter().cloned());
+
+        // Combining the resources is a little trickier because the resource
+        // array is indexed by bindless index, so we have to merge the two
+        // arrays, not just concatenate them.
+        let max_bindless_index = base_bindless_descriptor
+            .resources
+            .len()
+            .max(extended_bindless_descriptor.resources.len());
+        let mut resources = Vec::with_capacity(max_bindless_index);
+        for bindless_index in 0..max_bindless_index {
+            // In the event of a conflicting bindless index, we choose the
+            // base's binding.
+            match base_bindless_descriptor.resources.get(bindless_index) {
+                None | Some(&BindlessResourceType::None) => resources.push(
+                    extended_bindless_descriptor
+                        .resources
+                        .get(bindless_index)
+                        .copied()
+                        .unwrap_or(BindlessResourceType::None),
+                ),
+                Some(&resource_type) => resources.push(resource_type),
+            }
         }
 
-        None
+        Some(BindlessDescriptor {
+            resources: Cow::Owned(resources),
+            buffers: Cow::Owned(buffers),
+            index_tables: Cow::Owned(index_tables),
+        })
     }
 }