[CodeGen] Deduplicate restore blocks in branch relaxation

[piggynl]

Summary

For multiple far branches to the same destination, there are chances that some restore blocks could be shared if they clobber the same registers and share the same restore sequence. Since RISCV generates a same restore block consisting only ld s11, 0(sp) for all jumps to one destination (see D130560), deduplicating them would save a few bytes in that case.

This patch stores all restore blocks for a destination in RestoreBlocks in the destination's BasicBlockInfo, allowing the target to call DeduplicateRestoreBB() to find a previously inserted restore block identical to RestoreBB and insert an unconditional branch to it, then BranchRelaxation::fixupUnconditionalBranch() can erase RestoreBB from the function.

Suppose jump .dest is out of range so need a register to jump indirectly in the following code.

.label_0:
        foo     1
        jump    .dest
.label_1:
        foo     2
        jump    .dest
.label_2:
        bar
        bar
        bar
        bar
.dest:
        baz

Without this patch, it will be relaxed to:

.label_0:
        foo     1
        store   reg1, (sp)
        mov     reg1, .restore_1
        jump    reg1
.label_1:
        foo     2
        store   reg1, (sp)
        mov     reg1, .restore_2
        jump    reg1
.label_2:
        bar
        bar
        bar
        bar
.restore_1:
        load    reg1, (sp)
        jump    .dest
.restore_2:
        load    reg1, (sp)
.dest:
        baz

With this patch, restore blocks can be deduplicated, thus generating the following code:

.label_0:
        foo
        store   reg1, (sp)
        mov     reg1, .restore_1
        jump    reg1
.label_1:
        foo
        store   reg1, (sp)
        mov     reg1, .restore_1
        jump    reg1
.label_2:
        bar
        bar
        bar
        bar
.restore_1:
        load    reg1, (sp)
.dest:
        baz

Depends on D131865, which pre-commits the modified test.

[piggynl]

llvm/lib/CodeGen/BranchRelaxation.cpp

121

It is possible to use SmallSet instead of SmallVector for RestoreBlocks. It will enable this function to find the reusable block in O(log N) time, but It will take a lot of work. We have to implement a comparator for MachineBasicBlock like FunctionComparator. Also, we cannot reuse MachineInstr::isIdenticalTo() and MachineOperand::isIdenticalTo() because they can only return the equality.

I have no idea how frequently is register spill required for branch relaxation in AMDGPU, but in RISCV, it almost never happens. Also, since RISCV generates a same restore block consisting only ld s11, 0(sp) for all jumps to one destination (see D130560), there should be no difference for RISCV to use SmallVector or SmallSet.

llvm/test/CodeGen/AMDGPU/branch-relax-spill.ll

2036 ↗

