[InstrRef] Preserve debug instr num in aarch64-ldst-opt. (#136009)

The aarch64-ldst-opt pass tries to merge two load instructions
(LDR*) to a load pair instruction (LDP*).

When merging the instructions, there is a case where one of the 
loads would have to also be sign extended. In either case, 
(sign extend or not), the pass needs to preserve the debug-instr-number
from the original loads to the load pair instruction to make sure debug
info
isn't lost in the case where instruction referencing is being used.

For example:

We can have something like this:

```
debugValueSubstitutions:[]
$x1 = LDRXui $x0, 1, debug-instr-number 1
DBG_INSTR_REF !13, dbg-instr-ref(1, 0), debug-location !11
$x0 = LDRXui killed $x0, 0, debug-instr-number 2
DBG_INSTR_REF !14, dbg-instr-ref(2, 0), debug-location !11
```

This would be changed to:

```
debugValueSubstitutions: []
$x0, $x1 = LDPXi $x0, 0
DBG_INSTR_REF !12, dbg-instr-ref(1, 0), debug-location !14
DBG_INSTR_REF !13, dbg-instr-ref(2, 0), debug-location !14
```

In this case, we need to create a new debug instruction number
for the `LDP` instruction, we then need to add entries into the 
debugSubstitutions table to map the old instr-refs to the new ones.

After this patch, the result will be:
```
 debugValueSubstitutions:
 - { srcinst: 1, srcop: 0, dstinst: 3, dstop: 1, subreg: 0 }
 - { srcinst: 2, srcop: 0, dstinst: 3, dstop: 0, subreg: 0 }
 $x0, $x1 = LDPXi $x0, 0, debug-instr-number 3
 DBG_INSTR_REF !12, dbg-instr-ref(1, 0), debug-location !14
 DBG_INSTR_REF !12, dbg-instr-ref(2, 0), debug-location !14
```

However, this is not all, we also can have a case where there is a
sign-extend involved, let's look at the case:

```
debugValueSubstitutions:[]
$w1 = LDRWui $x0, 1, debug-instr-number 1
DBG_INSTR_REF !7, dbg-instr-ref(1, 0), debug-location !9
$x0 = LDRSWui $x0, 0, debug-instr-number 2
DBG_INSTR_REF !8, dbg-instr-ref(2, 0), debug-location !9
```

This will become:

```
debugValueSubstitutions:[]
$w0, $w1 = LDPWi $x0, 0
$w0 = KILL $w0, implicit-def $x0
$x0 = SBFMXri $x0, 0, 31
DBG_INSTR_REF !7, dbg-instr-ref(1, 0), debug-location !9
DBG_INSTR_REF !8, dbg-instr-ref(2, 0), debug-location !9
```
$x0 is where the final value is stored, so the sign extend (SBFMXri)
instruction contains the final value we care about we give it a new
debug-instr-number 3. Whereas, $w1 contains the final value that we care
about, therefore the LDP instruction is also given a new
debug-instr-number 4. We have to add these subsitutions to the
debugValueSubstitutions table. However, we also have to ensure that the
OpIndex that pointed to debug-instr-number 1 gets updated to 1, because
$w1 is the second operand of the LDP instruction.

The result after the patch looks like:

```
debugValueSubstitutions:
- { srcinst: 1, srcop: 0, dstinst: 4, dstop: 1, subreg: 0 }
- { srcinst: 2, srcop: 0, dstinst: 3, dstop: 0, subreg: 0 }
$w0, $w1 = LDPWi $x0, 0, debug-instr-number 4
$w0 = KILL $w0, implicit-def $x0
$x0 = SBFMXri $x0, 0, 31, debug-instr-number 3
DBG_INSTR_REF !7, dbg-instr-ref(1, 0), debug-location !9
DBG_INSTR_REF !8, dbg-instr-ref(2, 0), debug-location !9
```


This patch addresses that problem.

2 files changed, 229 insertions, 0 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
index 0e26005..ba3ffc2 100644
--- a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
@@ -964,6 +964,32 @@ static void updateDefinedRegisters(MachineInstr &MI, LiveRegUnits &Units,
       Units.addReg(MOP.getReg());
 }
 
+/// This function will add a new entry into the debugValueSubstitutions table
+/// when two instruction have been merged into a new one represented by \p
+/// MergedInstr.
+static void addDebugSubstitutionsToTable(MachineFunction *MF,
+                                         unsigned InstrNumToSet,
+                                         MachineInstr &OriginalInstr,
+                                         MachineInstr &MergedInstr) {
+
+  // Figure out the Operand Index of the destination register of the
+  // OriginalInstr in the new MergedInstr.
+  auto Reg = OriginalInstr.getOperand(0).getReg();
+  unsigned OperandNo = 0;
+  bool RegFound = false;
+  for (const auto Op : MergedInstr.operands()) {
+    if (Op.getReg() == Reg) {
+      RegFound = true;
+      break;
+    }
+    OperandNo++;
+  }
+
+  if (RegFound)
+    MF->makeDebugValueSubstitution({OriginalInstr.peekDebugInstrNum(), 0},
+                                   {InstrNumToSet, OperandNo});
+}
+
 MachineBasicBlock::iterator
 AArch64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
                                       MachineBasicBlock::iterator Paired,
@@ -1226,6 +1252,79 @@ AArch64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
             .addImm(0)
             .addImm(31);
     (void)MIBSXTW;
+
+    // In the case of a sign-extend, where we have something like:
+    // debugValueSubstitutions:[]
+    // $w1 = LDRWui $x0, 1, debug-instr-number 1
+    // DBG_INSTR_REF !7, dbg-instr-ref(1, 0), debug-location !9
+    // $x0 = LDRSWui $x0, 0, debug-instr-number 2
+    // DBG_INSTR_REF !8, dbg-instr-ref(2, 0), debug-location !9
+
+    // It will be converted to:
+    // debugValueSubstitutions:[]
+    // $w0, $w1 = LDPWi $x0, 0
+    // $w0 = KILL $w0, implicit-def $x0
+    // $x0 = SBFMXri $x0, 0, 31
+    // DBG_INSTR_REF !7, dbg-instr-ref(1, 0), debug-location !9
+    // DBG_INSTR_REF !8, dbg-instr-ref(2, 0), debug-location !9
+
+    // We want the final result to look like:
+    // debugValueSubstitutions:
+    // - { srcinst: 1, srcop: 0, dstinst: 4, dstop: 1, subreg: 0 }
+    // - { srcinst: 2, srcop: 0, dstinst: 3, dstop: 0, subreg: 0 }
+    // $w0, $w1 = LDPWi $x0, 0, debug-instr-number 4
+    // $w0 = KILL $w0, implicit-def $x0
+    // $x0 = SBFMXri $x0, 0, 31, debug-instr-number 3
+    // DBG_INSTR_REF !7, dbg-instr-ref(1, 0), debug-location !9
+    // DBG_INSTR_REF !8, dbg-instr-ref(2, 0), debug-location !9
+
+    // $x0 is where the final value is stored, so the sign extend (SBFMXri)
+    // instruction contains the final value we care about we give it a new
+    // debug-instr-number 3. Whereas, $w1 contains the final value that we care
+    // about, therefore the LDP instruction is also given a new
+    // debug-instr-number 4. We have to add these subsitutions to the
+    // debugValueSubstitutions table. However, we also have to ensure that the
+    // OpIndex that pointed to debug-instr-number 1 gets updated to 1, because
+    // $w1 is the second operand of the LDP instruction.
+
+    if (I->peekDebugInstrNum()) {
+      // If I is the instruction which got sign extended and has a
+      // debug-instr-number, give the SBFMXri instruction a new
+      // debug-instr-number, and update the debugValueSubstitutions table with
+      // the new debug-instr-number and OpIndex pair. Otherwise, give the Merged
+      // instruction a new debug-instr-number, and update the
+      // debugValueSubstitutions table with the new debug-instr-number and
+      // OpIndex pair.
+      unsigned NewInstrNum;
+      if (DstRegX == I->getOperand(0).getReg()) {
+        NewInstrNum = MIBSXTW->getDebugInstrNum();
+        addDebugSubstitutionsToTable(MBB->getParent(), NewInstrNum, *I,
+                                     *MIBSXTW);
+      } else {
+        NewInstrNum = MIB->getDebugInstrNum();
+        addDebugSubstitutionsToTable(MBB->getParent(), NewInstrNum, *I, *MIB);
+      }
+    }
+    if (Paired->peekDebugInstrNum()) {
+      // If Paired is the instruction which got sign extended and has a
+      // debug-instr-number, give the SBFMXri instruction a new
+      // debug-instr-number, and update the debugValueSubstitutions table with
+      // the new debug-instr-number and OpIndex pair. Otherwise, give the Merged
+      // instruction a new debug-instr-number, and update the
+      // debugValueSubstitutions table with the new debug-instr-number and
+      // OpIndex pair.
+      unsigned NewInstrNum;
+      if (DstRegX == Paired->getOperand(0).getReg()) {
+        NewInstrNum = MIBSXTW->getDebugInstrNum();
+        addDebugSubstitutionsToTable(MBB->getParent(), NewInstrNum, *Paired,
+                                     *MIBSXTW);
+      } else {
+        NewInstrNum = MIB->getDebugInstrNum();
+        addDebugSubstitutionsToTable(MBB->getParent(), NewInstrNum, *Paired,
+                                     *MIB);
+      }
+    }
+
     LLVM_DEBUG(dbgs() << "  Extend operand:\n    ");
     LLVM_DEBUG(((MachineInstr *)MIBSXTW)->print(dbgs()));
   } else if (Opc == AArch64::LDR_ZXI || Opc == AArch64::STR_ZXI) {
@@ -1239,6 +1338,45 @@ AArch64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
     MOp1.setReg(AArch64::Q0 + (MOp1.getReg() - AArch64::Z0));
     LLVM_DEBUG(((MachineInstr *)MIB)->print(dbgs()));
   } else {
+
+    // In the case that the merge doesn't result in a sign-extend, if we have
+    // something like:
+    // debugValueSubstitutions:[]
+    // $x1 = LDRXui $x0, 1, debug-instr-number 1
+    // DBG_INSTR_REF !13, dbg-instr-ref(1, 0), debug-location !11
+    // $x0 = LDRXui killed $x0, 0, debug-instr-number 2
+    // DBG_INSTR_REF !14, dbg-instr-ref(2, 0), debug-location !11
+
+    // It will be converted to:
+    // debugValueSubstitutions: []
+    // $x0, $x1 = LDPXi $x0, 0
+    // DBG_INSTR_REF !12, dbg-instr-ref(1, 0), debug-location !14
+    // DBG_INSTR_REF !13, dbg-instr-ref(2, 0), debug-location !14
+
+    // We want the final result to look like:
+    // debugValueSubstitutions:
+    // - { srcinst: 1, srcop: 0, dstinst: 3, dstop: 1, subreg: 0 }
+    // - { srcinst: 2, srcop: 0, dstinst: 3, dstop: 0, subreg: 0 }
+    // $x0, $x1 = LDPXi $x0, 0, debug-instr-number 3
+    // DBG_INSTR_REF !12, dbg-instr-ref(1, 0), debug-location !14
+    // DBG_INSTR_REF !12, dbg-instr-ref(2, 0), debug-location !14
+
+    // Here all that needs to be done is, that the LDP instruction needs to be
+    // updated with a new debug-instr-number, we then need to add entries into
+    // the debugSubstitutions table to map the old instr-refs to the new ones.
+
+    // Assign new DebugInstrNum to the Paired instruction.
+    if (I->peekDebugInstrNum()) {
+      unsigned NewDebugInstrNum = MIB->getDebugInstrNum();
+      addDebugSubstitutionsToTable(MBB->getParent(), NewDebugInstrNum, *I,
+                                   *MIB);
+    }
+    if (Paired->peekDebugInstrNum()) {
+      unsigned NewDebugInstrNum = MIB->getDebugInstrNum();
+      addDebugSubstitutionsToTable(MBB->getParent(), NewDebugInstrNum, *Paired,
+                                   *MIB);
+    }
+
     LLVM_DEBUG(((MachineInstr *)MIB)->print(dbgs()));
   }
   LLVM_DEBUG(dbgs() << "\n");
