[LoongArch] Implement branch analysis

This allows a number of optimisation passes to work.
E.g. BranchFolding and MachineBlockPlacement.

Differential Revision: https://reviews.llvm.org/D131316

Author: wangleiat

llvm/lib/Target/LoongArch/LoongArchInstrInfo.cpp

@@ -113,3 +113,192 @@ void LoongArchInstrInfo::loadRegFromStackSlot(
       .addImm(0)
       .addMemOperand(MMO);
 }
+
+unsigned LoongArchInstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
+  return MI.getDesc().getSize();
+}
+
+MachineBasicBlock *
+LoongArchInstrInfo::getBranchDestBlock(const MachineInstr &MI) const {
+  assert(MI.getDesc().isBranch() && "Unexpected opcode!");
+  // The branch target is always the last operand.
+  return MI.getOperand(MI.getNumExplicitOperands() - 1).getMBB();
+}
+
+static void parseCondBranch(MachineInstr &LastInst, MachineBasicBlock *&Target,
+                            SmallVectorImpl<MachineOperand> &Cond) {
+  // Block ends with fall-through condbranch.
+  assert(LastInst.getDesc().isConditionalBranch() &&
+         "Unknown conditional branch");
+  int NumOp = LastInst.getNumExplicitOperands();
+  Target = LastInst.getOperand(NumOp - 1).getMBB();
+
+  Cond.push_back(MachineOperand::CreateImm(LastInst.getOpcode()));
+  for (int i = 0; i < NumOp - 1; i++)
+    Cond.push_back(LastInst.getOperand(i));
+}
+
+bool LoongArchInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
+                                       MachineBasicBlock *&TBB,
+                                       MachineBasicBlock *&FBB,
+                                       SmallVectorImpl<MachineOperand> &Cond,
+                                       bool AllowModify) const {
+  TBB = FBB = nullptr;
+  Cond.clear();
+
+  // If the block has no terminators, it just falls into the block after it.
+  MachineBasicBlock::iterator I = MBB.getLastNonDebugInstr();
+  if (I == MBB.end() || !isUnpredicatedTerminator(*I))
+    return false;
+
+  // Count the number of terminators and find the first unconditional or
+  // indirect branch.
+  MachineBasicBlock::iterator FirstUncondOrIndirectBr = MBB.end();
+  int NumTerminators = 0;
+  for (auto J = I.getReverse(); J != MBB.rend() && isUnpredicatedTerminator(*J);
+       J++) {
+    NumTerminators++;
+    if (J->getDesc().isUnconditionalBranch() ||
+        J->getDesc().isIndirectBranch()) {
+      FirstUncondOrIndirectBr = J.getReverse();
+    }
+  }
+
+  // If AllowModify is true, we can erase any terminators after
+  // FirstUncondOrIndirectBR.
+  if (AllowModify && FirstUncondOrIndirectBr != MBB.end()) {
+    while (std::next(FirstUncondOrIndirectBr) != MBB.end()) {
+      std::next(FirstUncondOrIndirectBr)->eraseFromParent();
+      NumTerminators--;
+    }
+    I = FirstUncondOrIndirectBr;
+  }
+
+  // Handle a single unconditional branch.
+  if (NumTerminators == 1 && I->getDesc().isUnconditionalBranch()) {
+    TBB = getBranchDestBlock(*I);
+    return false;
+  }
+
+  // Handle a single conditional branch.
+  if (NumTerminators == 1 && I->getDesc().isConditionalBranch()) {
+    parseCondBranch(*I, TBB, Cond);
+    return false;
+  }
+
+  // Handle a conditional branch followed by an unconditional branch.
+  if (NumTerminators == 2 && std::prev(I)->getDesc().isConditionalBranch() &&
+      I->getDesc().isUnconditionalBranch()) {
+    parseCondBranch(*std::prev(I), TBB, Cond);
+    FBB = getBranchDestBlock(*I);
+    return false;
+  }
+
+  // Otherwise, we can't handle this.
+  return true;
+}
+
+unsigned LoongArchInstrInfo::removeBranch(MachineBasicBlock &MBB,
+                                          int *BytesRemoved) const {
+  if (BytesRemoved)
+    *BytesRemoved = 0;
+  MachineBasicBlock::iterator I = MBB.getLastNonDebugInstr();
+  if (I == MBB.end())
+    return 0;
+
+  if (!I->getDesc().isBranch())
+    return 0;
+
+  // Remove the branch.
+  if (BytesRemoved)
+    *BytesRemoved += getInstSizeInBytes(*I);
+  I->eraseFromParent();
+
+  I = MBB.end();
+
+  if (I == MBB.begin())
+    return 1;
+  --I;
+  if (!I->getDesc().isConditionalBranch())
+    return 1;
+
+  // Remove the branch.
+  if (BytesRemoved)
+    *BytesRemoved += getInstSizeInBytes(*I);
+  I->eraseFromParent();
+  return 2;
+}
+
+// Inserts a branch into the end of the specific MachineBasicBlock, returning
+// the number of instructions inserted.
+unsigned LoongArchInstrInfo::insertBranch(
+    MachineBasicBlock &MBB, MachineBasicBlock *TBB, MachineBasicBlock *FBB,
+    ArrayRef<MachineOperand> Cond, const DebugLoc &DL, int *BytesAdded) const {
+  if (BytesAdded)
+    *BytesAdded = 0;
+
+  // Shouldn't be a fall through.
+  assert(TBB && "insertBranch must not be told to insert a fallthrough");
+  assert(Cond.size() <= 3 && Cond.size() != 1 &&
+         "LoongArch branch conditions have at most two components!");
+
+  // Unconditional branch.
+  if (Cond.empty()) {
+    MachineInstr &MI = *BuildMI(&MBB, DL, get(LoongArch::PseudoBR)).addMBB(TBB);
+    if (BytesAdded)
+      *BytesAdded += getInstSizeInBytes(MI);
+    return 1;
+  }
+
+  // Either a one or two-way conditional branch.
+  MachineInstrBuilder MIB = BuildMI(&MBB, DL, get(Cond[0].getImm()));
+  for (unsigned i = 1; i < Cond.size(); ++i)
+    MIB.add(Cond[i]);
+  MIB.addMBB(TBB);
+  if (BytesAdded)
+    *BytesAdded += getInstSizeInBytes(*MIB);
+
+  // One-way conditional branch.
+  if (!FBB)
+    return 1;
+
+  // Two-way conditional branch.
+  MachineInstr &MI = *BuildMI(&MBB, DL, get(LoongArch::PseudoBR)).addMBB(FBB);
+  if (BytesAdded)
+    *BytesAdded += getInstSizeInBytes(MI);
+  return 2;
+}
+
+static unsigned getOppositeBranchOpc(unsigned Opc) {
+  switch (Opc) {
+  default:
+    llvm_unreachable("Unrecognized conditional branch");
+  case LoongArch::BEQ:
+    return LoongArch::BNE;
+  case LoongArch::BNE:
+    return LoongArch::BEQ;
+  case LoongArch::BEQZ:
+    return LoongArch::BNEZ;
+  case LoongArch::BNEZ:
+    return LoongArch::BEQZ;
+  case LoongArch::BCEQZ:
+    return LoongArch::BCNEZ;
+  case LoongArch::BCNEZ:
+    return LoongArch::BCEQZ;
+  case LoongArch::BLT:
+    return LoongArch::BGE;
+  case LoongArch::BGE:
+    return LoongArch::BLT;
+  case LoongArch::BLTU:
+    return LoongArch::BGEU;
+  case LoongArch::BGEU:
+    return LoongArch::BLTU;
+  }
+}
+
+bool LoongArchInstrInfo::reverseBranchCondition(
+    SmallVectorImpl<MachineOperand> &Cond) const {
+  assert((Cond.size() && Cond.size() <= 3) && "Invalid branch condition!");
+  Cond[0].setImm(getOppositeBranchOpc(Cond[0].getImm()));
+  return false;
+}

