[MIR] Fix out of bounds access in MIRPrinter.

[dfukalov]

Summary

Fixes: SWDEV-256460

[arsenm]

llvm/lib/CodeGen/MIRPrinter.cpp

91–92

Why does the index need to be stored? It seems like the problem is the printer/parser implicitly deleting dead objects

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

15–21

Don't need the IR section

117–181

Most of the mir fields can be eliminated

221

You should only need a single block with the spill instructions

[dfukalov]

Addressed comments.

[dfukalov]

llvm/lib/CodeGen/MIRPrinter.cpp

91–92

The issue was that IDs stored in FrameIndexOperand are not the same numbers as offsets in YMF.StackObjects and YMF.FixedStackObjects vectors, since they doesn't contain dead objects.

So I added mapping ID -> storage offset. They are stored in temporary FixedStackObjectsIdx and StackObjectsIdx within MIRPrinter::convertStackObjects.

[arsenm]

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

3

Should mention why with dead objects

[dfukalov]

Addressed comments.

llvm/lib/CodeGen/MIRPrinter.cpp

static cl::opt<bool> PrintLocations("mir-debug-loc", cl::Hidden, cl::init(true),
                                     cl::desc("Print MIR debug-locations"));
 
 namespace {
 
 /// This structure describes how to print out stack object references.
 struct FrameIndexOperand {
   std::string Name;
   unsigned ID;
   bool IsFixed;
 
   FrameIndexOperand(StringRef Name, unsigned ID, bool IsFixed)
       : Name(Name.str()), ID(ID), IsFixed(IsFixed) {}
 
   /// Return an ordinary stack object reference.
   static FrameIndexOperand create(StringRef Name, unsigned ID) {
     return FrameIndexOperand(Name, ID, /*IsFixed=*/false);
   }
 
void MIRPrinter::convert(ModuleSlotTracker &MST,
   }
 }
 
 void MIRPrinter::convertStackObjects(yaml::MachineFunction &YMF,
                                      const MachineFunction &MF,
                                      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;
 
     yaml::FixedMachineStackObject YamlObject;
     YamlObject.ID = ID;
     YamlObject.Type = MFI.isSpillSlotObjectIndex(I)
                           ? yaml::FixedMachineStackObject::SpillSlot
                           : yaml::FixedMachineStackObject::DefaultType;
     YamlObject.Offset = MFI.getObjectOffset(I);
     YamlObject.Size = MFI.getObjectSize(I);
     YamlObject.Alignment = MFI.getObjectAlign(I);
     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;
 
     yaml::MachineStackObject YamlObject;
     YamlObject.ID = ID;
     if (const auto *Alloca = MFI.getObjectAllocation(I))
       YamlObject.Name.Value = std::string(
           Alloca->hasName() ? Alloca->getName() : "<unnamed alloca>");
     YamlObject.Type = MFI.isSpillSlotObjectIndex(I)
                           ? yaml::MachineStackObject::SpillSlot
                           : MFI.isVariableSizedObjectIndex(I)
                                 ? yaml::MachineStackObject::VariableSized
                                 : yaml::MachineStackObject::DefaultType;
     YamlObject.Offset = MFI.getObjectOffset(I);
     YamlObject.Size = MFI.getObjectSize(I);
     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(StackObjectInfo != StackObjectOperandMapping.end() &&
+      assert(FrameIdx >= MFI.getObjectIndexBegin() &&
+             FrameIdx < MFI.getObjectIndexEnd() &&
              "Invalid stack object index");
-      const FrameIndexOperand &StackObject = StackObjectInfo->second;
-      if (StackObject.IsFixed) {
-        YMF.FixedStackObjects[StackObject.ID].CalleeSavedRegister = Reg;
-        YMF.FixedStackObjects[StackObject.ID].CalleeSavedRestored =
-          CSInfo.isRestored();
+      if (FrameIdx < 0) { // Negative index means fixed objects.
+        auto &Object =
+            YMF.FixedStackObjects
+                [FixedStackObjectsIdx[FrameIdx + MFI.getNumFixedObjects()]];
+        Object.CalleeSavedRegister = Reg;
+        Object.CalleeSavedRestored = CSInfo.isRestored();
       } else {
-        YMF.StackObjects[StackObject.ID].CalleeSavedRegister = Reg;
-        YMF.StackObjects[StackObject.ID].CalleeSavedRestored =
-          CSInfo.isRestored();
+        auto &Object = YMF.StackObjects[StackObjectsIdx[FrameIdx]];
+        Object.CalleeSavedRegister = Reg;
+        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(StackObjectInfo != StackObjectOperandMapping.end() &&
-           "Invalid stack object index");
-    const FrameIndexOperand &StackObject = StackObjectInfo->second;
-    assert(!StackObject.IsFixed && "Expected a locally mapped stack object");
-    YMF.StackObjects[StackObject.ID].LocalOffset = LocalObject.second;
+    assert(LocalObject.first >= 0 && "Expected a locally mapped stack object");
+    YMF.StackObjects[StackObjectsIdx[LocalObject.first]].LocalOffset =
+        LocalObject.second;
   }
 
   // Print the stack object references in the frame information class after
   // converting the stack objects.
   if (MFI.hasStackProtectorIndex()) {
     raw_string_ostream StrOS(YMF.FrameInfo.StackProtector.Value);
     MIPrinter(StrOS, MST, RegisterMaskIds, StackObjectOperandMapping)
         .printStackObjectReference(MFI.getStackProtectorIndex());
   }
 
   // Print the debug variable information.
   for (const MachineFunction::VariableDbgInfo &DebugVar :
        MF.getVariableDbgInfo()) {
-    auto StackObjectInfo = StackObjectOperandMapping.find(DebugVar.Slot);
-    assert(StackObjectInfo != StackObjectOperandMapping.end() &&
+    assert(DebugVar.Slot >= MFI.getObjectIndexBegin() &&
+           DebugVar.Slot < MFI.getObjectIndexEnd() &&
            "Invalid stack object index");
-    const FrameIndexOperand &StackObject = StackObjectInfo->second;
-    if (StackObject.IsFixed) {
-      auto &Object = YMF.FixedStackObjects[StackObject.ID];
+    if (DebugVar.Slot < 0) { // Negative index means fixed objects.
+      auto &Object =
+          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);
     }
   }
 }
 
 void MIRPrinter::convertCallSiteObjects(yaml::MachineFunction &YMF,
                                         const MachineFunction &MF,
                                         ModuleSlotTracker &MST) {

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

+# 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
+
+...
+