Escape analysis for Wasm GC

[fitzgen]

We could implement an escape analysis for Wasm GC types to figure out which allocations do not escape their stack frame. Then we can do unboxing/SROA, replacing a struct's fields or array elements with a set of local SSA values.

[fitzgen]

Thought about this a bit more recently.

When optimizing Wasm GC I think we will want the following pass ordering:

• Wasm --> translator --> CLIF w/ gc_{alloc,load,store} intrinsics
• (optional) CLIF w/ GC intrinsics --> mid-end --> CLIF w/ GC intrinsics
• (optional) CLIF w/ GC intrinsics --> inliner --> CLIF w/ GC intrinsics
• CLIF w/ GC intrinsics --> escape analysis & SROA --> CLIF w/ GC intrinsics
  • finds GC objects that don't escape
  • does SROA to remove the allocation and replace their fields with SSA values
• CLIF w/ GC intrinsics --> GC legalization & safepoint spilling --> CLIF
  • replaces GC intrinsics with their expanded regular CLIF instructions
  • does the same GC safepoint spilling we do today

From there the pipeline continues as it does today:

• CLIF --> mid-end --> CLIF
• CLIF --> backend --> VCode
• VCode --> regalloc2 --> VCode & register allocations
• VCode & register allocations --> emit --> machine code

The reason we want gc_{alloc,load,store} intrinsics is so that we can reason about them, and elide them when possible, during escape analysis. That is much harder to do if we've already exploded them into their inlined CLIF code (eg bump allocation requires branching on when we run out of capacity, GC loads and stores might require barriers which involve branching, etc).

Note that we want to do escape analysis before safepoint spilling. We don't want to make everything "escape" when we spill it to the stack, and making the analysis understand when spilling to the stack does vs doesn't count as escaping is a bit hairy. It's also unnecessary so long as escape analysis happens before safepoint spilling.

Note that we want to do escape analysis after inlining. An object might escape before inlining, because it is passed as an argument to a function call, but after we inline that function call we see that it no longer escapes.

Those two previous notes already imply a different phase ordering from what we do today. Today, we do safepoint spilling at the end of CLIF construction and before inlining. In this hypothetical future, we would do CLIF translation, then inlining, then safepoint spilling.

What would the GC intrinsics be and where would they be defined? I see two main alternatives:

1. If we keep the GC safepoint spilling in cranelift_frontend, and outside of cranelift_codegen, then I guess these would be well-known ir::ExtFunc sentinels. This is a bit hacky, and also means we wouldn't be able to use the regular filetest infrastructure, at least without a bunch of additional work. I've already been feeling this pain a bit with the lack of filetests for GC safepoint spilling, and including whole CLIF function dumps in regular rust unit tests. It is also not clear to me exactly how we would do inlining (which is part of cranelift_codegen) before GC safepoints (still part of cranelift_frontend). Also we wouldn't be able to attach alias regions to the GC intrinsics.

2. We could move GC safepoint spilling back into cranelift_codegen, and define the escape analysis in cranelift_codegen as well, which would allow the GC intrinsics to be CLIF instructions. This would allow us to attach alias regions to them. It would allow us to use filetests. Inlining before GC safepoint spilling is not a crate dependency issue anymore. It does mean that we need to reject the GC instructions during lowering and report a "you should have legalized these away by now" error. I think that is a fine trade off personally.

I know that (2) may raise questions like "wasn't the whole point of user stack maps to move this stuff out of cranelift_codegen?" Sort of. The key bit was making it so that the mid-end's GVN and alias analyis and all that saw the GC safepoint spilling and wouldn't do optimizations that were invalid in the face of e.g. a moving GC. The proposed phase ordering above still preserves that property (if we did the optional mid-end optimizations before inlining, we would need to disallow GVNing GC intrinsics, as that would open the original can of worms we were trying to avoid by introducint user stack maps back up again).

So yeah, I think (2) is the way to go here. And it is unfortunate that this is a bunch of new passes over the whole IR that we don't do today, but I think that this is somewhat unavoidable. That said, they are only necessary when we are optimizing GC code. Anything that doesn't use GC can completely skip all that.

I also have thoughts on the details of the static analysis that we do for escape analysis, but I'll leave them for a follow up comment.

[fitzgen]

Oh also one other small wrinkle from (2) above: how to do GC legalization and expand the GC instructions into regular CLIF? In (2) we are doing this in cranelift_codegen, so we can't give a cranelift_frontend::FunctionBuilder to the user and tell them to just insert what they need directly. Idea: reuse our inlining infrastructure and have the user supply a whole ir::Function that we inline to replace the gc_{alloc,load,store}.

[fitzgen]

Okay, here's some thoughts on the escape analysis itself, rather than the phase ordering.

(For simplicity, I'm imagining we won't bother doing SROA for arrays or any GC object that is dynamically indexed into, at least for quite a while.)

I don't think we need a sophisticated points-to analysis. I think we can get away with a simple backwards analysis where

• the lattice for each ir::Value is simply

  top = escapes
   |
  bottom = does not escape
  

• the lattice at a program point is then the power set of escaping ir::Values, with set-union as the join function
• and therefore the lattice for the whole program is the map from program point to the set of escaping ir::Values at that program point.

