SpirV inline linking and optimization complexity explosion

[PhenaOfMari]

Before I start: I'm not sure how much of this is actionable, but @Firestar99 suggested I summarize my recent findings in an issue so here I am.

I have recently been attempting to implement a set of compute shaders for a fairly complex algorithm. I finally got it to the point where I could attempt to compile it with spirv-builder and found that the initial successful compile took on the order of 30 minutes. Discussion with Firestar led me to flags for outputting the total timings of various parts of the compilation, finding the majority of the build time took place in the following steps:

• link_inline (about a minute)
• link_block_ordering_pass_and_mem2reg-after-inlining (about 3 minutes)
• link_spirv_opt (about 31 minutes)

This lead to discussion about how any functions that touch GPU buffers will be forced inline and optimized, which seems a fairly important thing to note. So, I started digging through my functions looking for things that don't make any sense with the knowledge that they would be inlined and here are the things that I found:

The Big Inline

let result = if condition {
    big_function(args..., mut_buffer_a, mut_buffer_b, more_args...)
} else {
    big_function(args..., mut_buffer_b, mut_buffer_a, more_args...)
}

Needing to access these swapping buffers before and after the function call limited my options. I couldn't just hand off the references to a set of first/second buffers because the code later wouldn't have any understanding of which one was which unless i also propagated the condition along with it. My solution here was to move the conditional inside of big_function and move them to a new set of variables in the proper order before anything else uses them. Doing this cut down the lengthy build calls by about 2/3rds resulting in a total compile time in the order of 10 minutes.

The Big Switch

match case {
    1 => {
        func(args1..., buffers...);
        func(args2..., buffers...);
        func(args3..., buffers...);
        func(args4..., buffers...);
        func(args5..., buffers...);
    },
    2 => {
        func(args6..., buffers...);
        func(args7..., buffers...);
        func(args8..., buffers...);
        func(args9..., buffers...);
    },
    // etc
}

This one I was debating about cleaning up before for other reasons, but it turns out doing it this way is a nightmare for the forced inlines too. What was essentially happening here is a 52 case switch where each case calls the same function either 4 or 5 times in a row with different set of arguments. The solution was pretty simple, instead of calling the function in the switch, use the switch to solve for the argument set and call the functions as applicable afterward. Changing this cut the compile time down to a nice 10 seconds.

However, attempting to solve this the first way I tried gave me another issue in the builder...

Some(tuple)

Since the cases in the switch above needed an inconsistent number of function calls, I had at first attempted to use an optional tuple for the arguments of the calls which may not happen in a given case:

let (first, second, third, fourth, fifth) = match case {
    1 => (args1, args2, args3, args4, Some(args5),
    2 => (args6, args7, args8, args9, None,
    // etc
}
func(first..., buffers...);
func(second..., buffers...);
func(third..., buffers...);
func(fourth..., buffers...);
if let Some(fifth) = fifth {
    func(fifth..., buffers...);
}

When compiling this, the compiler gave me a huge bunch of errors regarding trying to convert the optional tuple (in my case a set of u32s) back or forth to a slice of u8s. I didn't linger too much on this since it was more important to me to solve the switch inline problem, but it was pretty clear that the Optional was trying to do some buffer conversion. I don't know if this is supposed to happen or not, but it stuck out enough for me to think I ought to bring it up here.

[eddyb]

When compiling this, the compiler gave me a huge bunch of errors regarding trying to convert the optional tuple (in my case a set of u32s) back or forth to a slice of u8s.

To be clear on what happened here (since the errors are sadly emergent), Option<T> (or any other enum), will be, in the general case, represented as a tagged union, which rustc lowers to a bunch of pointer casts.

Since only one variant contributes any data fields and they don't overlap the tag, we could, in theory, turn Option<T> into something more like (bool, MaybeUninit<T>), but this can get pretty messy/fragile.

The reason, btw, that some special cases like Option<usize> happen to work, is that rustc already turns those into something akin to (bool, MaybeUninit<usize>) via the "scalar pair" layout optimization, whereas the general case looks more like union { tag: bool, some: (/* bool-sized gap */, T) }.

---

but it was pretty clear that the Optional was trying to do some buffer conversion

It's pointer casts, which you should be able to more directly observe by e.g. trying to use unions directly.
Right now, the untyped memory Rust expects clashes with the typed memory Vulkan SPIR-V mandates, and this semantic mismatch has been most of my focus for the past 2-3 years - though obviously the current state of Rust-GPU diagnostics in this area have been largely ignored (tbh a failure of imagination, in some cases).

Not to mention that most, if not all, "pointer casts" you see in errors are actually generated internally by rustc_codegen_spirv for declared-vs-accessed mismatches, so should probably say "cannot access memory declared as T with type U" or some such (feel free to throw that into another issue, or add to an existing one if it already exists - even Option<(u32, u32)> or Result<i32, u32> could work as an example of "confusing pointer cast error").

_{There are future Rust-GPU branches in various stages of completion, that could make some of your attempts work, or even compile much faster, but it's not always obvious what's even possible, or how much different your code should look like (e.g. most of my experiments have relied on letting the inliner run even more aggressively, and even qptr::legalize's current state benefits from casting pointers to integers to call functions, bypassing inlining, in non-panicking code at least) - but if you are interested in early testing, DM me on e.g. Discord (@eddyb).}

---

In general, having many calls is (quite literally) a recipe for amplification (both compile time and runtime), and even if the GPU driver doesn't end up doing any expensive inlining itself, non-uniform control-flow means that e.g.:

let i = local_id % 4;
match i {
    0 => (f(0), g(0), h(0)),
    1 => (f(1), g(1), h(1)),
    2 => (f(2), g(2), h(2)),
    3 => (f(3), g(3), h(3)),
    _ => unreachable!(),
}

will take (at least) 4x longer to execute than (f(i), g(i), h(i)), because invocations with different i cannot be active at the same time, and its execution is comparable to an unrolled version of this:

for j in 0..4 {
    if i == j {
        result = (f(j), g(j), h(j));
    }
}
// tbc, this is what I meant as the unrolling of the above loop:
if i == 0 { result = (f(0), g(0), h(0)); }
if i == 1 { result = (f(1), g(1), h(1)); }
if i == 2 { result = (f(2), g(2), h(2)); }
if i == 3 { result = (f(3), g(3), h(3)); }

If you started from a fully active subgroup, every iteration of that loop will call all 3 functions, with different subsets of invocations being active each time - same as the match, but without 12 inlined functions (well, unless unrolling duplicates the 3 calls/inlined callees again, ending up at the match).

The terminology for this varies a bit, but (f(i), g(i), h(i)) is the data-parallel version (whether or not inlining occurs), while the match (and equivalent forms) serializes your disjoint control-flow paths.

The solution was pretty simple, instead of calling the function in the switch, use the switch to solve for the argument set and call the functions as applicable afterward. Changing this cut the compile time down to a nice 10 seconds.

If you have a way to measure the runtime performance of that shader, I would be curious, because, if case actually varies between adjacent invocations, pulling the function calls out of the match would result in maximal data parallelism, and actually take advantage of GPU architecture (as opposed to accidentally serializing execution).