[WebAssembly] Add comments on local.tee transformation

[aheejin]

Summary

We have a good comment on TEE transformation in RegStackify:
https://github.com/llvm/llvm-project/blob/547e3456660000a16fc5c2a2f819f1a2b5d35b5d/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp#L613-L632

And I think it can be helpful to have some more comments on how the
TEEs created in RegStackify are converted into LOCAL_TEEs.

Variable OldReg is changed to DefReg to be consistent with
RegStackify's comment.

[tlively]

LGTM % language tweaks!

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

276

280

I think the past tense makes it clearer that this refers to the "Before" state.

288

289–290

I know you're just documenting the existing behavior, but wouldn't it be better to remove the tee and map both TeeReg and Reg to LocalId1 without introducing NewReg or LocalId2 at all?

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

289–290

Not sure what you are suggesting. Can you show what the code will look like (in terms of the existing INST, TeeReg, Reg, and such)?

TeeReg is stackified (and not mapped to a local) in both cwhen DefReg was stackified and not stackified). Are you suggesting TeeReg to be mapped to a local here? I don't think something can be mapped to a local and stackified at the same time.

When a reg is mapped to a local, that reg doesn't appear in the MIR code anymore. The reason Reg appears in INST ..., Reg, ... is because we are processing instructions sequentially and haven't processed that instruction yet. When we process each instruction, we replace each mapped reg to a local.

On the other hand, stackified registers (such as TeeReg or NewReg) will still appear in MIR code, but they in effect represent stack pushes and pops, and instructions will be stripped of them in MCInstLower or somewhere. So we create a new register whenever we stackify a register so that it won't clash with existing ones.
(Examples elsewhere:
https://github.com/llvm/llvm-project/blob/aba4e4d6c1c40ea7b8ba793627b62dc83a47a1d0/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp#L308 and
https://github.com/llvm/llvm-project/blob/aba4e4d6c1c40ea7b8ba793627b62dc83a47a1d0/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp#L380)

[aheejin]

Address comments

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

289–290

Ping 😀

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

289–290

Will land this, given that this is more of a suggestion for a follow-up.

[tlively]

Oh oops, I had written a response but then I never clicked "submit." Sorry about that.

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

289–290

Sure, here's the "after" case I have in mind for when DefReg is not stackified. Both TeeReg and Reg have been replaced with DefReg, which will be accessed using local.get instructions inserted later.

Removed: TeeReg, Reg = TEE DefReg
INST ..., DefReg, ...
INST ..., DefReg, ...
INST ..., DefReg, ...

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

289–290

I think, as you said, if we delete the TEE and replace TeeReg and Reg with DefReg, we can save one instruction. The code may look like this:

if (!MFI.isVRegStackified(DefReg)) {
  Register TeeReg = MI.getOperand(0).getReg();
  Register Reg = MI.getOperand(1).getReg();
  MRI.replaceRegWith(TeeReg, DefReg);
  MRI.replaceRegWith(Reg, DefReg);
  MI.eraseFromParent();
  Changed = true;
  continue;
}

But I'm wondering if we can replace a register with another register like this give that we don't have SSA at this point in the pipeline. Also there can be many Regs down the line, so they can be in other BBs, at which point we cannot guarantee those Regs have a unique def. The weird thing is, we are doing this in another part of this pass, which I don't understand: https://github.com/llvm/llvm-project/blob/7f3836150fd1429c20d39560a53ab02dbeaab643/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp#L422-L428

But that part of the code has existed for like 7 years now and hasn't caused any problems so far. Not sure why.

But anyway, for the TEE thing, given that it's not SSA anymore I'm not sure if I can replace a reg with another wholesale. But other that this it's hard to replace all Regs down the line, because we don't have SSA and thus no replaceAllUsesWith kind of thing.

But I think this, meaning DefReg unstackified by this point, currently happens in a very limited situation when we are making try-delegate and DefReg happens to be a stackified register in the boundary of newly splitted BBs so we unstackify it: https://github.com/llvm/llvm-project/blob/488185cca3871a0ef2ec3b9b4c642dc6db6eeea5/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp#L799

I think this case is rare enough, and not fixing this wouldn't make the code size meaningfully worse.
Also the weird thing is, this part of the code in ExplicitLocals has existed before that try-delegate handling, but I can't find anywhere we unstackify a register between RegStackify and here other that that. Not sure what was in Dan's mind 7(?) years ago at this point.

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

for (MachineInstr &MI : llvm::make_early_inc_range(MBB)) {
         MI.eraseFromParent();
         Changed = true;
         continue;
       }
 
       // Replace tee instructions with local.tee. The difference is that tee
       // instructions have two defs, while local.tee instructions have one def
       // and an index of a local to write to.
+      //
+      // - Before:
+      // TeeReg, Reg = TEE DefReg
+      // INST ..., TeeReg, ...
+      // INST ..., Reg, ...
+      // INST ..., Reg, ...
+      // * DefReg: may or may not be stackified
+      // * Reg: not stackified
+      // * TeeReg: stackified
+      //
+      // - After (when DefReg was already stackified):
+      // TeeReg = LOCAL_TEE LocalId1, DefReg
+      // INST ..., TeeReg, ...
+      // INST ..., Reg, ...
+      // INST ..., Reg, ...
+      // * Reg: mapped to LocalId1
+      // * TeeReg: stackified
+      //
+      // - After (when DefReg was not already stackified):
+      // NewReg = LOCAL_GET LocalId1
+      // TeeReg = LOCAL_TEE LocalId2, NewReg
+      // INST ..., TeeReg, ...
+      // INST ..., Reg, ...
+      // INST ..., Reg, ...
+      // * DefReg: mapped to LocalId1
+      // * Reg: mapped to LocalId2
+      // * TeeReg: stackified
       if (WebAssembly::isTee(MI.getOpcode())) {
         assert(MFI.isVRegStackified(MI.getOperand(0).getReg()));
         assert(!MFI.isVRegStackified(MI.getOperand(1).getReg()));
-        Register OldReg = MI.getOperand(2).getReg();
-        const TargetRegisterClass *RC = MRI.getRegClass(OldReg);
+        Register DefReg = MI.getOperand(2).getReg();
+        const TargetRegisterClass *RC = MRI.getRegClass(DefReg);
 
         // Stackify the input if it isn't stackified yet.
-        if (!MFI.isVRegStackified(OldReg)) {
-          unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, OldReg);
+        if (!MFI.isVRegStackified(DefReg)) {
+          unsigned LocalId = getLocalId(Reg2Local, MFI, CurLocal, DefReg);
           Register NewReg = MRI.createVirtualRegister(RC);
           unsigned Opc = getLocalGetOpcode(RC);
           BuildMI(MBB, &MI, MI.getDebugLoc(), TII->get(Opc), NewReg)
               .addImm(LocalId);
           MI.getOperand(2).setReg(NewReg);
           MFI.stackifyVReg(MRI, NewReg);
         }