(On Diff #451516)

I have no knowledge of AMDGPU and don't know how to construct some branches that can share or cann't share a same restore block, so please kindly let me know how to construct the test properly.

This is the first added function. I copied @spill() in this file and simply added another branch after the previous one. The two branches look to share the same restore block .LBB2_5.

2998 ↗

(On Diff #451516)

This is the second added function. I copied @spill_func() in this file and simply added another branch after the previous one. The two branches look don't have a same restore block.

3405–3415 ↗

(On Diff #451516)

and I don't know why this sequence is forming. Is this intended?

[piggynl]

• Rebase onto D130560.

• Add test for RISCV.

[craig.topper]

llvm/lib/CodeGen/BranchRelaxation.cpp

68

How often are there RestoreBlocks? I wonder if this makes sense to use TinyPtrVector so that we don't enlarge the BasicBlockInfo as much. This SmallVector is probably at least 24 bytes?

[craig.topper]

llvm/test/CodeGen/RISCV/branch-relaxation.ll

2653

Pre-commit the test?

[piggynl]

llvm/test/CodeGen/RISCV/branch-relaxation.ll

2787

I think this %dest_2 should be %dest_1. Creating D131863 to fix this. I'll pre-commit this @relax_jal_spill_32_deduplicate_restore_block() test based on D131863.

[piggynl]

• Change the type of BasicBlockInfo::RestoreBlocks from SmallVector to TinyPtrVector.

• Rebase onto pre-committed test D131865.

[piggynl]

llvm/test/CodeGen/AMDGPU/branch-relax-spill.ll

2036 ↗

(On Diff #451516)

Moved these comments about tests for AMDGPU to D131865.

[piggynl]

Ping for review :)

[piggynl]

Ping for review :)

[arsenm]

When would an inserted restore block ever not be identical with the same destination?

llvm/lib/CodeGen/BranchRelaxation.cpp

94

This loop seems overcomplicated for a comparison. Can you just do a range loop, comparing identical instructions until the first terminator?

98

s/has/have/

105

Does this also need to consider the various flags on the block (e.g. alignment)?

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

1002

Why is this a callback? Why doesn't the pass just deduplicate to begin with and pass in the new block?

[piggynl]

• Rebase

• Fix typo

• Simplify the loop for the comparison in BasicBlockInfo::isRestoreBlockIdentical()

[piggynl]

Currently SIInstrinfo::insertIndirectBranch do something like this:

1. Create a new virtual register PCReg and two symbols OffsetLo and OffsetHi.

2. Before the branch instruction, insert instructions to calculate the branch offset with PCReg, OffsetLo, and OffsetHi.

3. Scavenge a free SGPR, or spill an SGPR to VGPR to make it free.

4. Replace PCReg with the free SGPR.

5. Set the values of OffsetLo and OffsetHi.

Only when the spill happens, a restore block will be used and should be deduplicated.

In D131587#3792180, @arsenm wrote:

When would an inserted restore block ever not be identical with the same destination?

The restore block will not be identical when the used SGPR or VGPR is different.

llvm/lib/CodeGen/BranchRelaxation.cpp

105

No target is constructing the restore block with flagged instructions at this time. I've added a FIXME here.

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

1002

In TargetInstrinfo::insertIndirectBranch(), targets first optionally build the restore block, then build the indirect branch points to the restore block. The deduplication is only possible in the middle of the two processes. Both two processes may be heavy (for example SIInstrinfo::insertIndirectBranch()). I think splitting insertIndirectBranch() into two functions is not a good idea because it will increase the complexity for targets doesn't need any restore block.

[piggynl]

Update code and tests for LoongArch.

[piggynl]

Ping for review :)

[foad]

llvm/include/llvm/CodeGen/TargetInstrInfo.h

607–612

The new argument needs documentation.

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

2558

Why is this removed?

[piggynl]

Add a comment about parameter DeduplicateRestoreBB of function insertIndirectBranch.

[piggynl]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

2558

Because D106449 has already addressed this case.

[foad]

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

2558

Then please commit it separately.

[piggynl]

Split irrelevant comments change to 8fccdfa.

[piggynl]

Hi, is there anything else I can improve?

[qcolombet]

llvm/lib/CodeGen/BranchRelaxation.cpp

105

Could we bail out on that or at least assert?

[dhoekwater]

llvm/include/llvm/CodeGen/TargetInstrInfo.h

607–612

Is it necessary to modify the interface of this function in this way? Since you're delegating the "is the new restore block identical to any old ones?" logic to the caller anyway, why not leave the caller responsible for deduplicating blocks?

If you were to add a new virtual void retargetIndirectBranch(MachineBasicBlock &MBB, MachineBasicBlock &NewDestBB) hook to TargetInstrInfo, you could keep the deduplication logic one level up by doing something like:

insertIndirectBranch(...);
if (!restoreBlock->empty() && isDuplicate(restoreBlock))
  retargetIndirectBranch(MBB, duplicateRestoreBlock);

llvm/include/llvm/CodeGen/TargetInstrInfo.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_TARGETINSTRINFO_H
 #define LLVM_CODEGEN_TARGETINSTRINFO_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/Uniformity.h"
 #include "llvm/CodeGen/MIRFormatter.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineOutliner.h"
public:
   virtual MachineBasicBlock *getBranchDestBlock(const MachineInstr &MI) const {
     llvm_unreachable("target did not implement");
   }
 
   /// Insert an unconditional indirect branch at the end of \p MBB to \p
   /// NewDestBB. Optionally, insert the clobbered register restoring in \p
   /// RestoreBB. \p BrOffset indicates the offset of \p NewDestBB relative to
   /// the offset of the position to insert the new branch.
-  virtual void insertIndirectBranch(MachineBasicBlock &MBB,
-                                    MachineBasicBlock &NewDestBB,
-                                    MachineBasicBlock &RestoreBB,
-                                    const DebugLoc &DL, int64_t BrOffset = 0,
-                                    RegScavenger *RS = nullptr) const {
+  virtual void
+  insertIndirectBranch(MachineBasicBlock &MBB, MachineBasicBlock &NewDestBB,
+                       MachineBasicBlock *RestoreBB, const DebugLoc &DL,
+                       function_ref<MachineBasicBlock *(MachineBasicBlock *)>
+                           DeduplicateRestoreBB,
+                       int64_t BrOffset = 0, RegScavenger *RS = nullptr) const {
     llvm_unreachable("target did not implement");
   }
 
   /// Analyze the branching code at the end of MBB, returning
   /// true if it cannot be understood (e.g. it's a switch dispatch or isn't
   /// implemented for a target).  Upon success, this returns false and returns
   /// with the following information in various cases:
   ///

llvm/lib/CodeGen/BranchRelaxation.cpp

 //===- BranchRelaxation.cpp -----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"

 STATISTIC(NumConditionalRelaxed, "Number of conditional branches relaxed");
 STATISTIC(NumUnconditionalRelaxed, "Number of unconditional branches relaxed");
 
 #define BRANCH_RELAX_NAME "Branch relaxation pass"
 
 namespace {
 
 class BranchRelaxation : public MachineFunctionPass {
-  /// BasicBlockInfo - Information about the offset and size of a single
-  /// basic block.
+  /// BasicBlockInfo - Information of a single basic block.
   struct BasicBlockInfo {
     /// Offset - Distance from the beginning of the function to the beginning
     /// of this basic block.
     ///
     /// The offset is always aligned as required by the basic block.
     unsigned Offset = 0;
 
     /// Size - Size of the basic block in bytes.  If the block contains
     /// inline assembly, this is a worst case estimate.
     ///
     /// The size does not include any alignment padding whether from the
     /// beginning of the block, or from an aligned jump table at the end.
     unsigned Size = 0;
 
+    /// RestoreBlocks - Restore blocks for this basic block when this block
+    /// is a branch destination.
+    ///
+    /// This is used to deduplicate restore blocks. See deduplicateRestoreBlock
+    /// and fixupUnconditionalBranch for details.
+    TinyPtrVector<MachineBasicBlock *> RestoreBlocks;
+
     BasicBlockInfo() = default;
 
     /// Compute the offset immediately following this block. \p MBB is the next
     /// block.
     unsigned postOffset(const MachineBasicBlock &MBB) const {
       const unsigned PO = Offset + Size;
       const Align Alignment = MBB.getAlignment();
       const Align ParentAlign = MBB.getParent()->getAlignment();
       if (Alignment <= ParentAlign)
         return alignTo(PO, Alignment);
 
       // The alignment of this MBB is larger than the function's alignment, so we
       // can't tell whether or not it will insert nops. Assume that it will.
       return alignTo(PO, Alignment) + Alignment.value() - ParentAlign.value();
     }
+
+    /// Returns true if \p OldRestoreBB and \p NewRestoreBB are identical
+    /// restore blocks.
+    static bool isRestoreBlockIdentical(const MachineBasicBlock *OldRestoreBB,
+                                        const MachineBasicBlock *NewRestoreBB) {
+
+      auto ItNew = NewRestoreBB->begin();
+      auto ItNewEnd = NewRestoreBB->end();
+
+      for (auto &ItOld : *OldRestoreBB) {
+        // Since target shouldn't place terminators in the restore block, the
+        // only possible terminator is a unconditional branch is placed by
+        // fixupUnconditionalBranch when there are multiple different restore
+        // blocks. All but the last restore blocks have an unconditional branch
+        // to DestBB when they act as PrevBB. This branch doesn't take part in
+        // the check of whether two restore blocks are identical.
+        if (ItOld.isTerminator())
+          break;
+        // FIXME: consider the various flags on the block (e.g. alignment)
+        if (ItNew == ItNewEnd || !ItOld.isIdenticalTo(*ItNew))
+          return false;
+        ++ItNew;
+      }
+
+      return ItNew == ItNewEnd;
+    }
+
+    /// Iterate over previously inserted restore blocks, returns the one that
+    /// can be reused. That is target can insert an unconditional branch
+    /// to it and fixupUnconditionalBranch can erase \p RestoreBB from
+    /// the function. If no reusable restore block is found, insert \p RestoreBB
+    /// to RestoreBlocks and returns it.
+    MachineBasicBlock *deduplicateRestoreBlock(MachineBasicBlock *RestoreBB) {
+      assert(!RestoreBB->empty() &&
+             "RestoreBB to be deduplicated should not be empty");
+      for (auto &DeduplicatedBB : RestoreBlocks)
+        if (isRestoreBlockIdentical(DeduplicatedBB, RestoreBB))
+          return DeduplicatedBB;
+
+      RestoreBlocks.push_back(RestoreBB);
+      return RestoreBB;
+    }
   };
 
   SmallVector<BasicBlockInfo, 16> BlockInfo;
   std::unique_ptr<RegScavenger> RS;
   LivePhysRegs LiveRegs;
 
   MachineFunction *MF;
   const TargetRegisterInfo *TRI;
bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
   MI.eraseFromParent();
 
   // Create the optional restore block and, initially, place it at the end of
   // function. That block will be placed later if it's used; otherwise, it will
   // be erased.
   MachineBasicBlock *RestoreBB = createNewBlockAfter(MF->back(),
                                                      DestBB->getBasicBlock());
 
-  TII->insertIndirectBranch(*BranchBB, *DestBB, *RestoreBB, DL,
-                            DestOffset - SrcOffset, RS.get());
+  // For multiple far branches to the same destination, there are chances
+  // that some restore blocks could be shared if they clobber the same
+  // registers and share the same restore sequence. We allow target to call
+  // DeduplicateRestoreBB to find a previously inserted restore block
+  // identical to RestoreBB and insert an unconditional branch to it,
+  // then we can erase RestoreBB from the function.
+  BasicBlockInfo &DestBBInfo = BlockInfo[DestBB->getNumber()];
+  MachineBasicBlock *DeduplicatedBB = RestoreBB;
+  auto DeduplicateRestoreBB = [&DestBBInfo,
+                               &DeduplicatedBB](MachineBasicBlock *RestoreBB) {
+    return DeduplicatedBB = DestBBInfo.deduplicateRestoreBlock(RestoreBB);
+  };
+
+  TII->insertIndirectBranch(*BranchBB, *DestBB, RestoreBB, DL,
+                            DeduplicateRestoreBB, DestOffset - SrcOffset,
+                            RS.get());
 
   BlockInfo[BranchBB->getNumber()].Size = computeBlockSize(*BranchBB);
   adjustBlockOffsets(*MBB);
 
+  if (DeduplicatedBB != RestoreBB)
+    BranchBB->replaceSuccessor(DestBB, DeduplicatedBB);
+  if (DeduplicatedBB != RestoreBB || RestoreBB->empty()) {
+    // Remove restore block if it is deduplicated or is not required.
+    MF->erase(RestoreBB);
+    return true;
+  }
+
+  for (auto &RestoreI : *RestoreBB) {
+    assert(!RestoreI.isBranch() &&
+           "Restore blocks should not contain branch instruction");
+  }
+
   // If RestoreBB is required, try to place just before DestBB.
-  if (!RestoreBB->empty()) {
-    // TODO: For multiple far branches to the same destination, there are
-    // chances that some restore blocks could be shared if they clobber the
-    // same registers and share the same restore sequence. So far, those
-    // restore blocks are just duplicated for each far branch.
-    assert(!DestBB->isEntryBlock());
-    MachineBasicBlock *PrevBB = &*std::prev(DestBB->getIterator());
-    // Fall through only if PrevBB has no unconditional branch as one of its
-    // terminators.
-    if (auto *FT = PrevBB->getLogicalFallThrough()) {
-      assert(FT == DestBB);
-      TII->insertUnconditionalBranch(*PrevBB, FT, DebugLoc());
-      BlockInfo[PrevBB->getNumber()].Size = computeBlockSize(*PrevBB);
-    }
-    // Now, RestoreBB could be placed directly before DestBB.
-    MF->splice(DestBB->getIterator(), RestoreBB->getIterator());
-    // Update successors and predecessors.
-    RestoreBB->addSuccessor(DestBB);
-    BranchBB->replaceSuccessor(DestBB, RestoreBB);
-    if (TRI->trackLivenessAfterRegAlloc(*MF))
-      computeAndAddLiveIns(LiveRegs, *RestoreBB);
-    // Compute the restore block size.
-    BlockInfo[RestoreBB->getNumber()].Size = computeBlockSize(*RestoreBB);
-    // Update the offset starting from the previous block.
-    adjustBlockOffsets(*PrevBB);
-  } else {
-    // Remove restore block if it's not required.
-    MF->erase(RestoreBB);
-  }
+
 
   return true;
 }
 
 bool BranchRelaxation::relaxBranchInstructions() {
   bool Changed = false;
 
   // Relaxing branches involves creating new basic blocks, so re-eval

llvm/lib/Target/AMDGPU/SIInstrInfo.h

bool isBranchOffsetInRange(unsigned BranchOpc,
                              int64_t BrOffset) const override;
 
   MachineBasicBlock *getBranchDestBlock(const MachineInstr &MI) const override;
 
   /// Return whether the block terminate with divergent branch.
   /// Note this only work before lowering the pseudo control flow instructions.
   bool hasDivergentBranch(const MachineBasicBlock *MBB) const;
 
-  void insertIndirectBranch(MachineBasicBlock &MBB,
-                            MachineBasicBlock &NewDestBB,
-                            MachineBasicBlock &RestoreBB, const DebugLoc &DL,
-                            int64_t BrOffset, RegScavenger *RS) const override;
+  void
+  insertIndirectBranch(MachineBasicBlock &MBB, MachineBasicBlock &NewDestBB,
+                       MachineBasicBlock *RestoreBB, const DebugLoc &DL,
+                       function_ref<MachineBasicBlock *(MachineBasicBlock *)>
+                           DeduplicateRestoreBB,
 
   bool analyzeBranchImpl(MachineBasicBlock &MBB,
                          MachineBasicBlock::iterator I,
                          MachineBasicBlock *&TBB,
                          MachineBasicBlock *&FBB,
                          SmallVectorImpl<MachineOperand> &Cond,
                          bool AllowModify) const;
 

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

for (const MachineInstr &MI : MBB->terminators()) {
     if (MI.getOpcode() == AMDGPU::SI_NON_UNIFORM_BRCOND_PSEUDO ||
         MI.getOpcode() == AMDGPU::SI_IF || MI.getOpcode() == AMDGPU::SI_ELSE ||
         MI.getOpcode() == AMDGPU::SI_LOOP)
       return true;
   }
   return false;
 }
 
-void SIInstrInfo::insertIndirectBranch(MachineBasicBlock &MBB,
-                                       MachineBasicBlock &DestBB,
-                                       MachineBasicBlock &RestoreBB,
-                                       const DebugLoc &DL, int64_t BrOffset,
-                                       RegScavenger *RS) const {
+void SIInstrInfo::insertIndirectBranch(
+    MachineBasicBlock &MBB, MachineBasicBlock &DestBB,
+    MachineBasicBlock *RestoreBB, const DebugLoc &DL,
+    function_ref<MachineBasicBlock *(MachineBasicBlock *)> DeduplicateRestoreBB,
+    int64_t BrOffset, RegScavenger *RS) const {
   assert(RS && "RegScavenger required for long branching");
   assert(MBB.empty() &&
          "new block should be inserted for expanding unconditional branch");
   assert(MBB.pred_size() == 1);
-  assert(RestoreBB.empty() &&
+  assert(RestoreBB->empty() &&
          "restore block should be inserted for restoring clobbered registers");
 
   MachineFunction *MF = MBB.getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();
 
   // FIXME: Virtual register workaround for RegScavenger not working with empty
   // blocks.
   Register PCReg = MRI.createVirtualRegister(&AMDGPU::SReg_64RegClass);
BuildMI(MBB, I, DL, get(AMDGPU::S_ADDC_U32))
       .addReg(PCReg, RegState::Define, AMDGPU::sub1)
       .addReg(PCReg, 0, AMDGPU::sub1)
       .addSym(OffsetHi, MO_FAR_BRANCH_OFFSET);
 
   // Insert the indirect branch after the other terminator.
   BuildMI(&MBB, DL, get(AMDGPU::S_SETPC_B64))
     .addReg(PCReg);
 
-  // FIXME: If spilling is necessary, this will fail because this scavenger has
-  // no emergency stack slots. It is non-trivial to spill in this situation,
-  // because the restore code needs to be specially placed after the
-  // jump. BranchRelaxation then needs to be made aware of the newly inserted
-  // block.
-  //
   // If a spill is needed for the pc register pair, we need to insert a spill
   // restore block right before the destination block, and insert a short branch
   // into the old destination block's fallthrough predecessor.
   // e.g.:
   //
   // s_cbranch_scc0 skip_long_branch:
   //
   // long_branch_bb:
void SIInstrInfo::insertIndirectBranch(
   //
   // restore_bb:
   //  restore s[8:9]
   //  fallthrough dest_bb
   ///
   // dest_bb:
   //   buzz;
 
+  const MachineBasicBlock *BranchDest;
+
   RS->enterBasicBlockEnd(MBB);
   Register Scav = RS->scavengeRegisterBackwards(
       AMDGPU::SReg_64RegClass, MachineBasicBlock::iterator(GetPC),
       /* RestoreAfter */ false, 0, /* AllowSpill */ false);
   if (Scav) {
     RS->setRegUsed(Scav);
     MRI.replaceRegWith(PCReg, Scav);
     MRI.clearVirtRegs();
+    BranchDest = &DestBB;
   } else {
     // As SGPR needs VGPR to be spilled, we reuse the slot of temporary VGPR for
     // SGPR spill.
     const GCNSubtarget &ST = MF->getSubtarget<GCNSubtarget>();
     const SIRegisterInfo *TRI = ST.getRegisterInfo();
-    TRI->spillEmergencySGPR(GetPC, RestoreBB, AMDGPU::SGPR0_SGPR1, RS);
+
+    TRI->spillEmergencySGPR(GetPC, *RestoreBB, AMDGPU::SGPR0_SGPR1, RS);
+    BranchDest = DeduplicateRestoreBB(RestoreBB);
+
     MRI.replaceRegWith(PCReg, AMDGPU::SGPR0_SGPR1);
     MRI.clearVirtRegs();
   }
 
-  MCSymbol *DestLabel = Scav ? DestBB.getSymbol() : RestoreBB.getSymbol();
   // Now, the distance could be defined.
   auto *Offset = MCBinaryExpr::createSub(
-      MCSymbolRefExpr::create(DestLabel, MCCtx),
+      MCSymbolRefExpr::create(BranchDest->getSymbol(), MCCtx),
       MCSymbolRefExpr::create(PostGetPCLabel, MCCtx), MCCtx);
   // Add offset assignments.
   auto *Mask = MCConstantExpr::create(0xFFFFFFFFULL, MCCtx);
   OffsetLo->setVariableValue(MCBinaryExpr::createAnd(Offset, Mask, MCCtx));
   auto *ShAmt = MCConstantExpr::create(32, MCCtx);
   OffsetHi->setVariableValue(MCBinaryExpr::createAShr(Offset, ShAmt, MCCtx));
 }
 

llvm/lib/Target/AVR/AVRInstrInfo.h

public:
   bool
   reverseBranchCondition(SmallVectorImpl<MachineOperand> &Cond) const override;
 
   MachineBasicBlock *getBranchDestBlock(const MachineInstr &MI) const override;
 
   bool isBranchOffsetInRange(unsigned BranchOpc,
                              int64_t BrOffset) const override;
 
-  void insertIndirectBranch(MachineBasicBlock &MBB,
-                            MachineBasicBlock &NewDestBB,
-                            MachineBasicBlock &RestoreBB, const DebugLoc &DL,
-                            int64_t BrOffset, RegScavenger *RS) const override;
+  void
+  insertIndirectBranch(MachineBasicBlock &MBB, MachineBasicBlock &NewDestBB,
+                       MachineBasicBlock *RestoreBB, const DebugLoc &DL,
+                       function_ref<MachineBasicBlock *(MachineBasicBlock *)>
+                           DeduplicateRestoreBB,
 
 private:
   const AVRRegisterInfo RI;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_AVR_INSTR_INFO_H

llvm/lib/Target/AVR/AVRInstrInfo.cpp

bool AVRInstrInfo::isBranchOffsetInRange(unsigned BranchOp,
   case AVR::BRMIk:
   case AVR::BRPLk:
   case AVR::BRGEk:
   case AVR::BRLTk:
     return isIntN(7, BrOffset);
   }
 }
 
-void AVRInstrInfo::insertIndirectBranch(MachineBasicBlock &MBB,
-                                        MachineBasicBlock &NewDestBB,
-                                        MachineBasicBlock &RestoreBB,
-                                        const DebugLoc &DL, int64_t BrOffset,
-                                        RegScavenger *RS) const {
+void AVRInstrInfo::insertIndirectBranch(
+    MachineBasicBlock &MBB, MachineBasicBlock &NewDestBB,
+    MachineBasicBlock *RestoreBB, const DebugLoc &DL,
+    function_ref<MachineBasicBlock *(MachineBasicBlock *)> DeduplicateRestoreBB,
+    int64_t BrOffset, RegScavenger *RS) const {
   // This method inserts a *direct* branch (JMP), despite its name.
   // LLVM calls this method to fixup unconditional branches; it never calls
   // insertBranch or some hypothetical "insertDirectBranch".
   // See lib/CodeGen/RegisterRelaxation.cpp for details.
   // We end up here when a jump is too long for a RJMP instruction.
   BuildMI(&MBB, DL, get(AVR::JMPk)).addMBB(&NewDestBB);
 }
 
 } // end of namespace llvm

llvm/lib/Target/RISCV/RISCVInstrInfo.h

bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                      SmallVectorImpl<MachineOperand> &Cond,
                      bool AllowModify) const override;
 
   unsigned insertBranch(MachineBasicBlock &MBB, MachineBasicBlock *TBB,
                         MachineBasicBlock *FBB, ArrayRef<MachineOperand> Cond,
                         const DebugLoc &dl,
                         int *BytesAdded = nullptr) const override;
 
-  void insertIndirectBranch(MachineBasicBlock &MBB,
-                            MachineBasicBlock &NewDestBB,
-                            MachineBasicBlock &RestoreBB, const DebugLoc &DL,
-                            int64_t BrOffset, RegScavenger *RS) const override;
+  void
+  insertIndirectBranch(MachineBasicBlock &MBB, MachineBasicBlock &NewDestBB,
+                       MachineBasicBlock *RestoreBB, const DebugLoc &DL,
+                       function_ref<MachineBasicBlock *(MachineBasicBlock *)>
+                           DeduplicateRestoreBB,
 
   unsigned removeBranch(MachineBasicBlock &MBB,
                         int *BytesRemoved = nullptr) const override;
 
   bool
   reverseBranchCondition(SmallVectorImpl<MachineOperand> &Cond) const override;
 
   MachineBasicBlock *getBranchDestBlock(const MachineInstr &MI) const override;

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

 #include "RISCVMachineFunctionInfo.h"
 #include "RISCVSubtarget.h"
 #include "RISCVTargetMachine.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveVariables.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineCombinerPattern.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/MC/MCInstBuilder.h"
 #include "llvm/MC/TargetRegistry.h"
unsigned RISCVInstrInfo::insertBranch(
 
   // Two-way conditional branch.
   MachineInstr &MI = *BuildMI(&MBB, DL, get(RISCV::PseudoBR)).addMBB(FBB);
   if (BytesAdded)
     *BytesAdded += getInstSizeInBytes(MI);
   return 2;
 }
 
-void RISCVInstrInfo::insertIndirectBranch(MachineBasicBlock &MBB,
-                                          MachineBasicBlock &DestBB,
-                                          MachineBasicBlock &RestoreBB,
-                                          const DebugLoc &DL, int64_t BrOffset,
-                                          RegScavenger *RS) const {
+void RISCVInstrInfo::insertIndirectBranch(
+    MachineBasicBlock &MBB, MachineBasicBlock &DestBB,
+    MachineBasicBlock *RestoreBB, const DebugLoc &DL,
+    function_ref<MachineBasicBlock *(MachineBasicBlock *)> DeduplicateRestoreBB,
+    int64_t BrOffset, RegScavenger *RS) const {
   assert(RS && "RegScavenger required for long branching");
   assert(MBB.empty() &&
          "new block should be inserted for expanding unconditional branch");
   assert(MBB.pred_size() == 1);
-  assert(RestoreBB.empty() &&
+  assert(RestoreBB->empty() &&
          "restore block should be inserted for restoring clobbered registers");
 
   MachineFunction *MF = MBB.getParent();
   MachineRegisterInfo &MRI = MF->getRegInfo();
   RISCVMachineFunctionInfo *RVFI = MF->getInfo<RISCVMachineFunctionInfo>();
   const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
 
   if (!isInt<32>(BrOffset))
     report_fatal_error(
         "Branch offsets outside of the signed 32-bit range not supported");
 
   // FIXME: A virtual register must be used initially, as the register
   // scavenger won't work with empty blocks (SIInstrInfo::insertIndirectBranch
   // uses the same workaround).
   Register ScratchReg = MRI.createVirtualRegister(&RISCV::GPRRegClass);
   auto II = MBB.end();
-  // We may also update the jump target to RestoreBB later.
+  // We may also update the jump target later.
   MachineInstr &MI = *BuildMI(MBB, II, DL, get(RISCV::PseudoJump))
                           .addReg(ScratchReg, RegState::Define | RegState::Dead)
                           .addMBB(&DestBB, RISCVII::MO_CALL);
 
   RS->enterBasicBlockEnd(MBB);
   Register TmpGPR =
       RS->scavengeRegisterBackwards(RISCV::GPRRegClass, MI.getIterator(),
                                     /*RestoreAfter=*/false, /*SpAdj=*/0,
else {
     if (FrameIndex == -1)
       report_fatal_error("underestimated function size");
 
     storeRegToStackSlot(MBB, MI, TmpGPR, /*IsKill=*/true, FrameIndex,
                         &RISCV::GPRRegClass, TRI, Register());
     TRI->eliminateFrameIndex(std::prev(MI.getIterator()),
                              /*SpAdj=*/0, /*FIOperandNum=*/1);
 
-    MI.getOperand(1).setMBB(&RestoreBB);
-
-    loadRegFromStackSlot(RestoreBB, RestoreBB.end(), TmpGPR, FrameIndex,
+    loadRegFromStackSlot(*RestoreBB, RestoreBB->end(), TmpGPR, FrameIndex,
                          &RISCV::GPRRegClass, TRI, Register());
-    TRI->eliminateFrameIndex(RestoreBB.back(),
+    TRI->eliminateFrameIndex(RestoreBB->back(),
                              /*SpAdj=*/0, /*FIOperandNum=*/1);
+
+    MachineBasicBlock *JumpDest = DeduplicateRestoreBB(RestoreBB);
+    MI.getOperand(1).setMBB(JumpDest);
   }
 
   MRI.replaceRegWith(ScratchReg, TmpGPR);
   MRI.clearVirtRegs();
 }
 
 bool RISCVInstrInfo::reverseBranchCondition(
     SmallVectorImpl<MachineOperand> &Cond) const {

llvm/test/CodeGen/AMDGPU/branch-relax-spill-deduplication.ll

 ; CHECK-NEXT:    s_cbranch_scc0 .LBB0_2
 ; CHECK-NEXT:  .LBB0_6: ; %bb4
 ; CHECK-NEXT:    s_not_b64 exec, exec
 ; CHECK-NEXT:    buffer_store_dword v1, off, s[96:99], 0
 ; CHECK-NEXT:    v_writelane_b32 v1, s0, 0
 ; CHECK-NEXT:    v_writelane_b32 v1, s1, 1
 ; CHECK-NEXT:    s_getpc_b64 s[0:1]
 ; CHECK-NEXT:  .Lpost_getpc1:
-; CHECK-NEXT:    s_add_u32 s0, s0, (.LBB0_7-.Lpost_getpc1)&4294967295
-; CHECK-NEXT:    s_addc_u32 s1, s1, (.LBB0_7-.Lpost_getpc1)>>32
+; CHECK-NEXT:    s_add_u32 s0, s0, (.LBB0_5-.Lpost_getpc1)&4294967295
+; CHECK-NEXT:    s_addc_u32 s1, s1, (.LBB0_5-.Lpost_getpc1)>>32
 ; CHECK-NEXT:    s_setpc_b64 s[0:1]
 ; CHECK-NEXT:  .LBB0_2: ; %bb2
 ; CHECK-NEXT:    ;;#ASMSTART
 ; CHECK-NEXT:    v_nop_e64
 ; CHECK-NEXT:    v_nop_e64
 ; CHECK-NEXT:    v_nop_e64
 ; CHECK-NEXT:    v_nop_e64
 ; CHECK-NEXT:    v_nop_e64

 ; CHECK-NEXT:    v_nop_e64
 ; CHECK-NEXT:    ;;#ASMEND
 ; CHECK-NEXT:    s_branch .LBB0_3
 ; CHECK-NEXT:  .LBB0_5: ; %bb3
 ; CHECK-NEXT:    v_readlane_b32 s0, v1, 0
 ; CHECK-NEXT:    v_readlane_b32 s1, v1, 1
 ; CHECK-NEXT:    buffer_load_dword v1, off, s[96:99], 0
 ; CHECK-NEXT:    s_not_b64 exec, exec
-; CHECK-NEXT:    s_branch .LBB0_3
-; CHECK-NEXT:  .LBB0_7: ; %bb3
-; CHECK-NEXT:    v_readlane_b32 s0, v1, 0
-; CHECK-NEXT:    v_readlane_b32 s1, v1, 1
-; CHECK-NEXT:    buffer_load_dword v1, off, s[96:99], 0
-; CHECK-NEXT:    s_not_b64 exec, exec
 ; CHECK-NEXT:  .LBB0_3: ; %bb3
 ; CHECK-NEXT:    v_readlane_b32 s0, v0, 0
 ; CHECK-NEXT:    ;;#ASMSTART
 ; CHECK-NEXT:    ; reg use s0
 ; CHECK-NEXT:    ;;#ASMEND
 ; CHECK-NEXT:    ;;#ASMSTART
 ; CHECK-NEXT:    ; reg use s1
 ; CHECK-NEXT:    ;;#ASMEND
bb3:
   tail call void asm sideeffect "; reg use $0", "{vcc_hi}"(i32 %vcc_hi) #0
   ret void
 }
 
 define void @spill_func_duplicated_restore_block(i32 addrspace(1)* %arg) #0 {
 ; CHECK-LABEL: spill_func_duplicated_restore_block:
 ; CHECK:       ; %bb.0: ; %entry
 ; CHECK-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
-; CHECK-NEXT:    s_or_saveexec_b64 s[4:5], -1
+; CHECK-NEXT:    s_xor_saveexec_b64 s[4:5], -1
 ; CHECK-NEXT:    buffer_store_dword v0, off, s[0:3], s32 ; 4-byte Folded Spill
 ; CHECK-NEXT:    buffer_store_dword v1, off, s[0:3], s32 offset:4 ; 4-byte Folded Spill
 ; CHECK-NEXT:    s_mov_b64 exec, s[4:5]
 ; CHECK-NEXT:    s_waitcnt expcnt(1)
 ; CHECK-NEXT:    v_writelane_b32 v0, s30, 0
 ; CHECK-NEXT:    v_writelane_b32 v0, s31, 1
 ; CHECK-NEXT:    v_writelane_b32 v0, s33, 2
 ; CHECK-NEXT:    v_writelane_b32 v0, s34, 3

 ; CHECK-NEXT:    v_writelane_b32 v0, s90, 59
 ; CHECK-NEXT:    v_writelane_b32 v1, s98, 3
 ; CHECK-NEXT:    v_writelane_b32 v0, s91, 60
 ; CHECK-NEXT:    v_writelane_b32 v1, s99, 4
 ; CHECK-NEXT:    v_writelane_b32 v0, s92, 61
 ; CHECK-NEXT:    v_writelane_b32 v1, s100, 5
 ; CHECK-NEXT:    v_writelane_b32 v0, s93, 62
 ; CHECK-NEXT:    v_writelane_b32 v1, s101, 6
-; CHECK-NEXT:    s_mov_b32 s34, s12
+; CHECK-NEXT:    s_mov_b32 s49, s12
 ; CHECK-NEXT:    v_writelane_b32 v0, s94, 63
-; CHECK-NEXT:    v_writelane_b32 v1, s34, 7
-; CHECK-NEXT:    s_cmp_eq_u32 s34, 42
+; CHECK-NEXT:    v_writelane_b32 v1, s49, 7
+; CHECK-NEXT:    s_cmp_eq_u32 s49, 42
 ; CHECK-NEXT:    ;;#ASMSTART
 ; CHECK-NEXT:    s_mov_b32 s0, 0
 ; CHECK-NEXT:    ;;#ASMEND
 ; CHECK-NEXT:    ;;#ASMSTART
 ; CHECK-NEXT:    s_mov_b32 s1, 0
 ; CHECK-NEXT:    ;;#ASMEND
 ; CHECK-NEXT:    ;;#ASMSTART
 ; CHECK-NEXT:    s_mov_b32 s2, 0

 ; CHECK-NEXT:    v_readlane_b32 s38, v0, 7
 ; CHECK-NEXT:    v_readlane_b32 s37, v0, 6
 ; CHECK-NEXT:    v_readlane_b32 s36, v0, 5
 ; CHECK-NEXT:    v_readlane_b32 s35, v0, 4
 ; CHECK-NEXT:    v_readlane_b32 s34, v0, 3
 ; CHECK-NEXT:    v_readlane_b32 s33, v0, 2
 ; CHECK-NEXT:    v_readlane_b32 s31, v0, 1
 ; CHECK-NEXT:    v_readlane_b32 s30, v0, 0
-; CHECK-NEXT:    s_or_saveexec_b64 s[4:5], -1
+; CHECK-NEXT:    s_xor_saveexec_b64 s[4:5], -1
 ; CHECK-NEXT:    buffer_load_dword v0, off, s[0:3], s32 ; 4-byte Folded Reload
 ; CHECK-NEXT:    buffer_load_dword v1, off, s[0:3], s32 offset:4 ; 4-byte Folded Reload
 ; CHECK-NEXT:    s_mov_b64 exec, s[4:5]
 ; CHECK-NEXT:    s_waitcnt vmcnt(0)
 ; CHECK-NEXT:    s_setpc_b64 s[30:31]
 entry:
   %cnd = tail call i32 @llvm.amdgcn.workgroup.id.x() #0
   %sgpr0 = tail call i32 asm sideeffect "s_mov_b32 s0, 0", "={s0}"() #0

llvm/test/CodeGen/RISCV/branch-relaxation.ll

branch_2:
   call void asm sideeffect "# reg use $0", "{t3}"(i64 %t3)
   call void asm sideeffect "# reg use $0", "{t4}"(i64 %t4)
   call void asm sideeffect "# reg use $0", "{t5}"(i64 %t5)
   call void asm sideeffect "# reg use $0", "{t6}"(i64 %t6)
 
   ret void
 }
 
 define void @relax_jal_spill_32_restore_block_correspondence() {
 ; CHECK-RV32-LABEL: relax_jal_spill_32_restore_block_correspondence:
 ; CHECK-RV32:       # %bb.0: # %entry
 ; CHECK-RV32-NEXT:    addi sp, sp, -64
 ; CHECK-RV32-NEXT:    .cfi_def_cfa_offset 64
 ; CHECK-RV32-NEXT:    sw ra, 60(sp) # 4-byte Folded Spill
 ; CHECK-RV32-NEXT:    sw s0, 56(sp) # 4-byte Folded Spill
 ; CHECK-RV32-NEXT:    sw s1, 52(sp) # 4-byte Folded Spill
 ; CHECK-RV32-NEXT:    sw s2, 48(sp) # 4-byte Folded Spill

 ; CHECK-RV32-NEXT:  # %bb.11: # %cond_3
 ; CHECK-RV32-NEXT:    sw s11, 0(sp)
 ; CHECK-RV32-NEXT:    jump .LBB6_12, s11
 ; CHECK-RV32-NEXT:  .LBB6_3: # %space
 ; CHECK-RV32-NEXT:    #APP
 ; CHECK-RV32-NEXT:    .zero 1048576
 ; CHECK-RV32-NEXT:    #NO_APP
 ; CHECK-RV32-NEXT:    j .LBB6_4
 ; CHECK-RV32-NEXT:  .LBB6_8: # %dest_1
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)
 ; CHECK-RV32-NEXT:  .LBB6_4: # %dest_1
 ; CHECK-RV32-NEXT:    #APP
 ; CHECK-RV32-NEXT:    # dest 1
 ; CHECK-RV32-NEXT:    #NO_APP
 ; CHECK-RV32-NEXT:    j .LBB6_5
 ; CHECK-RV32-NEXT:  .LBB6_10: # %dest_2
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)

 ; CHECK-RV32-NEXT:    bne t3, t4, .LBB7_2
 ; CHECK-RV32-NEXT:  # %bb.12: # %cond_2
 ; CHECK-RV32-NEXT:    sw s11, 0(sp)
 ; CHECK-RV32-NEXT:    jump .LBB7_13, s11
 ; CHECK-RV32-NEXT:  .LBB7_2: # %cond_3
 ; CHECK-RV32-NEXT:    bne t1, t2, .LBB7_3
 ; CHECK-RV32-NEXT:  # %bb.14: # %cond_3
 ; CHECK-RV32-NEXT:    sw s11, 0(sp)
-; CHECK-RV32-NEXT:    jump .LBB7_15, s11
+; CHECK-RV32-NEXT:    jump .LBB7_11, s11
 ; CHECK-RV32-NEXT:  .LBB7_3: # %cond_4
 ; CHECK-RV32-NEXT:    bne s9, s10, .LBB7_4
 ; CHECK-RV32-NEXT:  # %bb.16: # %cond_4
 ; CHECK-RV32-NEXT:    sw s11, 0(sp)
-; CHECK-RV32-NEXT:    jump .LBB7_17, s11
+; CHECK-RV32-NEXT:    jump .LBB7_13, s11
 ; CHECK-RV32-NEXT:  .LBB7_4: # %cond_5
 ; CHECK-RV32-NEXT:    bne s7, s8, .LBB7_5
 ; CHECK-RV32-NEXT:  # %bb.18: # %cond_5
 ; CHECK-RV32-NEXT:    sw s11, 0(sp)
 ; CHECK-RV32-NEXT:    jump .LBB7_19, s11
 ; CHECK-RV32-NEXT:  .LBB7_5: # %cond_6
 ; CHECK-RV32-NEXT:    bne s5, s6, .LBB7_6
 ; CHECK-RV32-NEXT:  # %bb.20: # %cond_6
 ; CHECK-RV32-NEXT:    sw s11, 0(sp)
-; CHECK-RV32-NEXT:    jump .LBB7_21, s11
+; CHECK-RV32-NEXT:    jump .LBB7_11, s11
 ; CHECK-RV32-NEXT:  .LBB7_23: # %space
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)
 ; CHECK-RV32-NEXT:  .LBB7_6: # %space
 ; CHECK-RV32-NEXT:    #APP
 ; CHECK-RV32-NEXT:    .zero 1048576
 ; CHECK-RV32-NEXT:    #NO_APP
 ; CHECK-RV32-NEXT:    j .LBB7_7
 ; CHECK-RV32-NEXT:  .LBB7_11: # %dest_1
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)
-; CHECK-RV32-NEXT:    j .LBB7_7
-; CHECK-RV32-NEXT:  .LBB7_15: # %dest_1
-; CHECK-RV32-NEXT:    lw s11, 0(sp)
-; CHECK-RV32-NEXT:    j .LBB7_7
-; CHECK-RV32-NEXT:  .LBB7_21: # %dest_1
-; CHECK-RV32-NEXT:    lw s11, 0(sp)
 ; CHECK-RV32-NEXT:  .LBB7_7: # %dest_1
 ; CHECK-RV32-NEXT:    #APP
 ; CHECK-RV32-NEXT:    # dest 1
 ; CHECK-RV32-NEXT:    #NO_APP
 ; CHECK-RV32-NEXT:    j .LBB7_8
 ; CHECK-RV32-NEXT:  .LBB7_13: # %dest_2
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)
-; CHECK-RV32-NEXT:    j .LBB7_8
-; CHECK-RV32-NEXT:  .LBB7_17: # %dest_2
-; CHECK-RV32-NEXT:    lw s11, 0(sp)
 ; CHECK-RV32-NEXT:  .LBB7_8: # %dest_2
 ; CHECK-RV32-NEXT:    #APP
 ; CHECK-RV32-NEXT:    # dest 2
 ; CHECK-RV32-NEXT:    #NO_APP
 ; CHECK-RV32-NEXT:    j .LBB7_9
 ; CHECK-RV32-NEXT:  .LBB7_19: # %dest_3
 ; CHECK-RV32-NEXT:    lw s11, 0(sp)
 ; CHECK-RV32-NEXT:  .LBB7_9: # %dest_3