[llvm] Add late optimization pass for riscv (PR #117060)

[Philip Reames via llvm-commits]

================
@@ -0,0 +1,157 @@
+//===-- RISCVLatePeephole.cpp - Late stage peephole optimization ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file provides RISC-V specific target descriptions.
+///
+//===----------------------------------------------------------------------===//
+
+#include "MCTargetDesc/RISCVMCTargetDesc.h"
+#include "RISCV.h"
+#include "RISCVInstrInfo.h"
+#include "RISCVSubtarget.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
+#include "llvm/CodeGen/MachineDominators.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/RegisterScavenging.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/Debug.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "riscv-late-peephole"
+#define RISCV_LATE_PEEPHOLE_NAME "RISC-V Late Stage Peephole"
+
+namespace {
+
+struct RISCVLatePeepholeOpt : public MachineFunctionPass {
+  static char ID;
+
+  RISCVLatePeepholeOpt() : MachineFunctionPass(ID) {}
+
+  StringRef getPassName() const override { return RISCV_LATE_PEEPHOLE_NAME; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineDominatorTreeWrapperPass>();
+    AU.addPreserved<MachineDominatorTreeWrapperPass>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  bool runOnMachineFunction(MachineFunction &Fn) override;
+
+private:
+  void removeBlock(MachineBasicBlock *B, MachineBasicBlock *NewB);
+  bool removeSingleBranchBlock(MachineBasicBlock *B);
+
+  const RISCVInstrInfo *RII = nullptr;
+  MachineFunction *MFN = nullptr;
+  MachineDominatorTree *MDT = nullptr;
+};
+} // namespace
+
+char RISCVLatePeepholeOpt::ID = 0;
+INITIALIZE_PASS(RISCVLatePeepholeOpt, "riscv-late-peephole",
+                RISCV_LATE_PEEPHOLE_NAME, false, false)
+
+void RISCVLatePeepholeOpt::removeBlock(MachineBasicBlock *B, MachineBasicBlock *NewB) {
+  LLVM_DEBUG(dbgs() << "Removing block '#'" << B->getNumber() << "\n");
+
+  // Transfer the immediate dominator information from B to its descendants.
+  MachineDomTreeNode *N = MDT->getNode(B);
+  MachineDomTreeNode *IDN = N->getIDom();
+  if (IDN) {
+    MachineBasicBlock *IDB = IDN->getBlock();
+
+    using GTN = GraphTraits<MachineDomTreeNode *>;
+    using DTNodeVectType = SmallVector<MachineDomTreeNode *, 4>;
+
+    DTNodeVectType Cn(GTN::child_begin(N), GTN::child_end(N));
+    for (auto &I : Cn) {
+      MachineBasicBlock *SB = I->getBlock();
+      MDT->changeImmediateDominator(SB, IDB);
+    }
+  }
+
+  while (!B->succ_empty())
+    B->removeSuccessor(B->succ_begin());
+
+  for (MachineBasicBlock *Pred : B->predecessors()) {
+    Pred->removeSuccessor(B, true);
+    // TODO: how do I calculate the branch probability here?
+    if (Pred != NewB)
+      Pred->addSuccessor(NewB);
+  }
+
+  MDT->eraseNode(B);
+  MFN->erase(B->getIterator());
+}
+
+bool RISCVLatePeepholeOpt::removeSingleBranchBlock(MachineBasicBlock *B) {
+  LLVM_DEBUG(dbgs() << "Checking flow pattern at " << printMBBReference(*B)
+                    << "\n");
+
+  if (B->size() != 1)
+    return false;
+
+  MachineBasicBlock::const_iterator T1I = B->getFirstTerminator();
+  if (T1I == B->end())
+    return false;
+  unsigned Opc = T1I->getOpcode();
+  if (Opc != RISCV::BEQ && Opc != RISCV::BNE)
+    return false;
+
+  Register DstReg = T1I->getOperand(0).getReg();
+  Register SrcReg = T1I->getOperand(1).getReg();
+  if (DstReg != SrcReg)
----------------
preames wrote:

Even if we keep this as a separate late peephole, I'd suggest splitting this into two transforms:
1) Transform one does the bne zero, zero -> jump conversion.
2) Transform two does the jump threading of an edge through a single basis block.

Doing it that way should make the code easier to read, and test.  It'll also be more powerful.  

Note that we probably won't keep the separate pass, so I wouldn't necessarily spend time implementing this just yet.

https://github.com/llvm/llvm-project/pull/117060