diff --git a/llvm/test/CodeGen/AArch64/aarch64-ldst-opt-instr-ref.mir b/llvm/test/CodeGen/AArch64/aarch64-ldst-opt-instr-ref.mir
new file mode 100644
index 0000000..5a6b4c7
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/aarch64-ldst-opt-instr-ref.mir
@@ -0,0 +1,91 @@
+# RUN: llc -mtriple=aarch64-unknown-linux-gnu -o - %s -run-pass=aarch64-ldst-opt | FileCheck %s
+
+# This testcase was obtained by looking at FileCheck.cpp and reducing it down via llvm-reduce
+
+# The aarch64-ldst-opt pass tries to merge load instructions from LDR* to a load pair or LDP* instruction, in such a case, we must ensure that the debug-instr-number is properly preserved for instruction referencing.
+
+# Check that in the case of a sign extend, the debug instruction number is transferred to the sign extend instruction (SBFMXri in this case), whereas the LDP instruction gets the other debug instruction number for the load that doesn't get sign extended.
+
+# CHECK-LABEL: name:            _ZNK4llvm7Pattern5matchENS_9StringRefERKNS_9SourceMgrE
+# CHECK: debugValueSubstitutions:
+# CHECK-NEXT: - { srcinst: [[DBG_INSTR_NUM1:[0-9+]]], srcop: [[DBG_INSTR_OP1:[0-9+]]], dstinst: [[DBG_INSTR_NUM2:[0-9+]]], dstop: 1, subreg: 0 }
+# CHECK-NEXT: - { srcinst: [[DBG_INSTR_NUM3:[0-9+]]], srcop: [[DBG_INSTR_OP2:[0-9+]]], dstinst: [[DBG_INSTR_NUM4:[0-9+]]], dstop: 0, subreg: 0 }
+
+# CHECK: $w[[REG1:[0-9+]]], renamable $w[[REG2:[0-9+]]] = LDPWi renamable $x[[REG1]], 0, debug-instr-number [[DBG_INSTR_NUM2]]
+# CHECK-NEXT: $w[[REG1]] = KILL $w[[REG1]], implicit-def $x[[REG1]]
+# CHECK-NEXT: $x[[REG1]] = SBFMXri $x[[REG1]], 0, 31, debug-instr-number [[DBG_INSTR_NUM4]]
+# CHECK-NEXT: DBG_INSTR_REF !{{[0-9+]}}, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 64), dbg-instr-ref([[DBG_INSTR_NUM1]], [[DBG_INSTR_OP1]]), debug-location !{{[0-9+]}}
+# CHECK-NEXT: DBG_INSTR_REF !7, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 32), dbg-instr-ref([[DBG_INSTR_NUM3]], [[DBG_INSTR_OP2]]), debug-location !{{[0-9+]}}
+
+# Check that in the case there is no sign extend, the LDP instruction gets a new debug instruction number and both the DBG_INSTR_REFs use the new instruction number.
+
+# CHECK-LABEL: name:            _ZNK4llvm7Pattern5matchENS_9StringRefERKNS_9SourceMgrE2
+# CHECK: debugValueSubstitutions:
+# CHECK-NEXT: - { srcinst: [[DBG_INSTR_NUM5:[0-9+]]], srcop: [[DBG_INSTR_OP3:[0-9+]]], dstinst: [[DBG_INSTR_NUM6:[0-9+]]], dstop: 1, subreg: 0 }
+# CHECK-NEXT: - { srcinst: [[DBG_INSTR_NUM7:[0-9+]]], srcop: [[DBG_INSTR_OP4:[0-9+]]], dstinst: [[DBG_INSTR_NUM6]], dstop: 0, subreg: 0 }
+
+# CHECK: renamable $x[[REG3:[0-9+]]], renamable $x[[REG4:[0-9+]]] = LDPXi renamable $x[[REG3]], 0, debug-instr-number [[DBG_INSTR_NUM6]]
+# CHECK-NEXT: DBG_INSTR_REF !12, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 64), dbg-instr-ref([[DBG_INSTR_NUM5]], [[DBG_INSTR_OP3]]), debug-location !14
+# CHECK-NEXT: DBG_INSTR_REF !12, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 32), dbg-instr-ref([[DBG_INSTR_NUM7]], [[DBG_INSTR_OP4]]), debug-location !14
+
+--- |
+  define i64 @_ZNK4llvm9StringRef4sizeEv(ptr readonly captures(none) %this) local_unnamed_addr #0 {
+  entry:
+    %Length = getelementptr i8, ptr %this, i64 8
+    %0 = load i64, ptr %Length, align 4
+    ret i64 %0
+  }
+  define ptr @_ZNK4llvm9StringRef4dataEv(ptr readonly captures(none) %this) local_unnamed_addr #0 {
+  entry:
+    %0 = load ptr, ptr %this, align 8
+    ret ptr %0
+  }
+  define void @_ZNK4llvm7Pattern5matchENS_9StringRefERKNS_9SourceMgrE(ptr readonly captures(none) %agg.result) local_unnamed_addr !dbg !3 {
+    %call1541 = load volatile ptr, ptr null, align 4294967296, !dbg !9
+    %FullMatch.sroa.1.0.agg.result.sroa_idx = getelementptr inbounds nuw i8, ptr %agg.result, i64 8
+    ret void
+  }
+  define void @_ZNK4llvm7Pattern5matchENS_9StringRefERKNS_9SourceMgrE2(ptr readonly captures(none) %agg.result) local_unnamed_addr !dbg !10 {
+    %call1541 = load volatile ptr, ptr null, align 4294967296, !dbg !11
+    %FullMatch.sroa.1.0.agg.result.sroa_idx = getelementptr inbounds nuw i8, ptr %agg.result, i64 8
+    ret void
+  }
+  !llvm.module.flags = !{!0}
+  !llvm.dbg.cu = !{!1}
+  !0 = !{i32 2, !"Debug Info Version", i32 3}
+  !1 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !2, isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, sdk: "MacOSX15.3.sdk")
+  !2 = !DIFile(filename: "/Users/shubhamrastogi/Development/llvm-project-instr-ref/llvm-project/llvm/lib/FileCheck/FileCheck.cpp", directory: "/Users/shubhamrastogi/Development/llvm-project-instr-ref/llvm-project/build-baseline-stage2", checksumkind: CSK_MD5, checksum: "ac1d2352ab68b965fe7993c780cf92d7")
+  !3 = distinct !DISubprogram(scope: null, type: !4, spFlags: DISPFlagDefinition, unit: !1, retainedNodes: !6)
+  !4 = distinct !DISubroutineType(types: !5)
+  !5 = !{}
+  !6 = !{!7}
+  !7 = !DILocalVariable(name: "FullMatch", scope: !3, line: 1152, type: !8)
+  !8 = distinct !DICompositeType(tag: DW_TAG_class_type, size: 128, identifier: "_ZTSN4llvm9StringRefE")
+  !9 = !DILocation(line: 0, scope: !3)
+  !10 = distinct !DISubprogram(scope: null, type: !4, spFlags: DISPFlagDefinition, unit: !1, retainedNodes: !12)
+  !11 = !DILocation(line: 0, scope: !10)
+  !12 = !{!13}
+  !13 = !DILocalVariable(name: "FullMatch", scope: !10, line: 1152, type: !14)
+  !14 = distinct !DICompositeType(tag: DW_TAG_class_type, size: 128, identifier: "_ZTSN4llvm9StringRefE")
+  
+name:            _ZNK4llvm9StringRef4sizeEv
+---
+name:            _ZNK4llvm9StringRef4dataEv
+...
+name:            _ZNK4llvm7Pattern5matchENS_9StringRefERKNS_9SourceMgrE
+debugValueSubstitutions: []
+body:             |
+  bb.0 (%ir-block.0):
+    renamable $w1 = LDRWui renamable $x0, 1, debug-instr-number 1 :: (load (s64) from %ir.FullMatch.sroa.1.0.agg.result.sroa_idx, align 1)
+    DBG_INSTR_REF !7, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 64), dbg-instr-ref(1, 0), debug-location !9
+    renamable $x0 = LDRSWui killed renamable $x0, 0, debug-instr-number 2 :: (load (s64) from %ir.agg.result)
+    DBG_INSTR_REF !7, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 32), dbg-instr-ref(2, 0), debug-location !9
+...
+name:            _ZNK4llvm7Pattern5matchENS_9StringRefERKNS_9SourceMgrE2
+debugValueSubstitutions: []
+body:             |
+  bb.0 (%ir-block.0):
+    renamable $x1 = LDRXui renamable $x0, 1, debug-instr-number 1 :: (load (s64) from %ir.FullMatch.sroa.1.0.agg.result.sroa_idx, align 1)
+    DBG_INSTR_REF !13, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 64), dbg-instr-ref(1, 0), debug-location !11
+    renamable $x0 = LDRXui killed renamable $x0, 0, debug-instr-number 2 :: (load (s64) from %ir.agg.result)
+    DBG_INSTR_REF !13, !DIExpression(DW_OP_LLVM_arg, 0, DW_OP_LLVM_fragment, 64, 32), dbg-instr-ref(2, 0), debug-location !11