[Mips] Make MipsBranchExpansion aware of BBIT family of branch

Octeon branches (bbit0/bbit032/bbit1/bbit132) have an immediate operand,
so it is legal to have such replacement within
MipsBranchExpansion::replaceBranch().

According to the specification, a branch (e.g. bbit0 ) looks like:

bbit0  rs p offset  // p is an immediate operand
  if !rs<p> then branch

Without this patch, an assertion triggers in the method,
and the problem has been found in the real example.

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

llvm/lib/Target/Mips/MipsBranchExpansion.cpp

@@ -342,16 +342,25 @@ void MipsBranchExpansion::replaceBranch(MachineBasicBlock &MBB, Iter Br,
   for (unsigned I = 0, E = Br->getDesc().getNumOperands(); I < E; ++I) {
     MachineOperand &MO = Br->getOperand(I);
 
-    if (!MO.isReg()) {
-      assert(MO.isMBB() && "MBB operand expected.");
+    switch (MO.getType()) {
+    case MachineOperand::MO_Register:
+      MIB.addReg(MO.getReg());
       break;
+    case MachineOperand::MO_Immediate:
+      // Octeon BBIT family of branch has an immediate operand
+      // (e.g. BBIT0 $v0, 3, %bb.1).
+      if (!TII->isBranchWithImm(Br->getOpcode()))
+        llvm_unreachable("Unexpected immediate in branch instruction");
+      MIB.addImm(MO.getImm());
+      break;
+    case MachineOperand::MO_MachineBasicBlock:
+      MIB.addMBB(MBBOpnd);
+      break;
+    default:
+      llvm_unreachable("Unexpected operand type in branch instruction");
     }
-
-    MIB.addReg(MO.getReg());
   }
 
-  MIB.addMBB(MBBOpnd);
-
   if (Br->hasDelaySlot()) {
     // Bundle the instruction in the delay slot to the newly created branch
     // and erase the original branch.

llvm/lib/Target/Mips/MipsInstrInfo.h

@@ -106,6 +106,10 @@ class MipsInstrInfo : public MipsGenInstrInfo {
 
   virtual unsigned getOppositeBranchOpc(unsigned Opc) const = 0;
 
+  virtual bool isBranchWithImm(unsigned Opc) const {
+    return false;
+  }
+
   /// Return the number of bytes of code the specified instruction may be.
   unsigned getInstSizeInBytes(const MachineInstr &MI) const override;
 

llvm/lib/Target/Mips/MipsSEInstrInfo.cpp

@@ -483,6 +483,20 @@ bool MipsSEInstrInfo::expandPostRAPseudo(MachineInstr &MI) const {
   return true;
 }
 
+/// isBranchWithImm - Return true if the branch contains an immediate
+/// operand (\see lib/Target/Mips/MipsBranchExpansion.cpp).
+bool MipsSEInstrInfo::isBranchWithImm(unsigned Opc) const {
+  switch (Opc) {
+  default:
+    return false;
+  case Mips::BBIT0:
+  case Mips::BBIT1:
+  case Mips::BBIT032:
+  case Mips::BBIT132:
+    return true;
+  }
+}
+
 /// getOppositeBranchOpc - Return the inverse of the specified
 /// opcode, e.g. turning BEQ to BNE.
 unsigned MipsSEInstrInfo::getOppositeBranchOpc(unsigned Opc) const {

llvm/lib/Target/Mips/MipsSEInstrInfo.h

@@ -62,6 +62,8 @@ class MipsSEInstrInfo : public MipsInstrInfo {
 
   bool expandPostRAPseudo(MachineInstr &MI) const override;
 
+  bool isBranchWithImm(unsigned Opc) const override;
+
   unsigned getOppositeBranchOpc(unsigned Opc) const override;
 
   /// Adjust SP by Amount bytes.

llvm/test/CodeGen/Mips/longbranch/long-branch-octeon.ll

@@ -0,0 +1,105 @@
+;; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+;; Test that Octeon BBIT family of branch can be replaced by
+;; the long branch expansion pass.
+
+; RUN: llc -O3 -mtriple=mips64-octeon-linux -mcpu=octeon -force-mips-long-branch < %s -o - | FileCheck %s
+
+define i64 @bbit1(i64 %a) nounwind {
+; CHECK-LABEL: bbit1:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    bbit0 $4, 3, .LBB0_2
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  # %bb.1: # %entry
+; CHECK-NEXT:    j .LBB0_3
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  .LBB0_2: # %endif
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 12
+; CHECK-NEXT:  .LBB0_3: # %if
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 48
+entry:
+  %bit = and i64 %a, 8
+  %res = icmp eq i64 %bit, 0
+  br i1 %res, label %endif, label %if
+if:
+  ret i64 48
+
+endif:
+  ret i64 12
+}
+
+define i64 @bbit132(i64 %a) nounwind {
+; CHECK-LABEL: bbit132:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    bbit032 $4, 3, .LBB1_2
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  # %bb.1: # %entry
+; CHECK-NEXT:    j .LBB1_3
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  .LBB1_2: # %endif
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 12
+; CHECK-NEXT:  .LBB1_3: # %if
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 48
+entry:
+  %bit = and i64 %a, 34359738368
+  %res = icmp eq i64 %bit, 0
+  br i1 %res, label %endif, label %if
+if:
+  ret i64 48
+
+endif:
+  ret i64 12
+}
+
+define i64 @bbit0(i64 %a) nounwind {
+; CHECK-LABEL: bbit0:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    bbit1 $4, 3, .LBB2_2
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  # %bb.1: # %entry
+; CHECK-NEXT:    j .LBB2_3
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  .LBB2_2: # %endif
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 12
+; CHECK-NEXT:  .LBB2_3: # %if
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 48
+entry:
+  %bit = and i64 %a, 8
+  %res = icmp ne i64 %bit, 0
+  br i1 %res, label %endif, label %if
+if:
+  ret i64 48
+
+endif:
+  ret i64 12
+}
+
+define i64 @bbit032(i64 %a) nounwind {
+; CHECK-LABEL: bbit032:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    bbit132 $4, 3, .LBB3_2
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  # %bb.1: # %entry
+; CHECK-NEXT:    j .LBB3_3
+; CHECK-NEXT:    nop
+; CHECK-NEXT:  .LBB3_2: # %endif
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 12
+; CHECK-NEXT:  .LBB3_3: # %if
+; CHECK-NEXT:    jr $ra
+; CHECK-NEXT:    daddiu $2, $zero, 48
+entry:
+  %bit = and i64 %a, 34359738368
+  %res = icmp ne i64 %bit, 0
+  br i1 %res, label %endif, label %if
+if:
+  ret i64 48
+
+endif:
+  ret i64 12
+}