[x86] Fixing PR28755 by precomputing the address used in CMPXCHG8B

[aivchenk]

Summary

The bug arises during register allocation on i686 for
CMPXCHG8B instruction when base pointer is needed. CMPXCHG8B
needs 4 implicit registers (EAX, EBX, ECX, EDX) and a memory address,
plus ESI is reserved as the base pointer. With such constraints the only
way register allocator would do its job successfully is when the addressing
mode of the instruction requires only one register. If that is not the case
- we are emitting additional LEA instruction to compute the address

[qcolombet]

lib/Target/X86/X86ISelLowering.cpp

24614

Split the NFC changes in a different patch.

24733

Use an early return to reduce indentation per LLVM coding style.

24753

Shouldn't we check for is32Bit or something?
Indeed, LEA32r would not be valid for 64-bit arch and the helping regalloc shouldn't be needed for 64-bit arch, right?

test/CodeGen/X86/pr28755.ll

1 ↗

(On Diff #73030)

Couple of comment on the test case:

• Give a meaningful name to the file name (like cmpxchg something something)
• Add a comment on what does this test check (in particular the PR# should be help)
• Add check line to make sure we generate something sensible
• Add a run line for when the transformation should not happen (with the appropriate checks as well)

[qcolombet]

+llvm-commit

[aivchenk]

Addressing Quentin's comments:

[aivchenk]

lib/Target/X86/X86ISelLowering.cpp

24614

Uploaded as: https://reviews.llvm.org/D25192

24753

Right. I checked for "TRI->getBaseRegister() == X86::ESI", which is only true for i686, but I agree that it would be better to check directly for 32 bits

[qcolombet]

Looks pretty good, just one nitpick.

test/CodeGen/X86/cmpxchg8b_alloca_regalloc_handling.ll

22

Why to you need the regex here?

[aivchenk]

test/CodeGen/X86/cmpxchg8b_alloca_regalloc_handling.ll

22

You mean in the checks related to the first function? Well, I kinda wanted to express that the register that we are doing "lea" to is the same register we use for a memory location in cmpxchg8b. And since we don't hardcode the name of the register in EmitInstrWithCustomInserter, I did it as a regexp for esi/edi.

Although it makes me wonder whether it is valid or not to define/use regexp with the same name lower in the test

[qcolombet]

test/CodeGen/X86/cmpxchg8b_alloca_regalloc_handling.ll

22

That's fine. My question was regarding the xor or mov. Why do we expect this to change?
We usually pick whatever codegen does, even if both are valid, so that we know when something change. That way we can check it was intended.

[aivchenk]

Ok, got the point. Attached patch represents what currently is generated for those gen's

[qcolombet]

One last nitpick, don't need to go through another round of review, just commit directly.

Cheers,
Q.

lib/Target/X86/X86InstrBuilder.h

129

Could you add a comment on which operand are what?
I admit I don't remember at first when looking at the code why we need 1 for Operand + 1, until I realize, right this is the scale :).

[aivchenk]

thank you a lot for reviewing this! I will ask smb to commit it after I got approve for D24938 - I'm using getAddressFromInstr and it has to work correctly

lib/Target/X86/X86ISelLowering.cpp

case X86::FP80_TO_INT64_IN_MEM: {
     // Change the floating point control register to use "round towards zero"
     // mode when truncating to an integer value.
     int CWFrameIdx = MF->getFrameInfo().CreateStackObject(2, 2, false);
     addFrameReference(BuildMI(*BB, MI, DL,
                               TII->get(X86::FNSTCW16m)), CWFrameIdx);
 
     // Load the old value of the high byte of the control word...
     unsigned OldCW =
       MF->getRegInfo().createVirtualRegister(&X86::GR16RegClass);
     addFrameReference(BuildMI(*BB, MI, DL, TII->get(X86::MOV16rm), OldCW),
                       CWFrameIdx);
 
     // Set the high part to be round to zero...
     addFrameReference(BuildMI(*BB, MI, DL, TII->get(X86::MOV16mi)), CWFrameIdx)
       .addImm(0xC7F);
 
     // Reload the modified control word now...
case TargetOpcode::STATEPOINT:
     // As an implementation detail, STATEPOINT shares the STACKMAP format at
     // this point in the process.  We diverge later.
     return emitPatchPoint(MI, BB);
 
   case TargetOpcode::STACKMAP:
   case TargetOpcode::PATCHPOINT:
     return emitPatchPoint(MI, BB);
 
+  case X86::LCMPXCHG8B: {
+    const X86RegisterInfo *TRI = Subtarget.getRegisterInfo();
+    // In addition to 4 E[ABCD] registers implied by encoding, CMPXCHG8B
+    // requires a memory operand. If it happens that current architecture is
+    // i686 and for current function we need a base pointer
+    // - which is ESI for i686 - register allocator would not be able to
+    // allocate registers for an address in form of X(%reg, %reg, Y)
+    // - there never would be enough unreserved registers during regalloc
+    // (without the need for base ptr the only option would be X(%edi, %esi, Y).
+    // We are giving a hand to register allocator by precomputing the address in
+    // a new vreg using LEA.
+
+    // If it is not i686 or there is no base pointer - nothing to do here.
+    if (!Subtarget.is32Bit() || !TRI->hasBasePointer(*MF))
+      return BB;
+
+    // Even though this code does not necessarily needs the base pointer to
+    // be ESI, we check for that. The reason: if this assert fails, there are
+    // some changes happened in the compiler base pointer handling, which most
+    // probably have to be addressed somehow here.
+    assert(TRI->getBaseRegister() == X86::ESI &&
+           "LCMPXCHG8B custom insertion for i686 is written with X86::ESI as a "
+           "base pointer in mind");
+
+    MachineRegisterInfo &MRI = MF->getRegInfo();
+    MVT SPTy = getPointerTy(MF->getDataLayout());
+    const TargetRegisterClass *AddrRegClass = getRegClassFor(SPTy);
+    unsigned computedAddrVReg = MRI.createVirtualRegister(AddrRegClass);
+
+    X86AddressMode AM = getAddressFromInstr(&MI, 0);
+    // Regalloc does not need any help when the memory operand of CMPXCHG8B
+    // does not use index register.
+    if (AM.IndexReg == X86::NoRegister)
+      return BB;
+
+    // After X86TargetLowering::ReplaceNodeResults CMPXCHG8B is glued to its
+    // four operand definitions that are E[ABCD] registers. We skip them and
+    // then insert the LEA.
+    MachineBasicBlock::iterator MBBI(MI);
+    while (MBBI->definesRegister(X86::EAX) || MBBI->definesRegister(X86::EBX) ||
+           MBBI->definesRegister(X86::ECX) || MBBI->definesRegister(X86::EDX))
+      --MBBI;
+    addFullAddress(
+        BuildMI(*BB, *MBBI, DL, TII->get(X86::LEA32r), computedAddrVReg), AM);
+
+    setDirectAddressInInstr(&MI, 0, computedAddrVReg);
+
+    return BB;
+  }
+  case X86::LCMPXCHG16B:
+    return BB;
   case X86::LCMPXCHG8B_SAVE_EBX:
   case X86::LCMPXCHG16B_SAVE_RBX: {
     unsigned BasePtr =
         MI.getOpcode() == X86::LCMPXCHG8B_SAVE_EBX ? X86::EBX : X86::RBX;
     if (!BB->isLiveIn(BasePtr))
       BB->addLiveIn(BasePtr);
     return BB;
   }

lib/Target/X86/X86InstrBuilder.h

static inline X86AddressMode getAddressFromInstr(const MachineInstr *MI,
   if (Op3.isGlobal())
     AM.GV = Op3.getGlobal();
   else
     AM.Disp = Op3.getImm();
 
   return AM;
 }
 
+// Replace the address used in the instruction with the direct memory reference.
+static inline void setDirectAddressInInstr(MachineInstr *MI, unsigned Operand,
+                                           unsigned Reg) {
+  MI->getOperand(Operand).setReg(Reg);
+  MI->getOperand(Operand + 1).setImm(1);
+  MI->getOperand(Operand + 2).setReg(0);
+  MI->getOperand(Operand + 3).setImm(0);
+  MI->getOperand(Operand + 4).setReg(0);
+}
+
 /// addDirectMem - This function is used to add a direct memory reference to the
 /// current instruction -- that is, a dereference of an address in a register,
 /// with no scale, index or displacement. An example is: DWORD PTR [EAX].
 ///
 static inline const MachineInstrBuilder &
 addDirectMem(const MachineInstrBuilder &MIB, unsigned Reg) {
   // Because memory references are always represented with five
   // values, this adds: Reg, 1, NoReg, 0, NoReg to the instruction.

lib/Target/X86/X86InstrCompiler.td

 
 defm LOCK_INC    : LOCK_ArithUnOp<0xFE, 0xFF, MRM0m,  1, "inc">;
 defm LOCK_DEC    : LOCK_ArithUnOp<0xFE, 0xFF, MRM1m, -1, "dec">;
 
 // Atomic compare and swap.
 multiclass LCMPXCHG_UnOp<bits<8> Opc, Format Form, string mnemonic,
                          SDPatternOperator frag, X86MemOperand x86memop,
                          InstrItinClass itin> {
-let isCodeGenOnly = 1 in {
+let isCodeGenOnly = 1, usesCustomInserter = 1 in {
   def NAME : I<Opc, Form, (outs), (ins x86memop:$ptr),
                !strconcat(mnemonic, "\t$ptr"),
                [(frag addr:$ptr)], itin>, TB, LOCK;
 }
 }
 
 multiclass LCMPXCHG_BinOp<bits<8> Opc8, bits<8> Opc, Format Form,
                           string mnemonic, SDPatternOperator frag,

test/CodeGen/X86/cmpxchg8b_alloca_regalloc_handling.ll

+; RUN: llc < %s -march=x86 -stackrealign -O2 | FileCheck %s
+; PR28755
+
+; Check that register allocator is able to handle that
+; a-lot-of-fixed-and-reserved-registers case. We do that by
+; emmiting lea before 4 cmpxchg8b operands generators.
+
+define void @foo_alloca(i64* %a, i32 %off, i32 %n) {
+  %dummy = alloca i32, i32 %n
+  %addr = getelementptr inbounds i64, i64* %a, i32 %off
+
+  %res = cmpxchg i64* %addr, i64 0, i64 1 monotonic monotonic
+  ret void
+}
+
+; CHECK-LABEL: foo_alloca
+; CHECK: leal    {{\(%e..,%e..,.*\)}}, [[REGISTER:%e.i]]
+; CHECK-NEXT: xorl    %eax, %eax
+; CHECK-NEXT: xorl    %edx, %edx
+; CHECK-NEXT: xorl    %ecx, %ecx
+; CHECK-NEXT: movl    $1, %ebx
+; CHECK-NEXT: lock            cmpxchg8b       ([[REGISTER]])
+
+; If we don't use index register in the address mode -
+; check that we did not generate the lea.
+define void @foo_alloca_direct_address(i64* %addr, i32 %n) {
+  %dummy = alloca i32, i32 %n
+
+  %res = cmpxchg i64* %addr, i64 0, i64 1 monotonic monotonic
+  ret void
+}
+
+; CHECK-LABEL: foo_alloca_direct_address
+; CHECK-NOT: leal    {{\(%e.*\)}}, [[REGISTER:%e.i]]
+; CHECK: lock            cmpxchg8b       ([[REGISTER]])