This is a backwards analysis because it propagates backwards from a value's uses.

Our transfer functions would look something like this:

// A value escapes when it is an argument to a function.
escapes(S, (call f(x))) = S union {x}

// A value escapes when it is stored to non-GC memory.
escapes(S, (store x, p)) = S union {x}

// A value escapes when it is returned from the function.
escapes(S, (return x)) = S union {x}

// A value escapes if it is a block arg to a block param that escapes.
escapes(S, (jump block(x))) =
    if block.params[0] in S {
        S union {x}
    } else {
        S
    }

// A value escapes if it is stored into a GC object that escapes. 
escapes(S, (gc_store r[field], x)) =
    if r in S {
        S union {x}
    } else {
        S
    }

// Otherwise, an object escapes if it escapes in any successor.
escapes(S, x) = union( escapes(S, y) for y in successors(x) )

We would initialize S with our function parameters, since anything that came from outside this function has already escaped before we even began (and we can't elide allocations from this function's callers while optimizing this function anyways).

Note that we wouldn't actually compute this for every single instruction, but instead have escapes_in and escapes_out summary sets for each basic block, similar to what we do for liveness during safepoint spilling. Note that:

• An object escapes for the first time along that control-flow path when it appears in the escapes_out set but does not appear in the escapes_in set.
• Conversely, an object has already escaped when it is in both escapes_out and escapes_in.
• An object never escapes if on this control-flow path if it is not in escapes_out (and it will therefore also not be in escapes_in).

Given that information, we can do SROA of GC objects by walking the CFG in reverse post-order (so that visit defs before uses) and doing the following:

• When we define a GC object v3 = gc_alloc(v0, v1, v2) we optimistically remove that allocation (even if v3 escapes! we want to delay boxing until it actually escapes) and we set current_fields[v3] = [v0, v1, v2].
• When we have v7 = gc_load v5[field] before v5 escapes (if it ever does) then we remove the instruction and set current_aliases[v7] = current_fields[v5][field].
• When we find uses of a value in current_aliases, we replace it with the associated alias value.
• When we have gc_store v8[field], v6 before v8 escapes (if it ever does) then we update current_fields[v8][field] = v6.
• When an object escapes for the first time, we insert the delayed gc_alloc just before that escaping operation.

The current_fields and current_aliases maps will need to be scoped hash maps.

We will need to to do something about control flow join points to handle scenarios like this, where we need to insert new block parameters:

               block0(v0, v1):
                 ;; this allocation will get elided
                 v2 = gc_alloc(v0)
                 br_if v1, block1, block2
                /                    \
               /                      \
              /                        \
    block1:                             block2:
      v3 = call f()                       v4 = call g()
      ;; gc_store gets elided,            ;; gc_store gets elided,
      ;; current_fields[v2][0] = v3       ;; current_fields[v2][0] = v4
      gc_store v2[0], v3                  gc_store v2[0], v4
      jump block3                         jump block3
              \                        /
               \                      /
                \                    /
                 block3:
                   ;; gc_load gets elided, need
                   ;; to add a block param to
                   ;; block3 to get the current
                   ;; field
                   v5 = gc_load v2[0]
                   ;; v5 gets replaced with the
                   ;; block param
                   return v5

Not totally sure how we want to do this yet, given that our existing code to do this stuff is in cranelift_frontend rather than cranelift_codegen (the Variables and SSA construction stuff). Might be okay just to move that stuff into cranelift_codegen and publicly export it to make it usable in both places?

[cfallin]

Thanks for writing all of this up!

Regarding the pass order question and to write down here what we discussed already: yes, I think all of this is a good reason to pull safepoint spilling (in some form) back into the core compiler. The main simplification from my PoV that the earlier user stackmaps work brought (in addition to what you describe) is that it removed the load-bearingness of regalloc's and lowering's location tracking for GC correctness, and that would still be true here, so all is well.

I have some detailed comments on the escape analysis rules but on the whole, it sounds like a reasonable direction. To name another aspect explicitly, it's fairly interesting that it is a flow-sensitive analysis: typically I would imagine something more based on points-to analysis with heap abstractions, which would decide the escape status for each abstraction independent of program point. But I like the fact that we can tell where something escapes and box it only then (possibly e.g. on a rare exit path). Also: it's worth clarifying the aliasing invariants. Specifically, if we have e.g. p2 := p1 somewhere in the program (metaphorically at least -- say that p2 is offset into p1), and later we see that p1 escapes, we need to know that we escape p2 too, right?

My one uncertainty around your transfer-function definition: it's not clear to me how the "deep" escape works when a field is stored-to. Say that I allocate p1 = S1 { .x = ... }, p2 = S2 { .y = p1 }, p3 = S3 { .z = p2 }, ... and then eventually return or store p3 somewhere. We need to propagate all reachable objects to the escape-set, not just p3, right? This is where I guess heap abstractions would have come in handy. Interested to understand how this works in your analysis!

[tschneidereit]

That said, they are only necessary when we are optimizing GC code. Anything that doesn't use GC can completely skip all that.

Does that include exception handling? And if so, how fine-grained can we be about it? Would Rust or C++ code that uses exception handling all have to be compiled with these new passes running as well, or only some or even none of it?