bpf: Extends zero extension elimination beyond comparison instructions

The current zero extension elimination was restricted to operands of
comparison. It actually could be extended to more cases.

For example:

  int *inc_p (int *p, unsigned a)
  {
    return p + a;
  }

'a' will be promoted to i64 during addition, and the zero extension could
be eliminated as well.

For the elimination optimization, it should be much better to start
recognizing the candidate sequence from the SRL instruction instead of J*
instructions.

This patch makes it an generic zero extension elimination pass instead of
one restricted with comparison.

Signed-off-by: Jiong Wang <jiong.wang@netronome.com>
Signed-off-by: Yonghong Song <yhs@fb.com>
llvm-svn: 327367

2 files changed, 77 insertions, 93 deletions

diff --git a/llvm/lib/Target/BPF/BPFMIPeephole.cpp b/llvm/lib/Target/BPF/BPFMIPeephole.cpp
index 7d84b2834f8..b15a7e6663e 100644
--- a/llvm/lib/Target/BPF/BPFMIPeephole.cpp
+++ b/llvm/lib/Target/BPF/BPFMIPeephole.cpp
@@ -27,9 +27,9 @@
 
 using namespace llvm;
 
-#define DEBUG_TYPE "bpf-mi-promotion-elim"
+#define DEBUG_TYPE "bpf-mi-zext-elim"
 
-STATISTIC(CmpPromotionElemNum, "Number of shifts for CMP promotion eliminated");
+STATISTIC(ZExtElemNum, "Number of zero extension shifts eliminated");
 
 namespace {
 
@@ -48,10 +48,8 @@ private:
   // Initialize class variables.
   void initialize(MachineFunction &MFParm);
 
-  bool eliminateCmpPromotionSeq(void);
-  MachineInstr *getInsnDefZExtSubReg(unsigned Reg) const;
-  void updateInsnSeq(MachineBasicBlock &MBB, MachineInstr &MI,
-                     unsigned Reg) const;
+  bool isMovFrom32Def(MachineInstr *MovMI);
+  bool eliminateZExtSeq(void);
 
 public:
 
@@ -62,7 +60,7 @@ public:
 
     initialize(MF);
 
-    return eliminateCmpPromotionSeq();
+    return eliminateZExtSeq();
   }
 };
 
@@ -74,107 +72,77 @@ void BPFMIPeephole::initialize(MachineFunction &MFParm) {
   DEBUG(dbgs() << "*** BPF MI peephole pass ***\n\n");
 }
 
-MachineInstr *BPFMIPeephole::getInsnDefZExtSubReg(unsigned Reg) const {
-  MachineInstr *Insn = MRI->getVRegDef(Reg);
+bool BPFMIPeephole::isMovFrom32Def(MachineInstr *MovMI)
+{
+  MachineInstr *DefInsn = MRI->getVRegDef(MovMI->getOperand(1).getReg());
 
-  if (!Insn ||
-      Insn->isPHI() ||
-      Insn->getOpcode() != BPF::SRL_ri ||
-      Insn->getOperand(2).getImm() != 32)
-    return nullptr;
+  if (!DefInsn || DefInsn->isPHI())
+    return false;
 
-  Insn = MRI->getVRegDef(Insn->getOperand(1).getReg());
-  if (!Insn ||
-      Insn->isPHI() ||
-      Insn->getOpcode() != BPF::SLL_ri ||
-      Insn->getOperand(2).getImm() != 32)
-    return nullptr;
-
-  Insn = MRI->getVRegDef(Insn->getOperand(1).getReg());
-  if (!Insn ||
-      Insn->isPHI() ||
-      Insn->getOpcode() != BPF::MOV_32_64)
-    return nullptr;
-
-  Insn = MRI->getVRegDef(Insn->getOperand(1).getReg());
-  if (!Insn || Insn->isPHI())
-    return nullptr;
-
-  if (Insn->getOpcode() == BPF::COPY) {
-    MachineOperand &opnd = Insn->getOperand(1);
+  if (DefInsn->getOpcode() == BPF::COPY) {
+    MachineOperand &opnd = DefInsn->getOperand(1);
 
     if (!opnd.isReg())
-      return nullptr;
+      return false;
 
     unsigned Reg = opnd.getReg();
     if ((TargetRegisterInfo::isVirtualRegister(Reg) &&
-         MRI->getRegClass(Reg) == &BPF::GPRRegClass) ||
-        (TargetRegisterInfo::isPhysicalRegister(Reg) &&
-         BPF::GPRRegClass.contains(Reg)))
-      return nullptr;
+         MRI->getRegClass(Reg) == &BPF::GPRRegClass))
+       return false;
   }
 
-  return Insn;
+  return true;
 }
 
