Lack of unrolling and very high register usage

[nazar-pc]

I have many places in the code that can be condensed to something like this:

const MAX_BUCKET_SIZE: usize = 512;
const WORKGROUP_SIZE: u32 = 256;

unsafe {
    core::hint::assert_unchecked(matches_count <= MAX_BUCKET_SIZE);
}
for index in (local_invocation_id..matches_count as u32).step_by(WORKGROUP_SIZE as usize) {

Can also be rewritten in a way that will probably produce some helpful runtime code in SPIR-V since I suspect the hint will be lost in translation:

const MAX_BUCKET_SIZE: usize = 512;
const WORKGROUP_SIZE: u32 = 256;

for index in (local_invocation_id..matches_count.min(MAX_BUCKET_SIZE) as u32)
    .step_by(WORKGROUP_SIZE as usize)
{

It is not difficult for me to see that there will be at most two loop iterations here per invocation. However, it is not something compiler sees today.

The result is much higher register usage, impacting occupancy in a big way.

Rewriting it to inner function that is called twice with explicit bounds checks fixes register usage (though I hit #461 when doing so), but is far from idiomatic Rust and is quite painful to do manually in all such cases.

I wish end-to-end compilation chain was aware of things like this, it is a very important pattern for performance.

In fact loop unrolling is extremely bad right now, even fixed loops with 3-4 iterations and one or several simply ALU instructions in it are not unrolled and balloon register usage.

[Firestar99]

assert_unchecked calls crate::intrinsics::assume, so looking into our assume codegen, it's just getting discarded:

rust-gpu/crates/rustc_codegen_spirv/src/builder/intrinsics.rs

Lines 376 to 383 in 6bb1066

	// FIXME(eddyb) `assume` is not implemented atm, so all of its forms should
	// avoid computing its (potentially illegal) bool input in the first place.
	fn assume(&mut self, _val: Self::Value) {}
	fn expect(&mut self, cond: Self::Value, _expected: bool) -> Self::Value {
	// TODO: llvm.expect
	cond
	}

rust-gpu/crates/rustc_codegen_spirv/src/builder/builder_methods.rs

Lines 1950 to 1953 in d5c290e

	// FIXME(eddyb) `assume` is not implemented atm, so all of its forms should
	// avoid computing its (potentially illegal) bool input in the first place.
	fn assume_integer_range(&mut self, _imm: Self::Value, _ty: Self::Type, _range: WrappingRange) {}
	fn assume_nonnull(&mut self, _val: Self::Value) {}

However, what isn't getting ignored is unreachable:

rust-gpu/crates/rustc_codegen_spirv/src/builder/builder_methods.rs

Lines 1541 to 1543 in d5c290e

	fn unreachable(&mut self) {
	self.emit().unreachable().unwrap();
	}

Could you try something like this?

if matches_count <= MAX_BUCKET_SIZE {
    unsafe {
        core::hint::unreachable();
    }
}

If this works, I'm willing to just wire assume to be that

[Firestar99]

There's an extension called SPV_KHR_expect_assume that adds OpAssumeTrueKHR and OpExpectKHR. The Vulkan part VK_KHR_shader_expect_assume is promoted to be required in Vulkan1.4. So we could even emit those, and have a fallback to unreachable.

Problem: rspirv on crates.io doesn't provide those symbols yet... gfx-rs/rspirv#259

[nazar-pc]

Unreachable should have the same effect as the second loop example I have provided. Handling of loops feels completely broken right now in terms of optimizations.