Register fusion

[nazar-pc]

I think I mentioned this somewhere, but looks like I never created an actual issue.

Since registers are such a scarce resource on GPUs and have huge impact on occupancy, I want to propose a feature that helps with register usage.

I call it "register fusion", but there is probably a more established naming for it.

Here is an example extracted from the actual code base:

fn find_proofs_impl(
    // ...
) {
    let table_6_proof_targets = /* ... */;

    // `0` for left `1` for right
    let left_right = (subgroup_local_invocation_id % 2) as usize;

    let mut group_left_x_index = subgroup_local_invocation_id * 2;

    // Reading positions from table 6
    for table_6_chunk in 0..2 {
        let table_6_proof_targets = subgroup_shuffle(
            table_6_proof_targets,
            SUBGROUP_SIZE / 2 * table_6_chunk + subgroup_local_invocation_id / 2,
        );
        let table_6_proof_target = table_6_proof_targets.to_array()[left_right];

        let table_5_proof_targets = if table_6_proof_target == Position::SENTINEL {
            [Position::SENTINEL; 2]
        } else {
            table_6_positions[table_6_proof_target as usize]
        };

        // Reading positions from table 5
        for table_5_chunk in 0..2 {
            let table_5_proof_targets = subgroup_shuffle(
                table_5_proof_targets,
                SUBGROUP_SIZE / 2 * table_5_chunk + subgroup_local_invocation_id / 2,
            );
            let table_5_proof_target = table_5_proof_targets.to_array()[left_right];

            let table_4_proof_targets = if table_5_proof_target == Position::SENTINEL {
                [Position::SENTINEL; 2]
            } else {
                table_5_positions[table_5_proof_target as usize]
            };

            // Reading positions from table 4
            for table_4_chunk in 0..2 {
                let table_4_proof_targets = subgroup_shuffle(
                    table_4_proof_targets,
                    SUBGROUP_SIZE / 2 * table_4_chunk + subgroup_local_invocation_id / 2,
                );
                let table_4_proof_target = table_4_proof_targets.to_array()[left_right];

                let table_3_proof_targets = if table_4_proof_target == Position::SENTINEL {
                    [Position::SENTINEL; 2]
                } else {
                    table_4_positions[table_4_proof_target as usize]
                };

                // Reading positions from table 3
                for table_3_chunk in 0..2 {
                    let table_3_proof_targets = subgroup_shuffle(
                        table_3_proof_targets,
                        SUBGROUP_SIZE / 2 * table_3_chunk + subgroup_local_invocation_id / 2,
                    );
                    let table_3_proof_target = table_3_proof_targets.to_array()[left_right];

                    let table_2_proof_targets = if table_3_proof_target == Position::SENTINEL {
                        [Position::SENTINEL; 2]
                    } else {
                        table_3_positions[table_3_proof_target as usize]
                    };

                    // Reading positions from table 2
                    for table_2_chunk in 0..2 {
                        let table_2_proof_targets = subgroup_shuffle(
                            table_2_proof_targets,
                            SUBGROUP_SIZE / 2 * table_2_chunk + subgroup_local_invocation_id / 2,
                        );
                        let table_2_proof_target = table_2_proof_targets.to_array()[left_right];

                        let [x_left, x_right] = if table_2_proof_target == Position::SENTINEL {
                            [Position::SENTINEL; 2]
                        } else {
                            table_2_positions[table_2_proof_target as usize]
                        };

                        // Some logic
                    }
                }
            }
        }
    }
}

There are 5 nested loops here, each with just 2 iterations controlled by table_*_chunk variable and table_*_positions are in shared memory.

The high-level design here is that the whole data set is loaded into subgroup registers exactly once at the start to avoid touching memory afterwards and then invocations share local data with each other completely in registers for processing. It adapts to the subgroup sizes to always load unique data at the maximum possible width in parallel, while doing zero redundant/duplicated memory accesses.

---

The observation here is that with careful loop refactoring, each of the 5 nested loops needs exactly 1 bit to track its progress (do while). But what compiler faithfully generates right now is 5 separate registers, one use to track each loop separately.

This was a substantial occupancy constraint in my code base and after refactoring it to use a single register and extracting bits from it on demand resulted in massive performance improvement. Here is the PR and commit doing it, there are other commits of similar nature in that PR which improve register usage too: nazar-pc/abundance@eb1a8a3

