RFC: Add MIR-level inlining to the AMDGPU backend

Code Generation AMDGPU
1

Hello everyone,

I’ve recently been working on a machine-level inliner for AMDGPU and it’s been suggested that I start an RFC to discuss alternative solutions.

We have an increasing number of use-cases in the AMDGPU backend where we’d like to experiment with inlining in the backend, after register allocation, but some of them involve things that are not upstream yet so I’m a bit limited in what I can discuss here. Therefore I’ll focus on the use case that my current patch stack is dealing with (whole wave functions). What all the use cases seem to have in common is non-trivial register usage in the AMDGPU backend, which is currently difficult to explain to the register allocator, liveness analysis and other passes - hence the desire to put things in separate functions, let the backend work with the representation that it understands, and only inline late in the pipeline when things won’t be moved around so much.

Why do we have whole wave functions to begin with?

Instructions on the AMDGPU are run by a group of threads executing in lockstep (aka a wave). Oftentimes only a subset of the threads are running, the others being masked away via the EXEC register; but sometimes we really want to execute code with all the threads enabled (e.g. in order to do efficient reductions) [1].

The usual way to express this is via the WWM (Whole Wave Mode) intrinsics (llvm.amdgcn.strict.wwm, llvm.amdgcn.set.inactive etc), which are used by the SIWholeQuadMode pass to figure out regions of code that need to run in WWM. But of course these aren’t formally regions, so sometimes code gets moved around and then our imaginary regions can get split in unfortunate ways. Furthermore, the placement of these intrinsics in non-obvious, and especially where a WWM region should begin is a bit subtle. There are also concerns about the robustness of WWM regions that contain complex control flow.

For one of our simpler use cases, where we only need to run a code snippet in WWM at the end of a shader, we introduced an alternative mechanism, the llvm.amdgcn.init.whole.wave intrinsic. This starts the whole function in WWM, then switches to the original EXEC mask, and then reconverges at the end of the function to run the WWM section. This sidesteps issues with the WWM intrinsics, but it still leads to increased register usage. We have tried to mitigate this by introducing the llvm.amdgcn.dead intrinsic, which has helped but not entirely fixed the problem.

Which leads us to the whole wave functions (the amdgpu_gfx_whole_wave calling convention). These are meant to contain a WWM section and are called via a special intrinsic (llvm.amdgcn.call.whole.wave). Since it’s a separate function, it won’t be broken up by early optimizations, and because IPRA is enabled by default on AMDGPU we get pretty good register allocation too. In the long term, if the whole wave functions + inlining achieves good performance, we would like to consider replacing all the other WWM and init.whole.wave intrinsics with this mechanism (I would feel much better if we already had some real-world performance numbers, but alas we currently don’t and there’s probably a lot of optimization work left to do).

Some notes about the implementation of the inliner

Adding a machine-level inliner is not exactly trivial. A lot of things are tied to the MachineFunction (e.g. MachineMemOperands) and need to be re-created in the caller. The frame info and machine function info also need to be updated. Most of this isn’t very difficult conceptually, it’s just work. Another issue that @shiltian pointed out is that if we inline machine functions, then the MIR and the IR become out of sync, which can be a problem if we emit any kind of information based on the IR (which we currently do e.g. in the form of the IR-based callgraph being dumped into the ELF file and which will obviously become outdated after inlining).

There are also some complications with the pass manager which warrant some discussion. In the current flow, the whole MIR pipeline is run one function at a time – each MachineFunction is generated, optimized and emitted in one go, and then we move on to the next (in call-graph order because that’s enabled by default on AMDGPU). However, if we want to do inlining, we need to have access to the MIR of the callees while processing the caller. This means we need to pause compilation for the callees at the time of inlining, then process all the callers, and only then can we free the MIR of the callees. In the new pass manager, this is handled by making the inliner a ModulePass. This means the MachineFunctions are all generated, then the inliner is run (which removes any fully inlined MachineFunctions[2] and marks them via ShouldNotProcessFunctionPassesAnalysis so they don’t get processed by future function passes), and then the rest of the pipeline is run on each of the functions that are left. In the legacy pass manager, the support is a bit more involved since there’s no built-in mechanism for doing this. I also didn’t want to disturb other targets or introduce new concepts into legacy infra, so instead I added a custom inlining-aware function pass manager used only on AMDGPU and only when inlining is enabled [3].

