[RISCV] Avoid redundant branch-to-branch when expanding cmpxchg

[asb]

Summary

If the success value of a cmpxchg is used in a branch, the expanded cmpxchg sequence ends up with a redundant branch-to-branch (as the backend atomics expansion happens as late as possible, passes to optimise such cases have already run). This patch identifies this case and avoid it when expanding the cmpxchg.

Note that a similar optimisation is possible for a BEQ on the cmpxchg success value. As it's hard to imagine a case where real-world code may do that, this patch doens't handle that case.

[craig.topper]

llvm/lib/Target/RISCV/RISCVExpandAtomicPseudoInsts.cpp

531

I think this is skipDebugInstructionsForward?

535

Can we avoid an IsMasked parameter by checking MaskReg.isValid()?

540

80 columns

540

Might be a little better as

if (!(ANDOp1 == DestReg && ANDOp2 == MaskReg) &&
    !(ANDOp1 == MaskReg && ANDOp2 == DestReg))

but I'm not sure.

555

80 columns

555

Same Demorgan suggestion from above

583

Register(0) should be Register()

[asb]

Address all review comments (thanks Craig! and aplogies for missing I hadn't clang-formatted).

[asb]

Ping.

[craig.topper]

LGTM

llvm/lib/Target/RISCV/RISCVExpandAtomicPseudoInsts.cpp

bool RISCVExpandAtomicPseudo::expandAtomicMinMaxOp(
   computeAndAddLiveIns(LiveRegs, *LoopHeadMBB);
   computeAndAddLiveIns(LiveRegs, *LoopIfBodyMBB);
   computeAndAddLiveIns(LiveRegs, *LoopTailMBB);
   computeAndAddLiveIns(LiveRegs, *DoneMBB);
 
   return true;
 }
 
+// If a BNE on the cmpxchg comparison result immediately follows the cmpxchg
+// operation, it can be folded into the cmpxchg expansion by
+// modifying the branch within 'LoopHead' (which performs the same
+// comparison). This is a valid transformation because after altering the
+// LoopHead's BNE destination, the BNE following the cmpxchg becomes
+// redundant and and be deleted. In the case of a masked cmpxchg, an
+// appropriate AND and BNE must be matched.
+//
+// On success, returns true and deletes the matching BNE or AND+BNE, sets the
+// LoopHeadBNETarget argument to the target that should be used within the
+// loop head, and removes that block as a successor to MBB.
+bool tryToFoldBNEOnCmpXchgResult(MachineBasicBlock &MBB,
+                                 MachineBasicBlock::iterator MBBI,
+                                 Register DestReg, Register CmpValReg,
+                                 Register MaskReg,
+                                 MachineBasicBlock *&LoopHeadBNETarget) {
+  SmallVector<MachineInstr *> ToErase;
+  auto E = MBB.end();
+  if (MBBI == E)
+    return false;
+  MBBI = skipDebugInstructionsForward(MBBI, E);
+
+  // If we have a masked cmpxchg, match AND dst, DestReg, MaskReg.
+  if (MaskReg.isValid()) {
+    if (MBBI == E || MBBI->getOpcode() != RISCV::AND)
+      return false;
+    Register ANDOp1 = MBBI->getOperand(1).getReg();
+    Register ANDOp2 = MBBI->getOperand(2).getReg();
+    if (!(ANDOp1 == DestReg && ANDOp2 == MaskReg) &&
+        !(ANDOp1 == MaskReg && ANDOp2 == DestReg))
+      return false;
+    // We now expect the BNE to use the result of the AND as an operand.
+    DestReg = MBBI->getOperand(0).getReg();
+    ToErase.push_back(&*MBBI);
+    MBBI = skipDebugInstructionsForward(std::next(MBBI), E);
+  }
+
+  // Match BNE DestReg, MaskReg.
+  if (MBBI == E || MBBI->getOpcode() != RISCV::BNE)
+    return false;
+  Register BNEOp0 = MBBI->getOperand(0).getReg();
+  Register BNEOp1 = MBBI->getOperand(1).getReg();
+  if (!(BNEOp0 == DestReg && BNEOp1 == CmpValReg) &&
+      !(BNEOp0 == CmpValReg && BNEOp1 == DestReg))
+    return false;
+  ToErase.push_back(&*MBBI);
+  LoopHeadBNETarget = MBBI->getOperand(2).getMBB();
+  MBBI = skipDebugInstructionsForward(std::next(MBBI), E);
+  if (MBBI != E)
+    return false;
+
+  MBB.removeSuccessor(LoopHeadBNETarget);
+  for (auto *MI : ToErase)
+    MI->eraseFromParent();
+  return true;
+}
+
 bool RISCVExpandAtomicPseudo::expandAtomicCmpXchg(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, bool IsMasked,
     int Width, MachineBasicBlock::iterator &NextMBBI) {
   MachineInstr &MI = *MBBI;
   DebugLoc DL = MI.getDebugLoc();
   MachineFunction *MF = MBB.getParent();
   auto LoopHeadMBB = MF->CreateMachineBasicBlock(MBB.getBasicBlock());
   auto LoopTailMBB = MF->CreateMachineBasicBlock(MBB.getBasicBlock());
   auto DoneMBB = MF->CreateMachineBasicBlock(MBB.getBasicBlock());
 
+  Register DestReg = MI.getOperand(0).getReg();
+  Register ScratchReg = MI.getOperand(1).getReg();
+  Register AddrReg = MI.getOperand(2).getReg();
+  Register CmpValReg = MI.getOperand(3).getReg();
+  Register NewValReg = MI.getOperand(4).getReg();
+  Register MaskReg = IsMasked ? MI.getOperand(5).getReg() : Register();
+
+  MachineBasicBlock *LoopHeadBNETarget = DoneMBB;
+  tryToFoldBNEOnCmpXchgResult(MBB, std::next(MBBI), DestReg, CmpValReg, MaskReg,
+                              LoopHeadBNETarget);
+
   // Insert new MBBs.
   MF->insert(++MBB.getIterator(), LoopHeadMBB);
   MF->insert(++LoopHeadMBB->getIterator(), LoopTailMBB);
   MF->insert(++LoopTailMBB->getIterator(), DoneMBB);
 
   // Set up successors and transfer remaining instructions to DoneMBB.
   LoopHeadMBB->addSuccessor(LoopTailMBB);
-  LoopHeadMBB->addSuccessor(DoneMBB);
+  LoopHeadMBB->addSuccessor(LoopHeadBNETarget);
   LoopTailMBB->addSuccessor(DoneMBB);
   LoopTailMBB->addSuccessor(LoopHeadMBB);
   DoneMBB->splice(DoneMBB->end(), &MBB, MI, MBB.end());
   DoneMBB->transferSuccessors(&MBB);
   MBB.addSuccessor(LoopHeadMBB);
 
-  Register DestReg = MI.getOperand(0).getReg();
-  Register ScratchReg = MI.getOperand(1).getReg();
-  Register AddrReg = MI.getOperand(2).getReg();
-  Register CmpValReg = MI.getOperand(3).getReg();
-  Register NewValReg = MI.getOperand(4).getReg();
   AtomicOrdering Ordering =
       static_cast<AtomicOrdering>(MI.getOperand(IsMasked ? 6 : 5).getImm());
 
   if (!IsMasked) {
     // .loophead:
     //   lr.[w|d] dest, (addr)
     //   bne dest, cmpval, done
     BuildMI(LoopHeadMBB, DL, TII->get(getLRForRMW(Ordering, Width)), DestReg)
         .addReg(AddrReg);
     BuildMI(LoopHeadMBB, DL, TII->get(RISCV::BNE))
         .addReg(DestReg)
         .addReg(CmpValReg)
-        .addMBB(DoneMBB);
+        .addMBB(LoopHeadBNETarget);
     // .looptail:
     //   sc.[w|d] scratch, newval, (addr)
     //   bnez scratch, loophead
     BuildMI(LoopTailMBB, DL, TII->get(getSCForRMW(Ordering, Width)), ScratchReg)
         .addReg(AddrReg)
         .addReg(NewValReg);
     BuildMI(LoopTailMBB, DL, TII->get(RISCV::BNE))
         .addReg(ScratchReg)
         .addReg(RISCV::X0)
         .addMBB(LoopHeadMBB);
   } else {
     // .loophead:
     //   lr.w dest, (addr)
     //   and scratch, dest, mask
     //   bne scratch, cmpval, done
     Register MaskReg = MI.getOperand(5).getReg();
     BuildMI(LoopHeadMBB, DL, TII->get(getLRForRMW(Ordering, Width)), DestReg)
         .addReg(AddrReg);
     BuildMI(LoopHeadMBB, DL, TII->get(RISCV::AND), ScratchReg)
         .addReg(DestReg)
         .addReg(MaskReg);
     BuildMI(LoopHeadMBB, DL, TII->get(RISCV::BNE))
         .addReg(ScratchReg)
         .addReg(CmpValReg)
-        .addMBB(DoneMBB);
+        .addMBB(LoopHeadBNETarget);
 
     // .looptail:
     //   xor scratch, dest, newval
     //   and scratch, scratch, mask
     //   xor scratch, dest, scratch
     //   sc.w scratch, scratch, (adrr)
     //   bnez scratch, loophead
     insertMaskedMerge(TII, DL, LoopTailMBB, ScratchReg, DestReg, NewValReg,

llvm/test/CodeGen/RISCV/atomic-cmpxchg-branch-on-result.ll

 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:  .LBB0_1: # %do_cmpxchg
 ; CHECK-NEXT:    # =>This Loop Header: Depth=1
 ; CHECK-NEXT:    # Child Loop BB0_3 Depth 2
 ; CHECK-NEXT:  .LBB0_3: # %do_cmpxchg
 ; CHECK-NEXT:    # Parent Loop BB0_1 Depth=1
 ; CHECK-NEXT:    # => This Inner Loop Header: Depth=2
 ; CHECK-NEXT:    lr.w.aqrl a3, (a0)
-; CHECK-NEXT:    bne a3, a1, .LBB0_5
+; CHECK-NEXT:    bne a3, a1, .LBB0_1
 ; CHECK-NEXT:  # %bb.4: # %do_cmpxchg
 ; CHECK-NEXT:    # in Loop: Header=BB0_3 Depth=2
 ; CHECK-NEXT:    sc.w.aqrl a4, a2, (a0)
 ; CHECK-NEXT:    bnez a4, .LBB0_3
-; CHECK-NEXT:  .LBB0_5: # %do_cmpxchg
-; CHECK-NEXT:    # in Loop: Header=BB0_1 Depth=1
-; CHECK-NEXT:    bne a3, a1, .LBB0_1
+; CHECK-NEXT:  # %bb.5: # %do_cmpxchg
 ; CHECK-NEXT:  # %bb.2: # %exit
 ; CHECK-NEXT:    ret
 entry:
   br label %do_cmpxchg
 do_cmpxchg:
   %0 = cmpxchg i32* %ptr, i32 %cmp, i32 %val seq_cst seq_cst
   %1 = extractvalue { i32, i1 } %0, 1
   br i1 %1, label %exit, label %do_cmpxchg

 ; RV32IA-NEXT:  .LBB2_1: # %do_cmpxchg
 ; RV32IA-NEXT:    # =>This Loop Header: Depth=1
 ; RV32IA-NEXT:    # Child Loop BB2_3 Depth 2
 ; RV32IA-NEXT:  .LBB2_3: # %do_cmpxchg
 ; RV32IA-NEXT:    # Parent Loop BB2_1 Depth=1
 ; RV32IA-NEXT:    # => This Inner Loop Header: Depth=2
 ; RV32IA-NEXT:    lr.w.aqrl a4, (a3)
 ; RV32IA-NEXT:    and a5, a4, a0
-; RV32IA-NEXT:    bne a5, a1, .LBB2_5
+; RV32IA-NEXT:    bne a5, a1, .LBB2_1
 ; RV32IA-NEXT:  # %bb.4: # %do_cmpxchg
 ; RV32IA-NEXT:    # in Loop: Header=BB2_3 Depth=2
 ; RV32IA-NEXT:    xor a5, a4, a2
 ; RV32IA-NEXT:    and a5, a5, a0
 ; RV32IA-NEXT:    xor a5, a4, a5
 ; RV32IA-NEXT:    sc.w.aqrl a5, a5, (a3)
 ; RV32IA-NEXT:    bnez a5, .LBB2_3
-; RV32IA-NEXT:  .LBB2_5: # %do_cmpxchg
-; RV32IA-NEXT:    # in Loop: Header=BB2_1 Depth=1
-; RV32IA-NEXT:    and a4, a4, a0
-; RV32IA-NEXT:    bne a1, a4, .LBB2_1
+; RV32IA-NEXT:  # %bb.5: # %do_cmpxchg
 ; RV32IA-NEXT:  # %bb.2: # %exit
 ; RV32IA-NEXT:    ret
 ;
 ; RV64IA-LABEL: cmpxchg_masked_and_branch1:
 ; RV64IA:       # %bb.0: # %entry
 ; RV64IA-NEXT:    andi a3, a0, -4
 ; RV64IA-NEXT:    slliw a4, a0, 3
 ; RV64IA-NEXT:    li a0, 255
 ; RV64IA-NEXT:    sllw a0, a0, a4
 ; RV64IA-NEXT:    andi a1, a1, 255
 ; RV64IA-NEXT:    sllw a1, a1, a4
 ; RV64IA-NEXT:    andi a2, a2, 255
 ; RV64IA-NEXT:    sllw a2, a2, a4
 ; RV64IA-NEXT:  .LBB2_1: # %do_cmpxchg
 ; RV64IA-NEXT:    # =>This Loop Header: Depth=1
 ; RV64IA-NEXT:    # Child Loop BB2_3 Depth 2
 ; RV64IA-NEXT:  .LBB2_3: # %do_cmpxchg
 ; RV64IA-NEXT:    # Parent Loop BB2_1 Depth=1
 ; RV64IA-NEXT:    # => This Inner Loop Header: Depth=2
 ; RV64IA-NEXT:    lr.w.aqrl a4, (a3)
 ; RV64IA-NEXT:    and a5, a4, a0
-; RV64IA-NEXT:    bne a5, a1, .LBB2_5
+; RV64IA-NEXT:    bne a5, a1, .LBB2_1
 ; RV64IA-NEXT:  # %bb.4: # %do_cmpxchg
 ; RV64IA-NEXT:    # in Loop: Header=BB2_3 Depth=2
 ; RV64IA-NEXT:    xor a5, a4, a2
 ; RV64IA-NEXT:    and a5, a5, a0
 ; RV64IA-NEXT:    xor a5, a4, a5
 ; RV64IA-NEXT:    sc.w.aqrl a5, a5, (a3)
 ; RV64IA-NEXT:    bnez a5, .LBB2_3
-; RV64IA-NEXT:  .LBB2_5: # %do_cmpxchg
-; RV64IA-NEXT:    # in Loop: Header=BB2_1 Depth=1
-; RV64IA-NEXT:    and a4, a4, a0
-; RV64IA-NEXT:    bne a1, a4, .LBB2_1
+; RV64IA-NEXT:  # %bb.5: # %do_cmpxchg
 ; RV64IA-NEXT:  # %bb.2: # %exit
 ; RV64IA-NEXT:    ret
 entry:
   br label %do_cmpxchg
 do_cmpxchg:
   %0 = cmpxchg i8* %ptr, i8 %cmp, i8 %val seq_cst seq_cst
   %1 = extractvalue { i8, i1 } %0, 1
   br i1 %1, label %exit, label %do_cmpxchg