[llvm] [CodeGen] Introduce MIR-level target-independent rematerialization helper (PR #177080)
Lucas Ramirez via llvm-commits
llvm-commits at lists.llvm.org
Thu Feb 19 05:00:36 PST 2026
================
@@ -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);
+ unsigned UserRegion = MIRegion[&UserMI];
+ Regs[FromRegIdx].eraseUser(&UserMI, UserRegion);
+ Regs[ToRegIdx].addUser(&UserMI, UserRegion);
+ deleteRegIfUnused(FromRegIdx);
+}
+
+void Rematerializer::transferRegionUsers(unsigned FromRegIdx, unsigned ToRegIdx,
+ unsigned UseRegion) {
+ auto &FromRegUsers = Regs[FromRegIdx].Uses;
+ auto UsesIt = FromRegUsers.find(UseRegion);
+ if (UsesIt == FromRegUsers.end())
+ return;
+
+ const SmallDenseSet<MachineInstr *, 4> &RegionUsers = UsesIt->getSecond();
+ for (MachineInstr *UserMI : RegionUsers)
+ transferUserInternal(FromRegIdx, ToRegIdx, *UserMI);
+ Regs[ToRegIdx].addUsers(RegionUsers, UseRegion);
+ FromRegUsers.erase(UseRegion);
+ deleteRegIfUnused(FromRegIdx);
+}
+
+void Rematerializer::transferUserInternal(unsigned FromRegIdx,
+ unsigned ToRegIdx,
+ MachineInstr &UserMI) {
+ assert(MIRegion.contains(&UserMI) && "unknown user");
+ assert(getReg(FromRegIdx).Uses.at(MIRegion.at(&UserMI)).contains(&UserMI) &&
+ "not a user");
+ assert(FromRegIdx != ToRegIdx && "identical registers");
+ assert(getOriginOrSelf(FromRegIdx) == getOriginOrSelf(ToRegIdx) &&
+ "unrelated registers");
+
+ LLVM_DEBUG(rdbgs() << "User transfer from " << printID(FromRegIdx) << " to "
+ << printID(ToRegIdx) << ": " << printUser(&UserMI)
+ << '\n');
+
+ UserMI.substituteRegister(getReg(FromRegIdx).getDefReg(),
+ getReg(ToRegIdx).getDefReg(), 0, TRI);
+ LISUpdates.insert(FromRegIdx);
+ LISUpdates.insert(ToRegIdx);
+
+ // If the user is rematerializable, we must change its dependency to the
+ // new register.
+ if (unsigned UserRegIdx = getDefRegIdx(UserMI); UserRegIdx != NoReg) {
+ // Look for the user's dependency that matches the register.
+ for (Reg::Dependency &Dep : Regs[UserRegIdx].Dependencies) {
+ if (Dep.RegIdx == FromRegIdx) {
+ Dep.RegIdx = ToRegIdx;
+ return;
+ }
+ }
+ llvm_unreachable("broken dependency");
+ }
+}
+
+void Rematerializer::updateLiveIntervals() {
+ DenseSet<Register> SeenUnrematRegs;
+ for (unsigned RegIdx : LISUpdates) {
+ const Reg &UpdateReg = getReg(RegIdx);
+ assert(UpdateReg.DefMI || Rollbackable.contains(RegIdx) && "dead register");
+
+ Register DefReg = UpdateReg.getDefReg();
+ if (LIS.hasInterval(DefReg))
+ LIS.removeInterval(DefReg);
+ LIS.createAndComputeVirtRegInterval(DefReg);
+
+ LLVM_DEBUG({
+ rdbgs() << "Re-computed interval for " << printID(RegIdx) << ": ";
+ LIS.getInterval(DefReg).print(dbgs());
+ rdbgs() << '\n' << printRegUsers(RegIdx);
+ });
+
+ // Update intervals for unrematerializable operands.
+ for (unsigned MOIdx : getUnrematableOprds(RegIdx)) {
+ Register UnrematReg = UpdateReg.DefMI->getOperand(MOIdx).getReg();
+ if (!SeenUnrematRegs.insert(UnrematReg).second)
+ continue;
+ LIS.removeInterval(UnrematReg);
+ LIS.createAndComputeVirtRegInterval(UnrematReg);
+ LLVM_DEBUG(
+ dbgs() << " Re-computed interval for register "
+ << printReg(UnrematReg, &TRI,
+ UpdateReg.DefMI->getOperand(MOIdx).getSubReg(),
+ &MRI)
+ << '\n');
+ }
+ }
+ LISUpdates.clear();
+}
+
+void Rematerializer::commitRematerializations() {
+ for (auto &[RegIdx, _] : Rollbackable)
+ deleteReg(RegIdx);
+ Rollbackable.clear();
+}
+
+bool Rematerializer::isMOIdenticalAtUses(MachineOperand &MO,
+ ArrayRef<SlotIndex> Uses) const {
+ if (Uses.empty())
+ return true;
----------------
lucas-rami wrote:
I don't necessarily expect empty uses to be passed here but would prefer not to disallow it. To me this essentially checks that "property X is true for all uses", so if there are no uses the statement is vacuously true.
https://github.com/llvm/llvm-project/pull/177080
More information about the llvm-commits
mailing list