[AArch64] Fix a compiler crash in MachineSink

[momchil-velikov]

There were a couple of issues with maintaining register def/uses in MachineRegisterInfo:

• when an operand is changed from one register to another, the corresponding instructions must already be inserted into the function, or MRI won't be updated

• when traversing the set of all uses of a register, that set must not change

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Changes

There were a couple of issues with maintaining register def/uses in MachineRegisterInfo:

• when an operand is changed from one register to another, the corresponding instructions must already be inserted into the function, or MRI won't be updated

• when traversing the set of all uses of a register, that set must not change

---

Full diff: https://github.com/llvm/llvm-project/pull/67705.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/MachineSink.cpp (+19-17)
• (added) llvm/test/CodeGen/AArch64/sink-and-fold-dbg-value-crash.mir (+172)

diff --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index 480ac23d43ad879..9d4e0c647048f53 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -461,7 +461,7 @@ bool MachineSinking::PerformSinkAndFold(MachineInstr &MI,
       }
 
       if (UseInst.getParent() != MI.getParent()) {
-        // If the register class of the register we are replacingis a superset
+        // If the register class of the register we are replacing is a superset
         // of any of the register classes of the operands of the materialized
         // instruction don't consider that live range extended.
         const TargetRegisterClass *RCS = MRI->getRegClass(Reg);
@@ -527,8 +527,8 @@ bool MachineSinking::PerformSinkAndFold(MachineInstr &MI,
         if (U.getInt())
           continue;
         MachineInstr *NewDbgMI = SinkDst->getMF()->CloneMachineInstr(DbgMI);
-        NewDbgMI->getOperand(0).setReg(DstReg);
         SinkMBB.insertAfter(InsertPt, NewDbgMI);
+        NewDbgMI->getOperand(0).setReg(DstReg);
       }
     } else {
       // Fold instruction into the addressing mode of a memory instruction.
@@ -538,33 +538,35 @@ bool MachineSinking::PerformSinkAndFold(MachineInstr &MI,
     SinkDst->eraseFromParent();
   }
 
-  MI.eraseFromParent();
-
-  // Collect instructions that need to be deleted (COPYs). We cannot delete them
-  // while traversing register uses.
-  SmallVector<MachineInstr *> CleanupInstrs;
+  // Collect operands that need to be cleaned up because the registers no longer
+  // exist (in COPYs and debug instructions). We cannot delete instructions or
+  // clear operands while traversing register uses.
+  SmallVector<MachineOperand *> Cleanup;
   Worklist.push_back(DefReg);
   while (!Worklist.empty()) {
     Register Reg = Worklist.pop_back_val();
-
     for (MachineOperand &MO : MRI->use_operands(Reg)) {
       MachineInstr *U = MO.getParent();
       assert((U->isCopy() || U->isDebugInstr()) &&
              "Only debug uses and copies must remain");
-      if (U->isCopy()) {
+      if (U->isCopy())
         Worklist.push_back(U->getOperand(0).getReg());
-        CleanupInstrs.push_back(U);
-      } else {
-        MO.setReg(0);
-        MO.setSubReg(0);
-      }
+      Cleanup.push_back(&MO);
     }
   }
 
-  // Delete the dead COPYs.
-  for (MachineInstr *Del : CleanupInstrs)
-    Del->eraseFromParent();
+  // Delete the dead COPYs and clear operands in debug instructions
+  for (MachineOperand *MO : Cleanup) {
+    MachineInstr *I = MO->getParent();
+    if (I->isCopy()) {
+      I->eraseFromParent();
+    } else {
+      MO->setReg(0);
+      MO->setSubReg(0);
+    }
+  }
 
+  MI.eraseFromParent();
   return true;
 }
 
