[kripken]

This is a form of escape analysis, like Heap2Local, but it can handle multiple
allocations which Heap2Local can't, but otoh those must fit a simple, easy-to
identify pattern.

The pattern here is that while structs are allocated and used, all the uses
are simple struct.gets/sets. So we can just use the data in the fields
directly:

   x = new Foo(42);
   if (condition) x = new Foo(1337);
   else x.value = -1;
   bar(x.value, x.value);

=>

   x_value = (new Foo(42)).value;
   if (condition) x_value = (new Foo(1337)).value;
   else x_value = -1;
   bar(x_value, x_value);

We add a new x_value local here which stores the field, and write to
it from struct.new/struct.set, and read from it in struct.get. If those
are the only operations - no escape of the reference, and no ref.eq
etc. - then this is clearly valid to do. And after doing it, the allocations are
"untangled" - each is used just to write its fields, so their live ranges do
not overlap, and Heap2Local will remove them all. So this new pass,
when it optimizes, always leads to beneficial results, without tradeoffs.

This is a reboot of #9271, with the changes that it handles struct.sets,
and does not handle unknown values (the latter introduced unclear
tradeoffs). As a result the code in the pass is a lot simpler now.

This helps #8641 and almost fixes it. Specifically, once we merge
the two locals there, then this new pass will do all the rest. And we
should be getting those locals merged either through @stevenfontanella's
nullability pass, or my TypeGeneralizing pass.

Aside from helping that original testcase in #8641 (after I manually
merged the locals), this also helps 3/4 large wasm testcases I tried on,
and has no downside on 1/4.

[gemini-code-assist]

Code Review

This pull request introduces a new optimization pass, HeapUntangling, which performs multi-allocation escape analysis to cleanly separate multiple allocations written to a single local. Feedback on the implementation suggests clearing the localInfos and optimizedLocals collections at the start of doWalkFunction to prevent stale data from being carried over across functions. Additionally, it is recommended to explicitly mark a local as unoptimizable when a LocalSet (tee) is used as the reference of a StructSet instead of relying on an implicit count mismatch.

[tlively]

Comments so far. Still need to read the tests.

[tlively]

Alternative name for your consideration: LocalUnboxing

[kripken]

Hmm, "unboxing" is interesting, but "local" doesn't seem as clear as something with Heap. HeapUnboxing? But I like that untangling clarifies that we are separating/making something simpler.

[tlively]

"Local" works with two meanings: the pass is function-local (similar to LocalCSE, although that might mean block-local) and specifically unboxes values stored in locals.

HeapUnboxing or Unboxing both suggest a more global optimization to me.

Untangling certainly suggests that we are making something simpler, but it communicates almost nothing about how.

[kripken]

Hmm, I think what I don't like about unboxing is that it doesn't refer to the core aspect of handling multiple (tangled) allocations. That is, HeapUnboxing could have been the name of Heap2Local.

[tlively]

Yeah, good point that "unboxing" doesn't apply to this any more or less than it would to Heap2Local. Really this is just another version of Heap2Local that does a different escape analysis with different strengths and weaknesses. SimpleHeap2Local???

[kripken]

MultiHeap2Local...?

[tlively]

I wouldn't mind MultiHeap2Local.

[tlively]

Why even bother keeping the reference around for Heap2Local to clean up? Wouldn't it be simple to eliminate the struct.new completely because we know it is immediately set to a local we are unboxing?

[kripken]

Eliminating the struct.new could be done, but we'd need to handle effects and default values and maybe other stuff I'm forgetting (edit: we'd need to handle the descriptor too). It is trivial to just keep the struct.new and do struct.gets from it.

[tlively]

I feel like it would be worth it, though. Right now this pass is useless without a follow-on Heap2Local, but only because it leaves these allocations around. If it removed the allocations, it would stand completely on its own. This would at the very least simplify the tests, but would also let us experiment with running it in more places in the optimization pipeline.

And it should be pretty simple: ChildLocalizer + ref.as_non_null on the descriptor, if it exists.

(Unless there are more complications that I'm not thinking about.)

[kripken]

I let Gemini implement this, and it had to also handle local interactions:

(local.set $x
  (struct.new
    (i32.const 42)
    (struct.get 0
      (local.get $x)
    )
  )
)

=>

(local.set $x_field0
  (i32.const 42)
)
(local.set $x_field1
  (local.get $x_field0) ;; oops, this is 42, but shouldn't be
)

It is possible to detect such short-circuits but it seems messy... I think it's better to leave all that to Heap2Local?

[tlively]

Ok, fair enough. That is more complicated than I was thinking.

[tlively]

It would be good to see if we can get coverage for the remaining uncovered branches as well.

[tlively]

Do all the tests need this if-then-else structure? It's definitely important to show we can handle it, but then for some of these other tests, I think a sequence of (drop (struct.get (local.get $x))) and (struct.set (local.get $x)...) would be sufficient.

I guess this is showing that Heap2Local benefits? Again, I think we can cover that in just a couple tests.

[kripken]

A straightline sequence would have dead code, which I think makes it harder to read. That is, the if-else sequence makes it more realistic. Also more robust to possible future code changes.

[tlively]

How is this different from the previous test with a dropped local.get? Is it just that there's also a local.tee in this one?

[kripken]

This test has locals $yes and $no, and optimizes one and not the other. I'll add to the comment This tests that one local does not interfere with another.

[kripken]

It would be good to see if we can get coverage for the remaining uncovered branches as well.

Where are you seeing that data? Is the new bot running somewhere I can read its output?