From a964099ce5e560b0c4431938b2b28d131cd59bf5 Mon Sep 17 00:00:00 2001
From: Ivan Kosarev <ivan.kosarev@amd.com>
Date: Fri, 30 Sep 2022 14:24:08 +0100
Subject: [PATCH] [AMDGPU][SetWavePriority] Fix dealing with MBBInfo records.

Happened earlier than I anticipated. :)

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D134726
---
 llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp
index da2a2b9..64f1bd5 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSetWavePriority.cpp
@@ -122,13 +122,13 @@ bool AMDGPUSetWavePriority::runOnMachineFunction(MachineFunction &MF) {
   // executed provided no backedge is ever taken.
   MBBInfoSet MBBInfos;
   for (MachineBasicBlock *MBB : post_order(&MF)) {
-    MBBInfo &Info = MBBInfos[MBB];
     bool AtStart = true;
     unsigned MaxNumVALUInstsInMiddle = 0;
     unsigned NumVALUInstsAtEnd = 0;
     for (MachineInstr &MI : *MBB) {
       if (isVMEMLoad(MI)) {
         AtStart = false;
+        MBBInfo &Info = MBBInfos[MBB];
         Info.NumVALUInstsAtStart = 0;
         MaxNumVALUInstsInMiddle = 0;
         NumVALUInstsAtEnd = 0;
@@ -140,7 +140,7 @@ bool AMDGPUSetWavePriority::runOnMachineFunction(MachineFunction &MF) {
         NumVALUInstsAtEnd = 0;
       } else if (SIInstrInfo::isVALU(MI)) {
         if (AtStart)
-          ++Info.NumVALUInstsAtStart;
+          ++MBBInfos[MBB].NumVALUInstsAtStart;
         ++NumVALUInstsAtEnd;
       }
     }
@@ -152,6 +152,7 @@ bool AMDGPUSetWavePriority::runOnMachineFunction(MachineFunction &MF) {
       NumFollowingVALUInsts =
           std::max(NumFollowingVALUInsts, MBBInfos[Succ].NumVALUInstsAtStart);
     }
+    MBBInfo &Info = MBBInfos[MBB];
     if (AtStart)
       Info.NumVALUInstsAtStart += NumFollowingVALUInsts;
     NumVALUInstsAtEnd += NumFollowingVALUInsts;
-- 
2.7.4