[llvm] [ReachingDefAnalysis] Fix management of MBBFrameObjsReachingDefs (PR #124943)

[Michael Maitland via llvm-commits]

https://github.com/michaelmaitland updated https://github.com/llvm/llvm-project/pull/124943

>From f14a174452c752f5027929d66a6d6d31bb8d9cae Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Thu, 30 Jan 2025 07:20:05 -0800
Subject: [PATCH 1/2] precommit test case

---
 llvm/test/CodeGen/RISCV/rda-stack.mir | 37 +++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/llvm/test/CodeGen/RISCV/rda-stack.mir b/llvm/test/CodeGen/RISCV/rda-stack.mir
index 5f4974181c1cd5f..fc5bd8acd2840a8 100644
--- a/llvm/test/CodeGen/RISCV/rda-stack.mir
+++ b/llvm/test/CodeGen/RISCV/rda-stack.mir
@@ -149,3 +149,40 @@ body:             |
     $x10 = LD %stack.0, 0 :: (load (s64))
     PseudoRET implicit $x10
 ...
+---
+name:            test4
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: default, offset: 0, size: 4, alignment: 4,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+  - { id: 1, name: '', type: default, offset: 0, size: 4, alignment: 4,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+body:             |
+  ; CHECK: RDA results for test4
+  ; CHECK-NEXT: $x10:{ }
+  ; CHECK-NEXT: %stack.0:{ }
+  ; CHECK-NEXT: 0: SD $x10, %stack.0, 0 :: (store (s64))
+  ; CHECK-EMPTY:
+  ; CHECK-NEXT: $x11:{ }
+  ; CHECK-NEXT: %stack.0:{ 0 }
+  ; CHECK-NEXT: 1: SD $x11, %stack.0, 0 :: (store (s64))
+  ; CHECK-EMPTY:
+  ; CHECK-NEXT: $x10:{ }
+  ; CHECK-NEXT: %stack.1:{ }
+  ; CHECK-NEXT: 2: SD $x10, %stack.1, 0 :: (store (s64))
+  ; CHECK-EMPTY:
+  ; CHECK-NEXT: $x11:{ }
+  ; CHECK-NEXT: %stack.1:{ }
+  ; CHECK-NEXT: 3: SD $x11, %stack.1, 0 :: (store (s64))
+  ; CHECK-EMPTY:
+  ; CHECK-NEXT: 4: PseudoRET
+  bb.0.entry:
+    liveins: $x10, $x11
+    SD $x10, %stack.0, 0 :: (store (s64))
+    SD $x11, %stack.0, 0 :: (store (s64))
+    SD $x10, %stack.1, 0 :: (store (s64))
+    SD $x11, %stack.1, 0 :: (store (s64))
+    PseudoRET
+...

>From af0058bf5257908b181c00f25b7528ac7ea3b32c Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Wed, 29 Jan 2025 08:39:09 -0800
Subject: [PATCH 2/2] [ReachingDefAnalysis][NFC] Fix management of
 MBBFrameObjsReachingDefs

---
 llvm/lib/CodeGen/ReachingDefAnalysis.cpp | 28 ++++++++++++++----------
 llvm/test/CodeGen/RISCV/rda-stack.mir    |  2 +-
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/CodeGen/ReachingDefAnalysis.cpp b/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
index 08121f411640c87..a2adc5ede288903 100644
--- a/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
+++ b/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
@@ -147,16 +147,9 @@ void ReachingDefAnalysis::processDefs(MachineInstr *MI) {
       assert(FrameIndex >= 0 && "Can't handle negative frame indicies yet!");
       if (!isFIDef(*MI, FrameIndex, TII))
         continue;
-      if (MBBFrameObjsReachingDefs.contains(MBBNumber)) {
-        auto Frame2InstrIdx = MBBFrameObjsReachingDefs[MBBNumber];
-        if (Frame2InstrIdx.count(FrameIndex - ObjectIndexBegin) > 0)
-          Frame2InstrIdx[FrameIndex - ObjectIndexBegin].push_back(CurInstr);
-        else
-          Frame2InstrIdx[FrameIndex - ObjectIndexBegin] = {CurInstr};
-      } else {
-        MBBFrameObjsReachingDefs[MBBNumber] = {
-            {FrameIndex - ObjectIndexBegin, {CurInstr}}};
-      }
+
+      int Key = FrameIndex - ObjectIndexBegin;
+      MBBFrameObjsReachingDefs[MBBNumber][Key].push_back(CurInstr);
     }
     if (!isValidRegDef(MO))
       continue;
@@ -348,8 +341,19 @@ int ReachingDefAnalysis::getReachingDef(MachineInstr *MI, Register Reg) const {
 
   if (Reg.isStack()) {
     int FrameIndex = Reg.stackSlotIndex();
-    for (int Def : MBBFrameObjsReachingDefs.lookup(MBBNumber).lookup(
-             FrameIndex - ObjectIndexBegin)) {
+    int Key = FrameIndex - ObjectIndexBegin;
+
+    // Check that there was a reaching def.
+    auto Lookup1 = MBBFrameObjsReachingDefs.find(MBBNumber);
+    if (Lookup1 == MBBFrameObjsReachingDefs.end())
+      return LatestDef;
+    auto &InnerMap = Lookup1->second;
+    auto Lookup2 = InnerMap.find(Key);
+    if (Lookup2 == InnerMap.end())
+      return LatestDef;
+    auto &Defs = Lookup2->second;
+
+    for (int Def : Defs) {
       if (Def >= InstId)
         break;
       DefRes = Def;
diff --git a/llvm/test/CodeGen/RISCV/rda-stack.mir b/llvm/test/CodeGen/RISCV/rda-stack.mir
index fc5bd8acd2840a8..3ec17615fa67760 100644
--- a/llvm/test/CodeGen/RISCV/rda-stack.mir
+++ b/llvm/test/CodeGen/RISCV/rda-stack.mir
@@ -174,7 +174,7 @@ body:             |
   ; CHECK-NEXT: 2: SD $x10, %stack.1, 0 :: (store (s64))
   ; CHECK-EMPTY:
   ; CHECK-NEXT: $x11:{ }
-  ; CHECK-NEXT: %stack.1:{ }
+  ; CHECK-NEXT: %stack.1:{ 2 }
   ; CHECK-NEXT: 3: SD $x11, %stack.1, 0 :: (store (s64))
   ; CHECK-EMPTY:
   ; CHECK-NEXT: 4: PseudoRET