Works on simple testcase

Author: linuxrocks123

llvm/lib/Target/AMDGPU/AMDGPU.h

@@ -248,6 +248,9 @@ extern char &AMDGPUPreloadKernArgPrologLegacyID;
 void initializeAMDGPUPreloadKernelArgumentsLegacyPass(PassRegistry &);
 extern char &AMDGPUPreloadKernelArgumentsLegacyID;
 
+void initializeSIRestoreNormalEpilogLegacyPass(PassRegistry &);
+extern char &SIRestoreNormalEpilogLegacyID;
+
 // Passes common to R600 and SI
 FunctionPass *createAMDGPUPromoteAlloca();
 void initializeAMDGPUPromoteAllocaPass(PassRegistry&);

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

@@ -597,6 +597,7 @@ extern "C" LLVM_ABI LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
   initializeSIOptimizeExecMaskingLegacyPass(*PR);
   initializeSIPreAllocateWWMRegsLegacyPass(*PR);
   initializeSIFormMemoryClausesLegacyPass(*PR);
+  initializeSIRestoreNormalEpilogLegacyPass(*PR);
   initializeSIPostRABundlerLegacyPass(*PR);
   initializeGCNCreateVOPDLegacyPass(*PR);
   initializeAMDGPUUnifyDivergentExitNodesPass(*PR);
@@ -1595,6 +1596,8 @@ void GCNPassConfig::addOptimizedRegAlloc() {
 
   if (EnableRewritePartialRegUses)
     insertPass(&RenameIndependentSubregsID, &GCNRewritePartialRegUsesID);
+  
+  insertPass(&RenameIndependentSubregsID,&SIRestoreNormalEpilogLegacyID);
 
   if (isPassEnabled(EnablePreRAOptimizations))
     insertPass(&MachineSchedulerID, &GCNPreRAOptimizationsID);

llvm/lib/Target/AMDGPU/CMakeLists.txt

@@ -183,6 +183,7 @@ add_llvm_target(AMDGPUCodeGen
   SIPreEmitPeephole.cpp
   SIProgramInfo.cpp
   SIRegisterInfo.cpp
+  SIRestoreNormalEpilog.cpp
   SIShrinkInstructions.cpp
   SIWholeQuadMode.cpp
 

llvm/lib/Target/AMDGPU/SICustomBranchBundles.h

@@ -91,9 +91,7 @@ get_epilog_for_successor(MachineBasicBlock& pred_MBB, MachineBasicBlock& succ_MB
 
   for (MachineInstr &branch_MI : reverse(pred_MBB.instrs()))
     if (branch_MI.isBranch() && TII.getBranchDestBlock(branch_MI) == &succ_MBB)
-      return std::next(branch_MI.getIterator());
-    else
-      assert(branch_MI.isBranch() && "Shouldn't have fall-throughs here.");
+      return ++Epilog_Iterator(branch_MI.getIterator());
 
   llvm_unreachable("There should always be a branch to succ_MBB.");
 }
@@ -132,8 +130,8 @@ static inline void normalize_ir_post_phi_elimination(MachineFunction &MF) {
   for (MachineBasicBlock &MBB : MF) {
     CFG_Rewrite_Entry to_insert = {{}, &MBB, {}};
     for (MachineBasicBlock *pred_MBB : MBB.predecessors()) {
-      MachineBasicBlock::instr_iterator ep_it =
-          get_epilog_for_successor(*pred_MBB, MBB)->getIterator();
+      Epilog_Iterator ep_it =
+          get_epilog_for_successor(*pred_MBB, MBB);
 
       vector<MachineInstr *> epilog;
       while (!ep_it.isEnd())
@@ -175,6 +173,7 @@ static inline void normalize_ir_post_phi_elimination(MachineFunction &MF) {
   // Perform the journaled rewrites.
   for (auto &entry : cfg_rewrite_entries) {
     MachineBasicBlock *mezzanine_MBB = MF.CreateMachineBasicBlock();
+    MF.insert(MF.end(),mezzanine_MBB);
 
     // Deal with mezzanine to successor succession.
     BuildMI(mezzanine_MBB, DebugLoc(), TII.get(AMDGPU::S_BRANCH)).addMBB(entry.succ_MBB);
@@ -192,7 +191,7 @@ static inline void normalize_ir_post_phi_elimination(MachineFunction &MF) {
       pred_MBB->replaceSuccessor(entry.succ_MBB, mezzanine_MBB);
 
       // Delete instructions that were lowered from epilog
-      auto epilog_it = get_epilog_for_successor(*pred_MBB, *entry.succ_MBB);
+      auto epilog_it = ++Epilog_Iterator(branch_ins.getIterator());
       while (!epilog_it.isEnd())
         epilog_it++->eraseFromBundle();
     }
@@ -228,7 +227,8 @@ static inline void hoist_unrelated_copies(MachineFunction &MF) {
       while (!copy_move_it.isEnd()) {
         Epilog_Iterator next = copy_move_it; ++next;
         if (copy_move_it->getOpcode() == AMDGPU::COPY &&
-            !related_copy_sources.count(copy_move_it->getOperand(1).getReg())) {
+            !related_copy_sources.count(copy_move_it->getOperand(1).getReg())
+            || copy_move_it->getOpcode() == AMDGPU::IMPLICIT_DEF) {
           MachineInstr &MI_to_move = *copy_move_it;
           MI_to_move.removeFromBundle();
           MBB.insert(branch_MI.getIterator(),&MI_to_move);

llvm/lib/Target/AMDGPU/SIRestoreNormalEpilog.cpp

@@ -0,0 +1,46 @@
+#include "SICustomBranchBundles.h"
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/CodeGen/LiveIntervals.h"
+#include "llvm/CodeGen/LiveVariables.h"
+#include "llvm/CodeGen/MachineDominators.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
+#include "llvm/Target/TargetMachine.h"
+
+#define DEBUG_TYPE "si-restore-normal-epilog"
+
+namespace
+{
+
+class SIRestoreNormalEpilogLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  SIRestoreNormalEpilogLegacy() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override {
+    hoist_unrelated_copies(MF);
+    normalize_ir_post_phi_elimination(MF);
+    return true;
+  }
+
+  StringRef getPassName() const override {
+    return "SI Restore Normal Epilog Post PHI Elimination";
+  }
+
+  MachineFunctionProperties getRequiredProperties() const override {
+    return MachineFunctionProperties().setNoPHIs();
+  }
+
+};
+
+} // namespace
+
+INITIALIZE_PASS(SIRestoreNormalEpilogLegacy, DEBUG_TYPE,
+                "SI restore normal epilog", false, false)
+
+char SIRestoreNormalEpilogLegacy::ID;
+char &llvm::SIRestoreNormalEpilogLegacyID = SIRestoreNormalEpilogLegacy::ID;