llvm/lib/Target/LoongArch/LoongArchInstrInfo.h

@@ -40,6 +40,26 @@ class LoongArchInstrInfo : public LoongArchGenInstrInfo {
                             MachineBasicBlock::iterator MBBI, Register DstReg,
                             int FrameIndex, const TargetRegisterClass *RC,
                             const TargetRegisterInfo *TRI) const override;
+
+  unsigned getInstSizeInBytes(const MachineInstr &MI) const override;
+
+  MachineBasicBlock *getBranchDestBlock(const MachineInstr &MI) const override;
+
+  bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
+                     MachineBasicBlock *&FBB,
+                     SmallVectorImpl<MachineOperand> &Cond,
+                     bool AllowModify) const override;
+
+  unsigned removeBranch(MachineBasicBlock &MBB,
+                        int *BytesRemoved = nullptr) const override;
+
+  unsigned insertBranch(MachineBasicBlock &MBB, MachineBasicBlock *TBB,
+                        MachineBasicBlock *FBB, ArrayRef<MachineOperand> Cond,
+                        const DebugLoc &dl,
+                        int *BytesAdded = nullptr) const override;
+
+  bool
+  reverseBranchCondition(SmallVectorImpl<MachineOperand> &Cond) const override;
 };
 
 } // end namespace llvm

