[llvm] [CodeGen] Introduce MIR-level target-independent rematerialization helper (PR #177080)

[Quentin Colombet via llvm-commits]

================
@@ -0,0 +1,737 @@
+//=====-- Rematerializer.cpp - MIR rematerialization support ----*- C++ -*-===//
+//
+// 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
+//
+//==-----------------------------------------------------------------------===//
+//
+/// \file
+/// Implements helpers for target-independent rematerialization at the MIR
+/// level.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/Rematerializer.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/CodeGen/LiveIntervals.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/Register.h"
+#include "llvm/CodeGen/TargetOpcodes.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "rematerializer"
+
+using namespace llvm;
+
+static bool isAvailableAtUse(const VNInfo *OVNI, LaneBitmask Mask,
+                             SlotIndex UseIdx, const LiveInterval &LI) {
+  assert(OVNI);
+  if (OVNI != LI.getVNInfoAt(UseIdx))
+    return false;
+
+  // Check that subrange is live at user.
+  if (LI.hasSubRanges()) {
+    for (const LiveInterval::SubRange &SR : LI.subranges()) {
+      if ((SR.LaneMask & Mask).none())
+        continue;
+      if (!SR.liveAt(UseIdx))
+        return false;
+
+      // Early exit if all used lanes are checked. No need to continue.
+      Mask &= ~SR.LaneMask;
+      if (Mask.none())
+        break;
+    }
+  }
+  return true;
+}
+
+static Register isRegDependency(const MachineOperand &MO) {
+  if (!MO.isReg() || !MO.readsReg())
+    return Register();
+  Register Reg = MO.getReg();
+  if (Reg.isPhysical()) {
+    // By the requirements on trivially rematerializable instructions, a
+    // physical register use is either constant or ignorable.
+    return Register();
+  }
+  return Reg;
+}
+
+unsigned Rematerializer::rematerializeToRegion(unsigned RootIdx,
+                                               unsigned UseRegion,
+                                               DependencyReuseInfo &DRI) {
+
+  MachineInstr *FirstMI =
+      getReg(RootIdx).getRegionUseBounds(UseRegion, LIS).first;
+  unsigned NewRegIdx = rematerializeToPos(RootIdx, FirstMI, DRI);
+  transferRegionUsers(RootIdx, NewRegIdx, UseRegion);
+  return NewRegIdx;
+}
+
+unsigned
+Rematerializer::rematerializeToPos(unsigned RootIdx,
+                                   MachineBasicBlock::iterator InsertPos,
+                                   DependencyReuseInfo &DRI) {
+  LLVM_DEBUG({
+    rdbgs() << "Rematerializing " << printID(RootIdx) << " to "
+            << printUser(&*InsertPos) << '\n';
+    ++CallDepth;
+  });
+
+  // Create/identify dependencies for the new register. Copy the dependencies
+  // vector because underlying updates to the backing vector of registers may
+  // invalidate references.
+  SmallVector<Reg::Dependency, 2> NewDeps, Deps(Regs[RootIdx].Dependencies);
+  for (const Reg::Dependency &Dep : Deps) {
+    if (auto NewDep = DRI.DependencyMap.find(Dep.RegIdx);
+        NewDep != DRI.DependencyMap.end()) {
+      // We already have the version of the dependency we want to use.
+      NewDeps.emplace_back(Dep.MOIdx, NewDep->second);
+    } else {
+      // Dependencies must be rematerialized in def-use order.
+      unsigned NewDepIdx = rematerializeToPos(Dep.RegIdx, InsertPos, DRI);
+      DRI.DependencyMap.insert({Dep.RegIdx, NewDepIdx});
+      NewDeps.emplace_back(Dep.MOIdx, NewDepIdx);
+    }
+  }
+
+  LLVM_DEBUG(--CallDepth);
+  return rematerializeReg(RootIdx, InsertPos, std::move(NewDeps));
+}
+
+void Rematerializer::rollbackRematsOf(unsigned RootIdx) {
+  auto Remats = Rematerializations.find(RootIdx);
+  if (Remats == Rematerializations.end())
+    return;
+
+  LLVM_DEBUG({
+    rdbgs() << "Rolling back rematerializations of " << printID(RootIdx)
+            << '\n';
+  });
+
+  reviveRegIfDead(RootIdx);
+  // All of the rematerialization's users must use the revived register.
+  for (unsigned RematRegIdx : Remats->getSecond()) {
+    for (const auto &[UseRegion, RegionUsers] : Regs[RematRegIdx].Uses)
+      transferRegionUsers(RematRegIdx, RootIdx, UseRegion);
+  }
+  Rematerializations.erase(RootIdx);
+
+  LLVM_DEBUG({
+    rdbgs() << "** Rolled back rematerializations of " << printID(RootIdx)
+            << '\n';
+  });
+}
+
+void Rematerializer::rollback(unsigned RematIdx) {
+  assert(getReg(RematIdx).DefMI && !Rollbackable.contains(RematIdx) &&
+         "cannot rollback dead register");
+  const unsigned OriginRegIdx = getOriginOf(RematIdx);
+  reviveRegIfDead(OriginRegIdx);
+  for (const auto &[UseRegion, RegionUsers] : Regs[RematIdx].Uses)
+    transferRegionUsers(RematIdx, OriginRegIdx, UseRegion);
+}
+
+void Rematerializer::reviveRegIfDead(unsigned RootIdx) {
+  assert(!isRematerializedRegister(RootIdx) &&
+         "cannot revive rematerialization");
+
+  Reg &Root = Regs[RootIdx];
+  if (!Root.Uses.empty()) {
+    // The register still exists, nothing to do.
+    LLVM_DEBUG(rdbgs() << printID(RootIdx) << " still exists\n");
+    return;
+  }
+
+  assert(Rollbackable.contains(RootIdx) && "not marked rollbackable");
+  assert(Root.DefMI && Root.DefMI->getOpcode() == TargetOpcode::DBG_VALUE &&
+         "not the right opcode");
+  assert(Rematerializations.contains(RootIdx) && "no remats");
+
+  LLVM_DEBUG({
+    rdbgs() << "Reviving " << printID(RootIdx) << '\n';
+    ++CallDepth;
+  });
+
+  // Fully rematerialized dependencies need to be revived. All dependencies gain
+  // a new user.
+  for (const Reg::Dependency &Dep : Root.Dependencies) {
+    reviveRegIfDead(Dep.RegIdx);
+    Regs[Dep.RegIdx].addUser(Root.DefMI, Root.DefRegion);
+    LISUpdates.insert(Dep.RegIdx);
+  }
+
+  // Pick any rematerialization to retrieve the original opcode from.
+  unsigned RematIdx = *Rematerializations.at(RootIdx).begin();
+  Root.DefMI->setDesc(getReg(RematIdx).DefMI->getDesc());
+  for (const auto &[MOIdx, Reg] : Rollbackable.at(RootIdx))
+    Root.DefMI->getOperand(MOIdx).setReg(Reg);
+  Rollbackable.erase(RootIdx);
+  LISUpdates.insert(RootIdx);
+
+  LLVM_DEBUG({
+    rdbgs() << "** Revived " << printID(RootIdx) << " @ ";
+    LIS.getInstructionIndex(*Root.DefMI).print(dbgs());
+    dbgs() << '\n';
+    --CallDepth;
+  });
+}
+
+void Rematerializer::transferUser(unsigned FromRegIdx, unsigned ToRegIdx,
+                                  MachineInstr &UserMI) {
+  transferUserInternal(FromRegIdx, ToRegIdx, UserMI);
----------------
qcolombet wrote:

Nit: instead of `Internal` we usually go with `Impl`

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