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
+}