Alternative designs (that we haven’t prototyped yet)

One alternative that has been discussed internally is to introduce the concept of regions (similar to MLIR regions). This would make the WWM boundaries crystal clear but would probably involve a lot of heavy lifting to support properly. Another option that @arsenm suggested was to use WWM begin/end intrinsics with tokens. This sounds like an improvement over the current intrinsics, but I’m not sure how intuitive it is or how well it generalizes to other use cases [4].

What do you think? I hope this gives a good enough picture of the context. Should we explore other alternatives, or is the inliner palatable? Are there any other issues with the inliner that you would like to bring up, or anything that you consider a blocker?

Thanks for reading and I’m looking forward to your comments!

Footnotes:
[1] ISPC uses the ‘unmasked’ concept to represent these areas of code: Intel® ISPC User’s Guide. Note that ‘unmasked’ can also be applied to entire functions (which is kind of similar to our whole wave functions).
[2] At the moment this is all of them, but it in the future we should be able to support partial inlining.
[3] This still needs a tiny change to the infra, in order to make it possible for the function pass manager to be inherited. But otherwise it’s a pretty cool solution because it lets us keep only the MIR for the whole wave functions in memory; everything else is processed in one go, as before.
[4] As another example, consider dynamic VGPRs: Newer AMDGPU generations allow us to change the number of allocated VGPRs at runtime via the s_alloc_vgpr instruction. We currently take advantage of this by changing the allocation at function boundaries - the register allocator runs one function at a time, decides how many registers each function needs, and we call s_alloc_vgpr with a safe value before jumping to the next function (we use information about the raytracing pipeline to decide when to allocate down, since allocating back up might introduce a lot of stalls). It would be nice if we could inline some of these functions after regalloc. Doing things the other way around (i.e. identifying regions where it’s possible and desirable to change the VGPR allocation) would be much more involved.

2

I think this makes a lot more sense for solving the actual problem at hand. Doing regalloc on separate functions and then inlining post-RA feels like a hack, especially for handling regalloc related issues.

Doing inlining in the backend, especially this late, just seems like it would introduce a lot of problems. There are the mentioned problems with call graph sections, I feel like there could easily be issues with prolog/epilog/scalling conventions and maybe register clobbering, but that might be less of an issue on the GPU side. Introducing a module pass into the pipeline also requires instantiating all machine functions at once which could significantly increase heap memory usage. I guess that’s not an issue for AMDGPU though as big parts of the backend do run in a CGSCC pipeline already.

If this is going to be scoped to AMDGPU and not bother anyone else, I guess this seems fine to me.

3

It does sound like this whole wave mode calling convention is just a way of avoiding coming up with a solution to the representational problem of regions. For compute we still would like to be able to enable WWM for small regions for optimizations, and we should solve these problems anyway for that. I’d hope this isn’t actually that difficult compared to handling machine inlining. I think generally cleaning up the remnants of call sequences, particularly involving stack arguments, will be difficult to get right after register allocation.

If you’re operating after regalloc, you could probably manage to insert these instructions somewhere that isn’t a function boundary. You could still let the regular inliner make the inlining decisions?

4

It does sound like this whole wave mode calling convention is just a way of avoiding coming up with a solution to the representational problem of regions. For compute we still would like to be able to enable WWM for small regions for optimizations, and we should solve these problems anyway for that. I’d hope this isn’t actually that difficult compared to handling machine inlining.

Yes, it’s an alternative to introducing regions. It’s not easy to predict which path is easier. In my opinion, the difficulties with regions are figuring out all the parts of the backend that need to become aware of them, and also designing a representation for them (at the IR level we can probably get away with tokens like you suggested, but I’m not sure if the ENTRY/EXIT_WWM pseudos will be sufficient for the backend representation). I’d really like us to settle on a single representation for both compute and graphics - is anyone on the compute side actively working on this?

If you’re operating after regalloc, you could probably manage to insert these instructions somewhere that isn’t a function boundary. You could still let the regular inliner make the inlining decisions?

