diff --git a/llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp b/llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp
index 38f86ca3e9afc..2f1803ee6df59 100644
--- a/llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp
+++ b/llvm/lib/Target/AMDGPU/SIWholeQuadMode.cpp
@@ -218,8 +218,8 @@ class SIWholeQuadMode : public MachineFunctionPass {
                             bool IsWQM);
   MachineInstr *lowerKillF32(MachineBasicBlock &MBB, MachineInstr &MI);
 
-  void lowerBlock(MachineBasicBlock &MBB);
-  void processBlock(MachineBasicBlock &MBB, bool IsEntry);
+  void lowerBlock(MachineBasicBlock &MBB, BlockInfo &BI);
+  void processBlock(MachineBasicBlock &MBB, BlockInfo &BI, bool IsEntry);
 
   bool lowerLiveMaskQueries();
   bool lowerCopyInstrs();
@@ -1035,12 +1035,7 @@ MachineInstr *SIWholeQuadMode::lowerKillI1(MachineBasicBlock &MBB,
 // Replace (or supplement) instructions accessing live mask.
 // This can only happen once all the live mask registers have been created
 // and the execute state (WQM/StrictWWM/Exact) of instructions is known.
-void SIWholeQuadMode::lowerBlock(MachineBasicBlock &MBB) {
-  auto *BII = Blocks.find(&MBB);
-  if (BII == Blocks.end())
-    return;
-
-  const BlockInfo &BI = BII->second;
+void SIWholeQuadMode::lowerBlock(MachineBasicBlock &MBB, BlockInfo &BI) {
   if (!BI.NeedsLowering)
     return;
 
@@ -1052,8 +1047,9 @@ void SIWholeQuadMode::lowerBlock(MachineBasicBlock &MBB) {
 
   for (MachineInstr &MI : llvm::make_early_inc_range(
            llvm::make_range(MBB.getFirstNonPHI(), MBB.end()))) {
-    if (StateTransition.count(&MI))
-      State = StateTransition[&MI];
+    auto MIState = StateTransition.find(&MI);
+    if (MIState != StateTransition.end())
+      State = MIState->second;
 
     MachineInstr *SplitPoint = nullptr;
     switch (MI.getOpcode()) {
@@ -1260,13 +1256,8 @@ void SIWholeQuadMode::fromStrictMode(MachineBasicBlock &MBB,
   StateTransition[MI] = NonStrictState;
 }
 
-void SIWholeQuadMode::processBlock(MachineBasicBlock &MBB, bool IsEntry) {
-  auto *BII = Blocks.find(&MBB);
-  if (BII == Blocks.end())
-    return;
-
-  BlockInfo &BI = BII->second;
-
+void SIWholeQuadMode::processBlock(MachineBasicBlock &MBB, BlockInfo &BI,
+                                   bool IsEntry) {
   // This is a non-entry block that is WQM throughout, so no need to do
   // anything.
   if (!IsEntry && BI.Needs == StateWQM && BI.OutNeeds != StateExact) {
@@ -1809,11 +1800,11 @@ bool SIWholeQuadMode::runOnMachineFunction(MachineFunction &MF) {
     if (GlobalFlags & StateWQM)
       Blocks[&Entry].InNeeds |= StateWQM;
     // Wave mode switching requires full lowering pass.
-    for (auto BII : Blocks)
-      processBlock(*BII.first, BII.first == &Entry);
+    for (auto &BII : Blocks)
+      processBlock(*BII.first, BII.second, BII.first == &Entry);
     // Lowering blocks causes block splitting so perform as a second pass.
-    for (auto BII : Blocks)
-      lowerBlock(*BII.first);
+    for (auto &BII : Blocks)
+      lowerBlock(*BII.first, BII.second);
     Changed = true;
   }