diff --git a/llvm/test/CodeGen/AArch64/sink-and-fold-dbg-value-crash.mir b/llvm/test/CodeGen/AArch64/sink-and-fold-dbg-value-crash.mir
new file mode 100644
index 000000000000000..0fc645022780d61
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sink-and-fold-dbg-value-crash.mir
@@ -0,0 +1,172 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 3
+# RUN: llc --aarch64-enable-sink-fold=true --run-pass=machine-sink %s -o - | FileCheck %s
+--- |
+  source_filename = "x.ll"
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-linux"
+
+  %S = type <{ ptr, %T, i64, i64, i32, [4 x i8] }>
+  %T = type { ptr, ptr, ptr, ptr, i64 }
+
+  define void @f(ptr %p, i1 %c0, i1 %c1) {
+  entry:
+    %a = getelementptr %S, ptr %p, i64 0, i32 1
+    call void @llvm.dbg.value(metadata ptr %a, metadata !4, metadata !DIExpression()), !dbg !10
+    br i1 %c0, label %if.then, label %if.end
+
+  if.then:                                          ; preds = %entry
+    %v0 = tail call ptr @g(ptr %a)
+    br i1 %c1, label %exit, label %if.end
+
+  if.end:                                           ; preds = %if.then, %entry
+    %v1 = load i64, ptr %a, align 8
+    br label %exit
+
+  exit:                                             ; preds = %if.end, %if.then
+    ret void
+  }
+
+  declare ptr @g(ptr)
+
+  declare void @llvm.dbg.value(metadata, metadata, metadata) #0
+
+  attributes #0 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+
+  !llvm.dbg.cu = !{!0}
+  !llvm.module.flags = !{!2, !3}
+
+  !0 = distinct !DICompileUnit(language: DW_LANG_C11, file: !1, producer: "clang ", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, splitDebugInlining: false, nameTableKind: None)
+  !1 = !DIFile(filename: "f.c", directory: "/usr/rms")
+  !2 = !{i32 7, !"Dwarf Version", i32 4}
+  !3 = !{i32 2, !"Debug Info Version", i32 3}
+  !4 = !DILocalVariable(name: "x", arg: 1, scope: !5, file: !1, line: 2, type: !8)
+  !5 = distinct !DISubprogram(name: "f", scope: !1, file: !1, line: 2, type: !6, scopeLine: 2, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition, unit: !0, retainedNodes: !9)
+  !6 = !DISubroutineType(types: !7)
+  !7 = !{!8, !8}
+  !8 = !DIBasicType(name: "int", size: 32, encoding: DW_ATE_signed)
+  !9 = !{}
+  !10 = !DILocation(line: 2, column: 11, scope: !5)
+
+...
+---
+name:            f
+alignment:       4
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+failedISel:      false
+tracksRegLiveness: true
+hasWinCFI:       false
+callsEHReturn:   false
+callsUnwindInit: false
+hasEHCatchret:   false
+hasEHScopes:     false
+hasEHFunclets:   false
+isOutlined:      false
+debugInstrRef:   false
+failsVerification: false
+tracksDebugUserValues: false
+registers:
+  - { id: 0, class: gpr64all, preferred-register: '' }
+  - { id: 1, class: gpr64common, preferred-register: '' }
+  - { id: 2, class: gpr32, preferred-register: '' }
+  - { id: 3, class: gpr32, preferred-register: '' }
+  - { id: 4, class: gpr32, preferred-register: '' }
+  - { id: 5, class: gpr64sp, preferred-register: '' }
+  - { id: 6, class: gpr64all, preferred-register: '' }
+liveins:
+  - { reg: '$x0', virtual-reg: '%1' }
+  - { reg: '$w1', virtual-reg: '%2' }
+  - { reg: '$w2', virtual-reg: '%3' }
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   false
+  stackSize:       0
+  offsetAdjustment: 0
+  maxAlignment:    1
+  adjustsStack:    true
+  hasCalls:        true
+  stackProtector:  ''
+  functionContext: ''
+  maxCallFrameSize: 0
+  cvBytesOfCalleeSavedRegisters: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+  hasTailCall:     false
+  localFrameSize:  0
+  savePoint:       ''
+  restorePoint:    ''
+fixedStack:      []
+stack:           []
+entry_values:    []
+callSites:       []
+debugValueSubstitutions: []
+constants:       []
+machineFunctionInfo: {}
+body:             |
+  ; CHECK-LABEL: name: f
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT:   liveins: $x0, $w1, $w2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gpr32 = COPY $w2
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:gpr32 = COPY $w1
+  ; CHECK-NEXT:   [[COPY2:%[0-9]+]]:gpr64common = COPY $x0
+  ; CHECK-NEXT:   DBG_VALUE $noreg, $noreg, !4, !DIExpression(), debug-location !10
+  ; CHECK-NEXT:   TBZW [[COPY1]], 0, %bb.2
+  ; CHECK-NEXT:   B %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1.if.then:
+  ; CHECK-NEXT:   successors: %bb.3(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY3:%[0-9]+]]:gpr32 = COPY [[COPY]]
+  ; CHECK-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
+  ; CHECK-NEXT:   $x0 = ADDXri [[COPY2]], 8, 0
+  ; CHECK-NEXT:   DBG_VALUE $x0, $noreg, !4, !DIExpression(), debug-location !10
+  ; CHECK-NEXT:   BL @g, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit $x0, implicit-def $sp, implicit-def $x0
+  ; CHECK-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
+  ; CHECK-NEXT:   TBNZW [[COPY3]], 0, %bb.3
+  ; CHECK-NEXT:   B %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2.if.end:
+  ; CHECK-NEXT:   successors: %bb.3(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3.exit:
+  ; CHECK-NEXT:   RET_ReallyLR
+  bb.0.entry:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $x0, $w1, $w2
+
+    %3:gpr32 = COPY $w2
+    %2:gpr32 = COPY $w1
+    %1:gpr64common = COPY $x0
+    %4:gpr32 = COPY %3
+    %5:gpr64sp = ADDXri %1, 8, 0
+    DBG_VALUE %5, $noreg, !4, !DIExpression(), debug-location !10
+    %0:gpr64all = COPY %5
+    TBZW %2, 0, %bb.2
+    B %bb.1
+
+  bb.1.if.then:
+    successors: %bb.3(0x40000000), %bb.2(0x40000000)
+
+    ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
+    $x0 = COPY %0
+    BL @g, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit $x0, implicit-def $sp, implicit-def $x0
+    ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
+    TBNZW %4, 0, %bb.3
+    B %bb.2
+
+  bb.2.if.end:
+    successors: %bb.3(0x80000000)
+
+
+  bb.3.exit:
+    RET_ReallyLR
+
+...

[davemgreen]

Looks good to me.