2 files changed, 83 insertions, 0 deletions

diff --git a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
index 76555f99c58..df48c428cb1 100644
--- a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
@@ -193,6 +193,8 @@ MachineInstrBuilder MachineIRBuilder::buildExtract(ArrayRef<LLT> ResTys,
   assert(ResTys.size() == Results.size() && Results.size() == Indices.size() &&
          "inconsistent number of regs");
   assert(!Results.empty() && "invalid trivial extract");
+  assert(std::is_sorted(Indices.begin(), Indices.end()) &&
+         "extract offsets must be in ascending order");
 
   auto MIB = BuildMI(getMF(), DL, getTII().get(TargetOpcode::G_EXTRACT));
   for (unsigned i = 0; i < ResTys.size(); ++i)
@@ -222,6 +224,8 @@ MachineIRBuilder::buildSequence(LLT ResTy, unsigned Res,
   assert(OpTys.size() == Ops.size() && Ops.size() == Indices.size() &&
          "incompatible args");
   assert(!Ops.empty() && "invalid trivial sequence");
+  assert(std::is_sorted(Indices.begin(), Indices.end()) &&
+         "sequence offsets must be in ascending order");
 
   MachineInstrBuilder MIB =
       buildInstr(TargetOpcode::G_SEQUENCE, LLT::scalar(ResTy.getSizeInBits()));
diff --git a/llvm/lib/CodeGen/GlobalISel/MachineLegalizePass.cpp b/llvm/lib/CodeGen/GlobalISel/MachineLegalizePass.cpp
index ba56eb09db1..b0de76d075f 100644
--- a/llvm/lib/CodeGen/GlobalISel/MachineLegalizePass.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/MachineLegalizePass.cpp
@@ -19,6 +19,7 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetInstrInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 
 #define DEBUG_TYPE "legalize-mir"
@@ -46,6 +47,70 @@ void MachineLegalizePass::getAnalysisUsage(AnalysisUsage &AU) const {
 void MachineLegalizePass::init(MachineFunction &MF) {
 }
 
+bool MachineLegalizePass::combineExtracts(MachineInstr &MI,
+                                          MachineRegisterInfo &MRI,
+                                          const TargetInstrInfo &TII) {
+  bool Changed = false;
+  if (MI.getOpcode() != TargetOpcode::G_EXTRACT)
+    return Changed;
+
+  unsigned NumDefs = (MI.getNumOperands() - 1) / 2;
+  unsigned SrcReg = MI.getOperand(NumDefs).getReg();
+  MachineInstr &SeqI = *MRI.def_instr_begin(SrcReg);
+  if (SeqI.getOpcode() != TargetOpcode::G_SEQUENCE)
+      return Changed;
+
+  unsigned NumSeqSrcs = (SeqI.getNumOperands() - 1) / 2;
+  bool AllDefsReplaced = true;
+
+  // Try to match each register extracted with a corresponding insertion formed
+  // by the G_SEQUENCE.
+  for (unsigned Idx = 0, SeqIdx = 0; Idx < NumDefs; ++Idx) {
+    MachineOperand &ExtractMO = MI.getOperand(Idx);
+    assert(ExtractMO.isReg() && ExtractMO.isDef() &&
+           "unexpected extract operand");
+
+    unsigned ExtractReg = ExtractMO.getReg();
+    unsigned ExtractPos = MI.getOperand(NumDefs + Idx + 1).getImm();
+
+    while (SeqIdx < NumSeqSrcs &&
+           SeqI.getOperand(2 * SeqIdx + 2).getImm() < ExtractPos)
+      ++SeqIdx;
+
+    if (SeqIdx == NumSeqSrcs ||
+        SeqI.getOperand(2 * SeqIdx + 2).getImm() != ExtractPos ||
+        SeqI.getType(SeqIdx + 1) != MI.getType(Idx)) {
+      AllDefsReplaced = false;
+      continue;
+    }
+
+    unsigned OrigReg = SeqI.getOperand(2 * SeqIdx + 1).getReg();
+    assert(!TargetRegisterInfo::isPhysicalRegister(OrigReg) &&
+           "unexpected physical register in G_SEQUENCE");
+
+    // Finally we can replace the uses.
+    for (auto &Use : MRI.use_operands(ExtractReg)) {
+      Changed = true;
+      Use.setReg(OrigReg);
+    }
+  }
+
+  if (AllDefsReplaced) {
+    // If SeqI was the next instruction in the BB and we removed it, we'd break
+    // the outer iteration.
+    assert(std::next(MachineBasicBlock::iterator(MI)) != SeqI &&
+           "G_SEQUENCE does not dominate G_EXTRACT");
+
+    MI.eraseFromParent();
+
+    if (MRI.use_empty(SrcReg))
+      SeqI.eraseFromParent();
+    Changed = true;
+  }
+
+  return Changed;
+}
+
 bool MachineLegalizePass::runOnMachineFunction(MachineFunction &MF) {
   // If the ISel pipeline failed, do not bother running that pass.
   if (MF.getProperties().hasProperty(
@@ -94,5 +159,19 @@ bool MachineLegalizePass::runOnMachineFunction(MachineFunction &MF) {
 
       Changed |= Res == MachineLegalizeHelper::Legalized;
     }
+
+
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+  for (auto &MBB : MF) {
+    for (auto MI = MBB.begin(); MI != MBB.end(); MI = NextMI) {
+      // Get the next Instruction before we try to legalize, because there's a
+      // good chance MI will be deleted.
+      NextMI = std::next(MI);
+
+      Changed |= combineExtracts(*MI, MRI, TII);
+    }
+  }
+
   return Changed;
 }