-void
-BPFMIPeephole::updateInsnSeq(MachineBasicBlock &MBB, MachineInstr &MI,
-                             unsigned Reg) const {
-  MachineInstr *Mov, *Lshift, *Rshift;
-  unsigned SubReg;
-  DebugLoc DL;
-
-  Rshift = MRI->getVRegDef(Reg);
-  Lshift = MRI->getVRegDef(Rshift->getOperand(1).getReg());
-  Mov = MRI->getVRegDef(Lshift->getOperand(1).getReg());
-  SubReg = Mov->getOperand(1).getReg();
-  DL = MI.getDebugLoc();
-  BuildMI(MBB, Rshift, DL, TII->get(BPF::SUBREG_TO_REG), Reg)
-    .addImm(0).addReg(SubReg).addImm(BPF::sub_32);
-  Rshift->eraseFromParent();
-  Lshift->eraseFromParent();
-  Mov->eraseFromParent();
-
-  CmpPromotionElemNum++;
-}
-
-bool BPFMIPeephole::eliminateCmpPromotionSeq(void) {
+bool BPFMIPeephole::eliminateZExtSeq(void) {
+  MachineInstr* ToErase = nullptr;
   bool Eliminated = false;
-  MachineInstr *Mov;
-  unsigned Reg;
 
   for (MachineBasicBlock &MBB : *MF) {
     for (MachineInstr &MI : MBB) {
-      switch (MI.getOpcode()) {
-      default:
-        break;
-      case BPF::JUGT_rr:
-      case BPF::JUGE_rr:
-      case BPF::JULT_rr:
-      case BPF::JULE_rr:
-      case BPF::JEQ_rr:
-      case BPF::JNE_rr:
-        Reg = MI.getOperand(1).getReg();
-        Mov = getInsnDefZExtSubReg(Reg);
-        if (Mov) {
-          updateInsnSeq(MBB, MI, Reg);
-          Eliminated = true;
-        }
-        // Fallthrough
-      case BPF::JUGT_ri:
-      case BPF::JUGE_ri:
-      case BPF::JULT_ri:
-      case BPF::JULE_ri:
-      case BPF::JEQ_ri:
-      case BPF::JNE_ri:
-        Reg = MI.getOperand(0).getReg();
-        Mov = getInsnDefZExtSubReg(Reg);
-        if (!Mov)
-          break;
-
-	updateInsnSeq(MBB, MI, Reg);
-	Eliminated = true;
-        break;
+      // If the previous instruction was marked for elimination, remove it now.
+      if (ToErase) {
+        ToErase->eraseFromParent();
+        ToErase = nullptr;
+      }
+
+      // Eliminate the 32-bit to 64-bit zero extension sequence when possible.
+      //
+      //   MOV_32_64 rB, wA
+      //   SLL_ri    rB, rB, 32
+      //   SRL_ri    rB, rB, 32
+      if (MI.getOpcode() == BPF::SRL_ri &&
+          MI.getOperand(2).getImm() == 32) {
+        unsigned DstReg = MI.getOperand(0).getReg();
+        unsigned ShfReg = MI.getOperand(1).getReg();
+        MachineInstr *SllMI = MRI->getVRegDef(ShfReg);
+
+        if (!SllMI ||
+            SllMI->isPHI() ||
+            SllMI->getOpcode() != BPF::SLL_ri ||
+            SllMI->getOperand(2).getImm() != 32)
+          continue;
+
+        MachineInstr *MovMI = MRI->getVRegDef(SllMI->getOperand(1).getReg());
+        if (!MovMI ||
+            MovMI->isPHI() ||
+            MovMI->getOpcode() != BPF::MOV_32_64)
+          continue;
+
+        if (!isMovFrom32Def(MovMI))
+          continue;
+
+        unsigned SubReg = MovMI->getOperand(1).getReg();
+        BuildMI(MBB, MI, MI.getDebugLoc(), TII->get(BPF::SUBREG_TO_REG), DstReg)
+          .addImm(0).addReg(SubReg).addImm(BPF::sub_32);
+
+        SllMI->eraseFromParent();
+        MovMI->eraseFromParent();
+        // MI is the right shift, we can't erase it in it's own iteration.
+        // Mark it to ToErase, and erase in the next iteration.
+        ToErase = &MI;
+        ZExtElemNum++;
+        Eliminated = true;
       }
     }
   }
diff --git a/llvm/test/CodeGen/BPF/32-bit-subreg-peephole.ll b/llvm/test/CodeGen/BPF/32-bit-subreg-peephole.ll
index 78c344b1577..ed5268458f8 100644
--- a/llvm/test/CodeGen/BPF/32-bit-subreg-peephole.ll
+++ b/llvm/test/CodeGen/BPF/32-bit-subreg-peephole.ll
@@ -35,6 +35,11 @@
 ;   else
 ;     return c;
 ; }
+;
+; int *inc_p (int *p, unsigned a)
+; {
+;   return p + a;
+; }
 
 ; Function Attrs: norecurse nounwind readnone
 define dso_local i64 @select_u(i32 %a, i32 %b, i64 %c, i64 %d) local_unnamed_addr #0 {
@@ -89,3 +94,14 @@ entry:
 }
 
 declare dso_local i64 @bar(...) local_unnamed_addr #1
+
+; Function Attrs: norecurse nounwind readnone
+define dso_local i32* @inc_p(i32* readnone %p, i32 %a) local_unnamed_addr #0 {
+; CHECK-LABEL: inc_p:
+entry:
+  %idx.ext = zext i32 %a to i64
+; CHECK-NOT: r{{[0-9]+}} <<= 32
+; CHECK-NOT: r{{[0-9]+}} >>= 32
+  %add.ptr = getelementptr inbounds i32, i32* %p, i64 %idx.ext
+  ret i32* %add.ptr
+}