[RegAllocFast] insert additional spills along indirect edges of INLINEASM_BR

[nickdesaulniers]

Summary

When generating spills (stores) for values produced by INLINEASM_BR
instructions, make sure to insert one spill per indirect target.
Otherwise the reload generated may load from a stack slot that has not
yet been stored to (resulting in a load of an uninitialized stack slot).

Link: https://github.com/llvm/llvm-project/issues/53562
Fixes: https://github.com/llvm/llvm-project/issues/60855

[nickdesaulniers]

• add livein

[MatzeB]

llvm/lib/CodeGen/RegAllocFast.cpp

952

May be good to add a comment here to remind readers that we have the really unusual behavior of possibly jumping in the middle of a basic block.

952–963

Would it be possible to move this code block into the spill() function. To me it feels like it's just part of the spilling process. You could add a new parameter to pass along MI or a bool MayJump...

956

Just wondering: Can there be multiple INLINEASM_BR within the same block where you would need a check to only deal with the successors of the current INLINEASM_BR?

[MatzeB]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

or given there is only this use here, rename to spillAfter(), add the MI parameter and drop the SpillBefore parameter instead. And then just always compute SpillBefore within the function based on MI...

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

This works:

diff --git a/llvm/lib/CodeGen/RegAllocFast.cpp b/llvm/lib/CodeGen/RegAllocFast.cpp
index bf99b4b4a0ea..b2ef32a2328c 100644
--- a/llvm/lib/CodeGen/RegAllocFast.cpp
+++ b/llvm/lib/CodeGen/RegAllocFast.cpp
@@ -435,6 +435,7 @@ void RegAllocFast::spill(MachineBasicBlock::iterator Before, Register VirtReg,
   LLVM_DEBUG(dbgs() << " to stack slot #" << FI << '\n');
 
   const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
+  MachineBasicBlock *MBB = Before->getParent();// shadow MBB member.
   TII->storeRegToStackSlot(*MBB, Before, AssignedReg, Kill, FI, &RC, TRI,
                            VirtReg);
   ++NumStores;
@@ -954,9 +955,7 @@ void RegAllocFast::defineVirtReg(MachineInstr &MI, unsigned OpNum,
         const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
         for (MachineBasicBlock *Succ : MI.getParent()->successors()) {
           if (Succ->isInlineAsmBrIndirectTarget()) {
-            TII->storeRegToStackSlot(*Succ, Succ->begin(), PhysReg, Kill,
-                FI, &RC, TRI, VirtReg);
-            ++NumStores;
+            spill(Succ->begin(), VirtReg, PhysReg, Kill, LRI->LiveOut);
             Succ->addLiveIn(PhysReg);
           }
         }

[MatzeB]

llvm/lib/CodeGen/RegAllocFast.cpp

957

Probably better to use the getMBBBeginInsertionPoint() function instead of Succ->begin().

[MatzeB]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

Calling into spill() instead of directly using the storeRegToStackSlot callback seems sensible.
Still logically it seems somewhat unfortunate to have part of the spilling logic here and not in the spill() function...

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

We want to repeat the entirety of spill but for multiple destinations when encountering an INLINEASM_BR.

I could outline the body of RegAllocFast::spill into a new method (RegAllocFast::spillTo or something), and have RegAllocFast::spill check for the special case of INLINEASM_BR, calling RegAllocFast::spillTo 1-to-many times.

Perhaps it's clearer that RegAllocFast::defineVirtReg calls RegAllocFast::spill multiple times, rather than one call to RegAllocFast::spill resulting in 1-to-many spills?

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

Actually, I might be able to make RegAllocFast::spill recursive...let me give that a shot.

957

Would I need to do anything with the produced PrologLiveIns ?

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

Calling into spill() instead of directly using the storeRegToStackSlot callback seems sensible.

Looking at how LiveDbgValueMap is used then cleared in RegAllocFast::spill, I'm not sure that calling spill directly is actually better than the current diff (500885)

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

956

You're right, INLINEASM_BR is not a terminator, so it's possible to have more than one within a single MachineBasicBlock. Let me rework this...

[nickdesaulniers]

• add comment, iterate MBB operands rather than parent MBB successors

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

Actually, I might be able to make RegAllocFast::spill recursive...let me give that a shot.

NVM

[MatzeB]

Added some more comments. Note that we should be careful about debug info (in -O0 with fastregalloc the expectation is to have the best debugging experience possible), even though I am not aware of too many ways to test debug info except making sure the MIR output seems reasonable.

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

• It's not that big of a deal to keep it as two spill() calls here if that works better.
• The thing that did look useful to me was the placement of the extra DBG_VALUE when the value is live out. But I just realized that is not even that easy as the LiveOut is only compute for the current block, not the block where the new spill is placed in.

So maybe it would be best to go back to a direct storeRegToStackSlot call and then manually construct a DBG_VALUE immediate behind it to mark the value moving from register to stack?

[MatzeB]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

Sorry for the back and forth. After thinking some more about this, I haven't looked at fastregalloc in a while. Actually I think we likely can skip the DBG_VALUE as that is only necessary to mark transitions between register and stack. However as soon as we spill a value once we point all debug info in other blocks to the stack slot as far as I can tell. Meaning we shouldn't need to specially mark the transition for the extra spill.

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

If I regen the mir test from the initial report with debug info (clang x.c -g -S -o - -fno-asynchronous-unwind-tables -mllvm -stop-before=regallocfast > regallocfast.g.mir), this is the MIR I get from running regallocfast (llc -run-pass=regallocfast regallocfast.g.mir -o - -verify-machineinstrs): https://gist.github.com/nickdesaulniers/a6c234b56218bbcb21aabbd5521d9615

It looks like regardless of my change, the spills in the same block as the INLINEASM_BR are already missing debug-locations? (This is my first time looking at debug info in MIR).

[nickdesaulniers]

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

@MatzeB it looks like even with optimizations enabled we don't have DBG_VALUEs produced in the MIR stream. I spoke with @dblaikie about this briefly, but we couldn't spot what you may have been referring to.

Do you still have concerns about this code? If not, it's probably ready for review.

[MatzeB]

LGTM

llvm/lib/CodeGen/RegAllocFast.cpp

952–963

yeah for this to matter we would need dbg.value calls in the LLVM-IR to get corresponding DBG_VALUE statements tracking debug info... Though given we don't even run mem2reg in a -O0 -g setting we only have the debug.declare stuff relying on the code constantly doing load/store to the stack object corrpsonding to the local variable... I'm not sure how to even get debug.value with default clang then :-/ Maybe manually running mem2reg on the IR works.

952–963

I don't have any concerns left. I initially wondered whether we need to insert a DBG_VALUE similar to what spill() is doing for this newly inserted 2nd spill. But I convinced myself that this is likely not necessary ("Actually I think we likely can skip...") So all good.

(I thought you just wanted to sanity with debug info in general, which is a good thing to do).

[nickdesaulniers]

Thanks for the review!

llvm/lib/CodeGen/RegAllocFast.cpp

void RegAllocFast::defineVirtReg(MachineInstr &MI, unsigned OpNum,
   if (LRI->Reloaded || LRI->LiveOut) {
     if (!MI.isImplicitDef()) {
       MachineBasicBlock::iterator SpillBefore =
           std::next((MachineBasicBlock::iterator)MI.getIterator());
       LLVM_DEBUG(dbgs() << "Spill Reason: LO: " << LRI->LiveOut << " RL: "
                         << LRI->Reloaded << '\n');
       bool Kill = LRI->LastUse == nullptr;
       spill(SpillBefore, VirtReg, PhysReg, Kill, LRI->LiveOut);
+
+      // We need to place additional spills for each indirect destination of an
+      // INLINEASM_BR.
+      if (MI.getOpcode() == TargetOpcode::INLINEASM_BR) {
+        int FI = StackSlotForVirtReg[VirtReg];
+        const TargetRegisterClass &RC = *MRI->getRegClass(VirtReg);
+        for (MachineOperand &MO : MI.operands()) {
+          if (MO.isMBB()) {
+            MachineBasicBlock *Succ = MO.getMBB();
+            TII->storeRegToStackSlot(*Succ, Succ->begin(), PhysReg, Kill,
+                FI, &RC, TRI, VirtReg);
+            ++NumStores;
+            Succ->addLiveIn(PhysReg);
+          }
+        }
+      }
+
       LRI->LastUse = nullptr;
     }
     LRI->LiveOut = false;
     LRI->Reloaded = false;
   }
   if (MI.getOpcode() == TargetOpcode::BUNDLE) {
     BundleVirtRegsMap[VirtReg] = PhysReg;
   }

llvm/test/CodeGen/X86/callbr-asm-outputs-regallocfast.mir

body:             |
   ; CHECK-NEXT:   MOV32mi %stack.1.x, 1, $noreg, 0, $noreg, 6 :: (store (s32) into %ir.x)
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.2.label:
   ; CHECK-NEXT:   renamable $eax = MOV32rm %stack.1.x, 1, $noreg, 0, $noreg :: (load (s32) from %ir.x)
   ; CHECK-NEXT:   RET64 implicit killed $eax
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.3.label.split (machine-block-address-taken, inlineasm-br-indirect-target):
   ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT:   liveins: $eax
   ; CHECK-NEXT: {{  $}}
-  ; FIXME: this is a load from a stack slot that hasn't been stored to yet!
+  ; CHECK-NEXT:   MOV32mr %stack.3, 1, $noreg, 0, $noreg, $eax :: (store (s32) into %stack.3)
   ; CHECK-NEXT:   $eax = MOV32rm %stack.3, 1, $noreg, 0, $noreg :: (load (s32) from %stack.3)
   ; CHECK-NEXT:   MOV32mr %stack.1.x, 1, $noreg, 0, $noreg, killed renamable $eax :: (store (s32) into %ir.x)
   ; CHECK-NEXT:   JMP_1 %bb.2
   bb.0.entry:
     successors: %bb.1(0x40000000), %bb.3(0x40000000)
 
     MOV32mi %stack.0.retval, 1, $noreg, 0, $noreg, 0 :: (store (s32) into %ir.retval)
     MOV32mi %stack.1.x, 1, $noreg, 0, $noreg, 123 :: (store (s32) into %ir.x)