llvm/test/CodeGen/LoongArch/analyze-branch.ll

@@ -0,0 +1,76 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc --mtriple=loongarch64 < %s | FileCheck %s
+
+;; This test checks that LLVM can do basic stripping and reapplying of branches
+;; to basic blocks.
+
+declare void @test_true()
+declare void @test_false()
+
+;; !0 corresponds to a branch being taken, !1 to not being taken.
+!0 = !{!"branch_weights", i32 64, i32 4}
+!1 = !{!"branch_weights", i32 4, i32 64}
+
+define void @test_bcc_fallthrough_taken(i64 %in) nounwind {
+; CHECK-LABEL: test_bcc_fallthrough_taken:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi.d $sp, $sp, -16
+; CHECK-NEXT:    st.d $ra, $sp, 8 # 8-byte Folded Spill
+; CHECK-NEXT:    ori $a1, $zero, 42
+; CHECK-NEXT:    bne $a0, $a1, .LBB0_3
+; CHECK-NEXT:  # %bb.1: # %true
+; CHECK-NEXT:    bl test_true
+; CHECK-NEXT:  .LBB0_2: # %true
+; CHECK-NEXT:    ld.d $ra, $sp, 8 # 8-byte Folded Reload
+; CHECK-NEXT:    addi.d $sp, $sp, 16
+; CHECK-NEXT:    jirl $zero, $ra, 0
+; CHECK-NEXT:  .LBB0_3: # %false
+; CHECK-NEXT:    bl test_false
+; CHECK-NEXT:    b .LBB0_2
+  %tst = icmp eq i64 %in, 42
+  br i1 %tst, label %true, label %false, !prof !0
+
+;; Expected layout order is: Entry, TrueBlock, FalseBlock
+;; Entry->TrueBlock is the common path, which should be taken whenever the
+;; conditional branch is false.
+
+true:
+  call void @test_true()
+  ret void
+
+false:
+  call void @test_false()
+  ret void
+}
+
+define void @test_bcc_fallthrough_nottaken(i64 %in) nounwind {
+; CHECK-LABEL: test_bcc_fallthrough_nottaken:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    addi.d $sp, $sp, -16
+; CHECK-NEXT:    st.d $ra, $sp, 8 # 8-byte Folded Spill
+; CHECK-NEXT:    ori $a1, $zero, 42
+; CHECK-NEXT:    beq $a0, $a1, .LBB1_1
+; CHECK-NEXT:  # %bb.3: # %false
+; CHECK-NEXT:    bl test_false
+; CHECK-NEXT:  .LBB1_2: # %true
+; CHECK-NEXT:    ld.d $ra, $sp, 8 # 8-byte Folded Reload
+; CHECK-NEXT:    addi.d $sp, $sp, 16
+; CHECK-NEXT:    jirl $zero, $ra, 0
+; CHECK-NEXT:  .LBB1_1: # %true
+; CHECK-NEXT:    bl test_true
+; CHECK-NEXT:    b .LBB1_2
+  %tst = icmp eq i64 %in, 42
+  br i1 %tst, label %true, label %false, !prof !1
+
+;; Expected layout order is: Entry, FalseBlock, TrueBlock
+;; Entry->FalseBlock is the common path, which should be taken whenever the
+;; conditional branch is false.
+
+true:
+  call void @test_true()
+  ret void
+
+false:
+  call void @test_false()
+  ret void
+}