From 61c1f9a2901ebdcba6090b69f889a1c9267d5ceb Mon Sep 17 00:00:00 2001
From: Edd Barrett <vext01@gmail.com>
Date: Thu, 12 Dec 2024 13:37:43 +0000
Subject: [PATCH] Ensure spillmap entries are properly killed where necessary.

In a couple of places we were failing to clear associations between
locations where an MIR instruction defined or killed a register. This
lead to us deopting dead stale values at runtime.

Co-authored-by: Lukas Diekmann <lukas.diekmann@gmail.com>
---
 llvm/lib/Target/X86/X86AsmPrinter.cpp         | 36 +++++++++++++++----
 .../test/CodeGen/X86/yk-stackmaps-extrareg.ll |  4 +--
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Target/X86/X86AsmPrinter.cpp b/llvm/lib/Target/X86/X86AsmPrinter.cpp
index 93c473f17ea32a..99335f5152c082 100644
--- a/llvm/lib/Target/X86/X86AsmPrinter.cpp
+++ b/llvm/lib/Target/X86/X86AsmPrinter.cpp
@@ -46,6 +46,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Target/TargetMachine.h"
 
+#include <map>
 #include <cstdint>
 #include <set>
 
@@ -62,6 +63,7 @@ X86AsmPrinter::X86AsmPrinter(TargetMachine &TM,
 
 const TargetInstrInfo *TII;
 const TargetRegisterInfo *TRI;
+const MCRegisterInfo *MRI;
 
 /// Clear any mappings that map to the given register.
 void clearRhs(Register Reg, std::map<Register, std::set<int64_t>> &SpillMap) {
@@ -82,6 +84,12 @@ void clearOffset(unsigned int Reg, int Offset, std::map<Register, std::set<int64
   }
 }
 
+/// Remove all mentions of the DWARF register `DwReg` from `SpillMap`.
+void killRegister(const Register DwReg, std::map<Register, std::set<int64_t>> &SpillMap) {
+    SpillMap.erase(DwReg);
+    clearRhs(DwReg, SpillMap);
+}
+
 /// Given a MachineBasicBlock, analyse its instructions to build a mapping of
 /// where duplicate values live. This can be either in another register or on
 /// the stack. Since registers are always positive and stack offsets negative,
@@ -98,6 +106,12 @@ void processInstructions(
     // encoded in the stackmap when it is lowered.
     if (Instr.getOpcode() == TargetOpcode::STACKMAP) {
       StackmapSpillMaps[&Instr] = SpillMap;
+      for (const MachineOperand MO : Instr.uses()) {
+        if (MO.isReg() && MO.isKill()) {
+          auto DwReg = getDwarfRegNum(MO.getReg(), TRI);
+          killRegister(DwReg, SpillMap);
+        }
+      }
       continue;
     }
 
@@ -105,6 +119,12 @@ void processInstructions(
     // moment we desire, there's no need to compute a spillmap for them nor do
     // we have to "patch them up". We can just skip them.
     if (Instr.getOpcode() == TargetOpcode::PATCHPOINT) {
+      for (const MachineOperand MO : Instr.uses()) {
+        if (MO.isReg() && MO.isKill()) {
+          auto DwReg = getDwarfRegNum(MO.getReg(), TRI);
+          killRegister(DwReg, SpillMap);
+        }
+      }
       continue;
     }
 
@@ -133,6 +153,9 @@ void processInstructions(
       SpillMap[LhsDwReg].insert(Other.begin(), Other.end());
       // Also add Lhs to the mapping of Rhs.
       SpillMap[RhsDwReg].insert(LhsDwReg);
+      if (Rhs.isKill()) {
+        killRegister(RhsDwReg, SpillMap);
+      }
       continue;
     }
 
@@ -175,18 +198,18 @@ void processInstructions(
     for (const MachineOperand MO : Instr.defs()) {
       assert(MO.isReg() && "Is register.");
       auto DwReg = getDwarfRegNum(MO.getReg(), TRI);
-      SpillMap.erase(DwReg);
-      clearRhs(DwReg, SpillMap);
+      killRegister(DwReg, SpillMap);
     }
 
     // Delete registers that are "killed" (no longer live) after this
     // intsruction. This prevents us deopting values at runtime that will never
     // be used.
     for (const MachineOperand MO : Instr.uses()) {
-      if (MO.isReg() && MO.isKill()) {
-        auto DwReg = getDwarfRegNum(MO.getReg(), TRI);
-        SpillMap.erase(DwReg);
-        clearRhs(DwReg, SpillMap);
+      if (MO.isReg() && (MO.isKill() || MO.isDef())) {
+        int DwReg = MRI->getDwarfRegNum(MO.getReg(), false);
+        if (DwReg >= 0) {
+          killRegister(DwReg, SpillMap);
+        }
       }
     }
   }
@@ -218,6 +241,7 @@ void findSpillLocations(
 ///
 bool X86AsmPrinter::runOnMachineFunction(MachineFunction &MF) {
   TRI = MF.getSubtarget().getRegisterInfo();
+  MRI = TM.getMCRegisterInfo();
   Subtarget = &MF.getSubtarget<X86Subtarget>();
 
   SMShadowTracker.startFunction(MF);
diff --git a/llvm/test/CodeGen/X86/yk-stackmaps-extrareg.ll b/llvm/test/CodeGen/X86/yk-stackmaps-extrareg.ll
index 36f278925aac8e..760b7e06c89492 100644
--- a/llvm/test/CodeGen/X86/yk-stackmaps-extrareg.ll
+++ b/llvm/test/CodeGen/X86/yk-stackmaps-extrareg.ll
@@ -12,11 +12,9 @@
 ; NOTE: Reserved
 ; CHECK-NEXT: .short 0
 ; NOTE: Number of extra locations.
-; CHECK-NEXT: .short 2
+; CHECK-NEXT: .short 1
 ; NOTE: Stack offset this value is stored in.
 ; CHECK-NEXT: .short -80
-; NOTE: Extra register this value is stored in.
-; CHECK-NEXT: .short 8
 
 source_filename = "ld-temp.o"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"