You could, but it would be harder to tell where it’s a good idea to insert them.

There are the mentioned problems with call graph sections

Actually those might have been a red herring. The generic code computes those based on the [MIR](https://github.com/llvm/llvm-project/blob/2563006f3178a6443cd2f05b9e4383ce2a4e7003/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp#L2274), it was only some internal code that looks at the IR and it’s not clear if it really needs to be that way.

Thanks for the comments!

5

at the IR level we can probably get away with tokens like you suggested

No. I strongly believe that anything built on types of data flow analysis is misguided. If somebody could sketch out more explicitly what is even meant by a token approach, perhaps that could change my mind. Failing that, my assumption is that tokens ultimately rely on some data flow type analysis, which is exactly what the current representation of WWM already does. The WWM use is bounded by the amdgcn.set.inactive and amdgcn.wwm intrinsics, which effectively already play the role of wwm.begin and wwm.end. I don’t see how adding tokens into the mix is anything more than a cosmetic change.

And actually, I think that @rovka understated the problems with the current representation. There aren’t just “concerns” about mixing WWM with complex control flow; I would just assume that that’s pretty much broken.

We also just found a new issue with the current representation. The AMDGPUUniformIntrinsicCombine pass makes some assumptions based on workgroup size attributes that don’t apply for intrinsics executed in WWM. With the current representation, there is fundamentally no cheap way to check whether an instruction is supposed to be executed in WWM (and control flow makes the even worse).

The bottom line is: We know that all of these problems can go away with a proper region representation, and we know from MLIR practice that regions are often a good idea.

If we had the prior buy-in from the community, then maybe the proper way to go about this would be to boil the ocean and properly add regions to LLVM IR and MIR. That would certainly have secondary benefits, e.g. it could be used for loop optimizations.

Using separate functions is the pragmatic second-best approach for now.

Perhaps the framing as “inlining” is a bit unfortunate since inlining is mostly something heuristic. Whenever we have a function that uses WWM, where the WWM code lives outside of the llvm::{Machine}Function, we should still try to think of it as a single function as much as possible.

1 Like
6

The idea of having “sub-function” regions in LLVM IR has come up in other contexts, so it does seem worth exploring supporting it properly in IR.

Two examples:

#1: We have long had a bunch of weird special-cases to implement windows exception handling, which uses a “funclets” mechanism – this effectively already uses special regions of an LLVM IR function, and those regions need to be kept separate and tracked all the way through to codegen, and it is invalid to move code into or out of them.

Right now, it’s handled using tokens – and a “parent token” chain, to handle multiple levels of subfunction. I believe this mainly works out in existing IR because the regions are only entered/exited through special EH LLVM IR instructions, so most optimizations don’t attempt to do wrong things. While the current representation works, I suspect having subfunction regions would allow this to seem less of a weird one-off hack.

(I suspect you could actually get a similar hack to work for the AMDGPU use case, but adding a second special-case version of the region idea is probably not a good idea.)

#2: For llvm.assume (and in the more recent intrinsic for speculative load), there’s a sort of concept of “ghost values” – where you have some IR instructions computing values which are only needed inside the compiler to give the compiler extra information, and should never be emitted for codegen. And these computations should not negatively impact optimizations of the “real” code.

With current IR, that’s difficult to ensure. Especially since sometimes these computations need conditionals, and now you have basic block structure that may not be obviously removable – but should be removed/ignored. It might be nice if we could have those computations live in distinct subfunction regions designated for ghost value computation, so that it’s easier to ensure that they will not influence other optimizations without needing to actually outline to a separate function.

2 Likes
7

@jyknight Thanks for those examples, they’re really interesting.

I’d like to at least prototype adding proper region support to the IRs, but honestly it sounds like a pretty long and intensive project. If I did that I’d probably start with MIR first, since IMO that’s where it would make the biggest difference for AMDGPU (and if we had MIR support, then we could at least in the beginning use whole wave functions at the IR level, inline right before ISel with maybe some start/end intrinsics and generate MIR regions during ISel based on those).

Would it be ok to have the MIR inliner in tree in the meantime to unblock other work? @arsenm