`NonUniform` decoration support

[expenses]

Occasionally, a shader variable requires the NonUniform decoration. One example of this is when doing a texture lookup in a ray-tracing shader, where the texture index might not be uniform with the neighbouring rays. See: https://github.com/expenses/ray-tracing-gallery/blob/cb0cb06cc60a5632150e606261a02ba26862ac09/shaders/closesthit.glsl#L136-L138.

I've attempted to do this with asm!:

#[spirv(descriptor_set = 0, binding = 1)] textures: &RuntimeArray<SampledImage<
    Image!(2D, type=f32, sampled=true),
>>,
...

unsafe {
    asm! {
        "OpDecorate {} NonUniform",
        in(reg) &texture_index
    }
};

let uv = Vec2::splat(0.0);
let texture = unsafe {
    textures.index(texture_index as usize)
};
let colour: Vec4 = unsafe { texture.sample_by_lod(uv, 0.0) };

but the decoration wasn't included in the final SPIR-V module. Perhaps it was optimized out? The u32 -> usize conversion that is needed for accessing indexing RuntimeArrays could be a problem as well.

[DJMcNab]

I mentioned this on discord, but I believe the reason that the decoration is optimised out here is because the entry point is generated, which calls the actual function; this call of the function then gets inlined, removing the decoration.

That is, within the function context, your decoration is applying to the function argument, not the entry point variable.

I however don't have a good suggestion for how this could be worked around, at least not without some support in rust-gpu.

[expenses]

@DJMcNab Right, thanks for investigating! And to clarify, this inlining happens in rust-gpu and not via the spirv-opt passes, correct?

[DJMcNab]

I haven't checked it through either way. I'd be inclined to think it's a rust-gpu pass, but I can't be sure.

I suppose one algorithm we could use is to decorate entry point variables with any variables which apply to the entry point function arguments in entry point decoration. This would make the asm! workaround work, without having to try and work out a safe interface to these parts.

[expenses]

I did some testing by writing out the module between each pass. It seems that the DCE pass is removing the decoration. It's possible that it would get removed in other places as well.

[expenses]

Oh but before that, the inlining pass inlines the variable being decorated without updating the decoration. That should be the first thing to fix.

[khyperia]

Just a quick note, I don't think decorating NonUniform via asm! is the way we want to support NonUniform - asm! is intended to be an internal implementation detail, rather than something users write. We probably want to look into properly designing something instead of trying to hack up asm! to do something it was not intended to do.

[expenses]

Just a quick note, I don't think decorating NonUniform via asm! is the way we want to support NonUniform - asm! is intended to be an internal implementation detail, rather than something users write. We probably want to look into properly designing something instead of trying to hack up asm! to do something it was not intended to do.

Yeah, I guess you want fn non_uniform<T>(value: T) -> T, similar to https://github.com/KhronosGroup/GLSL/blob/master/extensions/ext/GL_EXT_nonuniform_qualifier.txt.

[ElectronicRU]

What about uniformity analysis? We already do rather heavy block-graph analyses in linker
Assigning nonuniform automatically could be pretty useful wrt Just Working (will also help catch errors in places where non-uniform values are not actually supported)

[expenses]

I looked into implemented the support via an intrinsic and emit_global: main...expenses:non-uniform-intrinsic but it seems that the decoration is still being silently ignored/removed. I don't know what I'm doing though and this could still be the right approach.

[Jasper-Bekkers]

What about uniformity analysis? We already do rather heavy block-graph analyses in linker
Assigning nonuniform automatically could be pretty useful wrt Just Working (will also help catch errors in places where non-uniform values are not actually supported)

For that please see #49

[SiebenCorgie]

I tried to find out why @expenses decoration was removed before. For that, I added a index_nonuniform function to RuntimeArray.

Apparently, the Inliner was responsible for removing the decoration. Since the decoration happens in an inlined function, the decoration's IDs need to be rewritten as well, not just the inlined part of the function. This is implemented here.

However, this does not fix everything. Unfortunately it is not enough to keep the NonUniform decoration just for the result of OpAccessChain. Image::sample does an OpLoad for the Image's ID. In our non-uniform case, this ID must be decorated with NonUniform as well. However, it seems non-trivial to find out when to decorate this load and when not. In my experience, the Image::sample function is often not inlined. So you'd need one version with a non-uniform decorated %image and one without. I'm not sure how that would be implemented. For now, I just added a special sample_nonuniform function here.

With all that in place, the correctly decorated code gets emitted for the following Rust code:

fn some_function(..., image_array: &mut RuntimeArray<Image!(2D, type = f32, sampled = true)>){
	let nonuniformly_calculated_index = ...;
	let value = image_array
		.index_nonuniform(nonuniformly_calculated_index)
		.sample_nonuniform(
		    some_sampler,
		    some_vec2,
		);
}

A fix for now could be to just decorate that OpLoad in all Image::sample-like functions. Together with the Inliner-fix that would enable some non-uniform access. AFAIK SpirT will track decorations better later anyway.

I didn't test RuntimeArrays of buffers (RuntimeArray<T> and RuntimeArray<TypedBuffer<T>> / RuntimeArray<TypedBuffer<[T]>>) but I think I could come up with a similar workaround for those.