[MIR] Fix out of bounds access in MIRPrinter.

Fixes: SWDEV-256460

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D90239

llvm/lib/CodeGen/MIRPrinter.cpp

@@ -367,9 +367,17 @@ void MIRPrinter::convertStackObjects(yaml::MachineFunction &YMF,
                                      ModuleSlotTracker &MST) {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
+
   // Process fixed stack objects.
+  assert(YMF.FixedStackObjects.empty());
+  SmallVector<int, 32> FixedStackObjectsIdx;
+  const int BeginIdx = MFI.getObjectIndexBegin();
+  if (BeginIdx < 0)
+    FixedStackObjectsIdx.reserve(-BeginIdx);
+
   unsigned ID = 0;
-  for (int I = MFI.getObjectIndexBegin(); I < 0; ++I, ++ID) {
+  for (int I = BeginIdx; I < 0; ++I, ++ID) {
+    FixedStackObjectsIdx.push_back(-1); // Fill index for possible dead.
     if (MFI.isDeadObjectIndex(I))
       continue;
 
@@ -384,14 +392,22 @@ void MIRPrinter::convertStackObjects(yaml::MachineFunction &YMF,
     YamlObject.StackID = (TargetStackID::Value)MFI.getStackID(I);
     YamlObject.IsImmutable = MFI.isImmutableObjectIndex(I);
     YamlObject.IsAliased = MFI.isAliasedObjectIndex(I);
+    // Save the ID' position in FixedStackObjects storage vector.
+    FixedStackObjectsIdx[ID] = YMF.FixedStackObjects.size();
     YMF.FixedStackObjects.push_back(YamlObject);
     StackObjectOperandMapping.insert(
         std::make_pair(I, FrameIndexOperand::createFixed(ID)));
   }
 
   // Process ordinary stack objects.
+  assert(YMF.StackObjects.empty());
+  SmallVector<unsigned, 32> StackObjectsIdx;
+  const int EndIdx = MFI.getObjectIndexEnd();
+  if (EndIdx > 0)
+    StackObjectsIdx.reserve(EndIdx);
   ID = 0;
-  for (int I = 0, E = MFI.getObjectIndexEnd(); I < E; ++I, ++ID) {
+  for (int I = 0; I < EndIdx; ++I, ++ID) {
+    StackObjectsIdx.push_back(-1); // Fill index for possible dead.
     if (MFI.isDeadObjectIndex(I))
       continue;
 
@@ -410,41 +426,42 @@ void MIRPrinter::convertStackObjects(yaml::MachineFunction &YMF,
     YamlObject.Alignment = MFI.getObjectAlign(I);
     YamlObject.StackID = (TargetStackID::Value)MFI.getStackID(I);
 
+    // Save the ID' position in StackObjects storage vector.
+    StackObjectsIdx[ID] = YMF.StackObjects.size();
     YMF.StackObjects.push_back(YamlObject);
     StackObjectOperandMapping.insert(std::make_pair(
         I, FrameIndexOperand::create(YamlObject.Name.Value, ID)));
   }
 
   for (const auto &CSInfo : MFI.getCalleeSavedInfo()) {
-    if (!CSInfo.isSpilledToReg() && MFI.isDeadObjectIndex(CSInfo.getFrameIdx()))
+    const int FrameIdx = CSInfo.getFrameIdx();
+    if (!CSInfo.isSpilledToReg() && MFI.isDeadObjectIndex(FrameIdx))
       continue;
 
     yaml::StringValue Reg;
     printRegMIR(CSInfo.getReg(), Reg, TRI);
     if (!CSInfo.isSpilledToReg()) {
-      auto StackObjectInfo = StackObjectOperandMapping.find(CSInfo.getFrameIdx());
+      assert(FrameIdx >= MFI.getObjectIndexBegin() &&
-      assert(StackObjectInfo != StackObjectOperandMapping.end() &&
+             FrameIdx < MFI.getObjectIndexEnd() &&
              "Invalid stack object index");
-      const FrameIndexOperand &StackObject = StackObjectInfo->second;
+      if (FrameIdx < 0) { // Negative index means fixed objects.
-      if (StackObject.IsFixed) {
+        auto &Object =
-        YMF.FixedStackObjects[StackObject.ID].CalleeSavedRegister = Reg;
+            YMF.FixedStackObjects
-        YMF.FixedStackObjects[StackObject.ID].CalleeSavedRestored =
+                [FixedStackObjectsIdx[FrameIdx + MFI.getNumFixedObjects()]];
-          CSInfo.isRestored();
+        Object.CalleeSavedRegister = Reg;
+        Object.CalleeSavedRestored = CSInfo.isRestored();
       } else {
-        YMF.StackObjects[StackObject.ID].CalleeSavedRegister = Reg;
+        auto &Object = YMF.StackObjects[StackObjectsIdx[FrameIdx]];
-        YMF.StackObjects[StackObject.ID].CalleeSavedRestored =
+        Object.CalleeSavedRegister = Reg;
-          CSInfo.isRestored();
+        Object.CalleeSavedRestored = CSInfo.isRestored();
       }
     }
   }
   for (unsigned I = 0, E = MFI.getLocalFrameObjectCount(); I < E; ++I) {
     auto LocalObject = MFI.getLocalFrameObjectMap(I);
-    auto StackObjectInfo = StackObjectOperandMapping.find(LocalObject.first);
+    assert(LocalObject.first >= 0 && "Expected a locally mapped stack object");
-    assert(StackObjectInfo != StackObjectOperandMapping.end() &&
+    YMF.StackObjects[StackObjectsIdx[LocalObject.first]].LocalOffset =
-           "Invalid stack object index");
+        LocalObject.second;
-    const FrameIndexOperand &StackObject = StackObjectInfo->second;
-    assert(!StackObject.IsFixed && "Expected a locally mapped stack object");
-    YMF.StackObjects[StackObject.ID].LocalOffset = LocalObject.second;
   }
 
   // Print the stack object references in the frame information class after
@@ -458,15 +475,16 @@ void MIRPrinter::convertStackObjects(yaml::MachineFunction &YMF,
   // Print the debug variable information.
   for (const MachineFunction::VariableDbgInfo &DebugVar :
        MF.getVariableDbgInfo()) {
-    auto StackObjectInfo = StackObjectOperandMapping.find(DebugVar.Slot);
+    assert(DebugVar.Slot >= MFI.getObjectIndexBegin() &&
-    assert(StackObjectInfo != StackObjectOperandMapping.end() &&
+           DebugVar.Slot < MFI.getObjectIndexEnd() &&
            "Invalid stack object index");
-    const FrameIndexOperand &StackObject = StackObjectInfo->second;
+    if (DebugVar.Slot < 0) { // Negative index means fixed objects.
-    if (StackObject.IsFixed) {
+      auto &Object =
-      auto &Object = YMF.FixedStackObjects[StackObject.ID];
+          YMF.FixedStackObjects[FixedStackObjectsIdx[DebugVar.Slot +
+                                                     MFI.getNumFixedObjects()]];
       printStackObjectDbgInfo(DebugVar, Object, MST);
     } else {
-      auto &Object = YMF.StackObjects[StackObject.ID];
+      auto &Object = YMF.StackObjects[StackObjectsIdx[DebugVar.Slot]];
       printStackObjectDbgInfo(DebugVar, Object, MST);
     }
   }

llvm/test/CodeGen/MIR/AMDGPU/stack-id-assert.mir

@@ -0,0 +1,63 @@
+# This test used to crash MIRPrinter::convertStackObjects():
+# MFI can contain some dead stack objects after PEI pass, but objects storage
+# contains not dead objects only. So using objects IDs as offset in the storage
+# caused out of bounds access.
+
+# RUN: llc -march=amdgcn -run-pass=si-lower-sgpr-spills,prologepilog -verify-machineinstrs -o - %s | FileCheck %s
+
+# CHECK-LABEL: name:            foo
+# CHECK: {{^}}fixedStack:      []
+# CHECK: stack:           []
+
+# CHECK-LABEL: name:            bar
+# CHECK: fixedStack:      []
+# CHECK-NEXT: {{^}}stack:
+# CHECK-NEXT:  - { id:
+
+
+---
+name:            foo
+liveins:
+  - { reg: '$sgpr30_sgpr31', virtual-reg: '' }
+body:             |
+  bb.0:
+    S_SETPC_B64_return killed renamable $sgpr30_sgpr31
+
+...
+---
+name:            bar
+tracksRegLiveness: true
+liveins:
+  - { reg: '$vgpr0', virtual-reg: '' }
+  - { reg: '$vgpr1', virtual-reg: '' }
+  - { reg: '$vgpr2', virtual-reg: '' }
+  - { reg: '$sgpr30_sgpr31', virtual-reg: '' }
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: 0, size: 8, alignment: 4,
+      stack-id: sgpr-spill, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+machineFunctionInfo:
+  scratchRSrcReg:  '$sgpr0_sgpr1_sgpr2_sgpr3'
+  stackPtrOffsetReg: '$sgpr32'
+body:             |
+  bb.0:
+    liveins: $vgpr0, $vgpr1, $vgpr2, $sgpr30_sgpr31
+
+    renamable $vgpr41 = COPY $vgpr2, implicit $exec
+    renamable $vgpr40 = COPY $vgpr1, implicit $exec
+    renamable $vcc = V_CMP_NE_U32_e64 42, killed $vgpr0, implicit $exec
+    $sgpr4_sgpr5 = S_AND_SAVEEXEC_B64 $vcc, implicit-def $exec, implicit-def $scc, implicit $exec
+    renamable $sgpr34_sgpr35 = S_XOR_B64 $exec, killed renamable $sgpr4_sgpr5, implicit-def dead $scc
+
+    ADJCALLSTACKUP 0, 0, implicit-def dead $scc, implicit-def $sgpr32, implicit $sgpr32
+    renamable $sgpr4_sgpr5 = SI_PC_ADD_REL_OFFSET target-flags(amdgpu-gotprel32-lo) @foo + 4, target-flags(amdgpu-gotprel32-hi) @foo + 12, implicit-def dead $scc
+    renamable $sgpr4_sgpr5 = S_LOAD_DWORDX2_IMM killed renamable $sgpr4_sgpr5, 0, 0, 0
+    SI_SPILL_S64_SAVE killed renamable $sgpr30_sgpr31, %stack.0, implicit $exec, implicit $sgpr0_sgpr1_sgpr2_sgpr3, implicit $sgpr32
+    dead $sgpr30_sgpr31 = SI_CALL killed renamable $sgpr4_sgpr5, @foo, csr_amdgpu_highregs, implicit $sgpr0_sgpr1_sgpr2_sgpr3
+    renamable $sgpr30_sgpr31 = SI_SPILL_S64_RESTORE %stack.0, implicit $exec, implicit $sgpr0_sgpr1_sgpr2_sgpr3, implicit $sgpr32
+    ADJCALLSTACKDOWN 0, 0, implicit-def dead $scc, implicit-def $sgpr32, implicit $sgpr32
+
+    S_SETPC_B64_return killed renamable $sgpr30_sgpr31
+
+...
+