Is there a way to directly specify descriptor sets for custom material uniforms?

[pedrotrschneider]

Hi! I'm in the process of integrating the Slang shading language into Bevy. It's all going super smoothly. I can compile to either GLSL or SPIR-V and it works flawlessly most of the time. My main problem is that there doesn't seem to exist a way to force slang to use a specific descriptor set when compiling, it always binds to descriptor set 0. The problem with that is, for custom materials, Bevy seems to use descriptor set 2 by default for user-defined uniforms. For example, this WGSL shader creates a fullscreen triangle in the vertex shader and, in the fragment shader, draws a pulsating green circle at the center of the screen:

struct FullscreenVertexOutput {
    @builtin(position)
    position: vec4<f32>,
    @location(0)
    uv: vec2<f32>,
};

@vertex
fn vertex(@builtin(vertex_index) vertex_index: u32) -> FullscreenVertexOutput {
    let uv = vec2<f32>(f32(vertex_index >> 1u), f32(vertex_index & 1u)) * 2.0;
    let clip_position = vec4<f32>(uv * vec2<f32>(2.0, -2.0) + vec2<f32>(-1.0, 1.0), 0.0, 1.0);

    return FullscreenVertexOutput(clip_position, uv);
}

struct PostProcessSettings {
    time: f32,
    window_dimensions: vec3<f32>,
}
@group(2) @binding(0) var<uniform> settings: PostProcessSettings;

@fragment
fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
    let uv = (in.uv * 2.0 - 1.0) * settings.window_dimensions.xy;
    let circle_sdf = length(uv) - (240 + 60 * sin(settings.time));

    let color = smoothstep(3.0, 0.0, circle_sdf) * vec3<f32>(0.0, 1.0, 0.0);
    return vec4<f32>(color, 1.0);
}

It works much like a post-processing effect but is implemented as a custom shader on a 2D triangle mesh. The problem is that, as I mentioned, there seems to be no way in Slang to bind uniforms to a specific descriptor set, so I can't use uniforms because they always bind to descriptor set 0. I know there was a way to do that in previous versions of Bevy (before AsBindGroup), but I can't find any documentation on how to do it in the latest version. Is there something I'm missing?

For reference, I'm using Bevy version 13.2, and compiling the Slang code to SPIR-V. I know that the compilation step is not the issue because the same code works if I implement it as a post-processing effect, since post-processing effects allow me to bind stuff to descriptor set 0. I also used naga to translate the compiled SPIR-V code back to WGSL, and manually changed the @group(0) to @group(2), which also worked. This is the code naga generated for me, based on the SPIR-V code the Slang compiler generated:

struct PostprocessingParams_std140_ {
    time: f32,
    window_dimmensions: vec3<f32>,
}

struct GlobalParams_std140_ {
    settings: PostprocessingParams_std140_,
}

var<private> inputuv_1: vec2<f32>;
@group(0) @binding(0) 
var<uniform> globalParams: GlobalParams_std140_;
var<private> fragment_1: vec4<f32>;

fn fragment() {
    let _e12 = inputuv_1;
    let _e16 = globalParams.settings;
    let _e21 = globalParams.settings;
    fragment_1 = vec4<f32>((vec3<f32>(0f, 1f, 0f) * smoothstep(3f, 0f, (length((((_e12 * 2f) - vec2<f32>(1f, 1f)) * _e16.window_dimmensions.xy)) - fma(60f, sin(_e21.time), 240f)))), 1f);
    return;
}

@fragment 
fn main(@location(0) inputuv: vec2<f32>) -> @location(0) vec4<f32> {
    inputuv_1 = inputuv;
    fragment();
    let _e3 = fragment_1;
    return _e3;
}

As I mentioned, this doesn't work, but if I change the @group(0) to @group(2), it works. This is the rust code I'm using. Basically the standard example for CustomMaterial:

use bevy::{
    math::vec3,
    prelude::*,
    render::{
        mesh::MeshVertexBufferLayout,
        render_resource::{
            AsBindGroup, RenderPipelineDescriptor, ShaderRef, ShaderType,
            SpecializedMeshPipelineError,
        },
    },
    sprite::{Material2d, Material2dKey, Material2dPlugin, MaterialMesh2dBundle, Mesh2dHandle},
};

#[derive(Default, Clone, Copy, ShaderType)]
struct PostProcessSettings {
    time: f32,
    window_dimensions: Vec3,
}

#[derive(Component)]
struct FullscreenTriangle;

#[derive(Asset, TypePath, AsBindGroup, Clone)]
struct CustomMaterial {
    #[uniform(0)]
    settings: PostProcessSettings,
}

impl Material2d for CustomMaterial {
    fn vertex_shader() -> ShaderRef {
        return "shaders\\slang_post_processing.vert.spv".into();
    }

    fn fragment_shader() -> ShaderRef {
        return "shaders\\slang_post_processing.frag.spv".into();
    }

    fn specialize(
        descriptor: &mut RenderPipelineDescriptor,
        _layout: &MeshVertexBufferLayout,
        _key: Material2dKey<Self>,
    ) -> Result<(), SpecializedMeshPipelineError> {
        descriptor.vertex.entry_point = "main".into();
        descriptor.fragment.as_mut().unwrap().entry_point = "main".into();
        return Ok(());
    }
}

fn main() {
    App::new()
        .add_plugins((
            DefaultPlugins,
            Material2dPlugin::<CustomMaterial>::default(),
        ))
        .add_systems(Startup, setup)
        .add_systems(Update, update)
        .run();
}

fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<CustomMaterial>>,
) {
    commands.spawn(Camera2dBundle::default());

    commands.spawn((
        MaterialMesh2dBundle {
            mesh: Mesh2dHandle(meshes.add(Triangle2d::new(Vec2::ZERO, Vec2::ZERO, Vec2::ZERO))),
            material: materials.add(CustomMaterial {
                settings: PostProcessSettings {
                    window_dimensions: vec3(0.0, 0.0, 0.0),
                    time: 0.0,
                },
            }),
            transform: Transform::from_xyz(0.0, 0.0, 0.0),
            ..default()
        },
        FullscreenTriangle,
    ));
}

fn update(
    time: Res<Time>,
    window: Query<&Window>,
    material_handles: Query<&Handle<CustomMaterial>>,
    mut materials: ResMut<Assets<CustomMaterial>>,
) {
    let window = window.single();
    for material_handle in material_handles.iter() {
        if let Some(material) = materials.get_mut(material_handle) {
            let settings = &mut material.settings;
            settings.time = time.elapsed_seconds();
            settings.window_dimensions.x = window.width();
            settings.window_dimensions.y = window.height();
            settings.window_dimensions.z = window.width() / window.height();
        }
    }
}

Any help would be greatly appreciated!

[kitlith]

My main problem is that there doesn't seem to exist a way to force slang to use a specific descriptor set when compiling, it always binds to descriptor set 0.

slang has an attribute for specifying the descriptor set for spir-v: https://shader-slang.com/slang/user-guide/spirv-target-specific#vkbindingbinding-int-set-int--0

[[vk::binding(binding: int, set: int = 0)]]

Similar to binding layout qualifier in Vulkan. It specifies the uniform buffer binding point, and the descriptor set for Vulkan.

or if i'm understanding things correctly, you can use register in a similar way that it's used in hlsl.

[pedrotrschneider]

Yes and no. This was actually not working (it was, but the docs were wrong). I oppened an issue in the repo here and there's a discussion describing the whole problem.