My feature request is to make it possible to write idiomatic Rust code and have compiler fuse multiple registers that have only a few bits occupied with a single register + a bit of code to extract necessary values on demand.

I did this with multiple data structures to reduce memory usage and fit into shared memory, which is much harder to do automatically without range types and such. But for registers specifically compiler should have full visibility and a lot more opportunities for automatic transparent rewriting. The compute cost of packing and unpacking small values from registers is almost always worth it when compared to reduced occupancy or spilling into any kind of memory.

---

As for ways to implement this, I think the transparent way when developer doesn't have to do anything at all is ideal, but may be too complex, at least at first.

Maybe some compiler pseudo-macro could be used to tag the variables that must be packed together, something like pack_value!(var_name, num_bits, "tag_name"). In example above it'll be something like this:

for pack_value!(table_6_chunk, 2, "table_chunk") in 0..2 {

I used 2 for the width since the loop uses "while do" pattern rather "do while". But it'll already be a huge improvement that way.

Then all variables tagged with tag_name will be (if possible) packed into the same register.

[Firestar99]

First, Have you tried to replace your loops with bool?:

for table_6_chunk_bool in [false, true] {
    let table_6_chunk = table_6_chunk_bool as u32;

A bool in LLVM types is an i1 in registers (and a typical i8 on stack or heap) instead of a full i32, which makes it a lot easier for some optimization pass in your driver to pack them for you.

Second, can the driver reasonably prove that this code section is in uniform control flow? Meaning that all threads of the workgroup / subgroup are executing these instructions, and no threads have hit an if or a panic by which control flow may have diverged?

I'm asking this from the perspective of an AMD's GCN/RDNA µarch, which have separate 32bit scalar (sgpr) and 32bit vector general purpose registers (vgpr). So their compiler does some heavy analysis whether a value is scalar across the workgroup / subgroup (not sure) and can be put in a single scalar register per workgroup / subgroup, or needs to be placed in vector registers with one register per thread. Consider inspecting the ISA with tools like rga, which primarily target AMD Windows, but RADV also has some flags to trigger it, to see if it correctly figures out that these are scalar.

But storing vector bools in 32/256 vector registers is very inefficient, since you just need one bit, so instead they bitpack them in 1-2 scalar registers (depending on 32/64 wide subgroup configuration, which is dynamic but defaults to 64 for compute iirc). And all of their bool instructions are made for bit manipulation of these scalar registers, even their execution mask is a bool bitpacked scalar register. So my first suggestion of bools could also be a massive improvement.

I'm not sure what happens if it can prove that the bools are scalar, and how it stores them. But usually you're under vector register pressure, not scalar register pressure, so this may just be insignificant.

[nazar-pc]

I have not tried replacing it with bool and it is not a natural thing to do anyway. It would be nice to write something idiomatic, which means using integers rather than booleans. Also this is just a special case that needs 0/1, I believe there are many cases with nested loops that have small number of iterations in them.

The code is intentionally written to maintain uniform control flow to the best of my understanding. There are no panics anywhere and there is one if, but it is at the end of the iteration and synchronization is used right after that, so all threads should be executing in lock step most of the time (you can see the whole code in the link I attached originally).

The issue was with VGPRs and I did use RGP to figure out that I had occupancy issues and VGPRS where the primary constraint for it. Reducing register usage allowed me to increase occupancy dramatically.

From my experience of how dumb/basic GPU optimizations seem to be, it doesn't really do a whole lot of analysis and I have to do even relatively basic transformations that I'd expect compiler to handle automatically my hand. I would not expect the GPU driver to optimize booleans that way and even then it is just one special case.

For example, this basic loop unrolling helped to reduce VGPR usage further, even though it is an extremely basic loop with 0/1/2 iterations based on runtime information: nazar-pc/abundance@23b14a1

I wrote about issues of this nature in #462 already. I don't know if it is Rust-GPU's code generation that hinders these basic optimizations or something else, but it is a big problem for high-performance code.

---

Just to be clear, I don't need help with that particular piece of code since I found the reasons for performance issues and optimized it, but I wish I didn't have to write such awkward code to get there and for these optimizations to happen automatically or with minimal disruptions to ergonomics.