[InstrRef][AArch64][3/4] Recover variables from aarch64 load/store combiner

[jmorse]

Summary

This pass combines several load or store instructions into a single instruction that defines several registers at a time. With the instruction referencing model of tracking variable values for debug-info, we need to record that this optimisation took place, mapping the instruction numbers of any optimised-out loads onto the relevant operands of the new LDP instruction. Stores should not be labelled.

The test added borrows functions from other aarch64 tests. Right now, the "case4" function produces a LDPXpost instruction that I'm not quite sure what to do with. That's left unlabelled (meaning variable locations are safely dropped) to be implemented as a future exercise.

llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

#endif
 
   MIB.add(RegOp0)
       .add(RegOp1)
       .add(BaseRegOp)
       .addImm(OffsetImm)
       .cloneMergedMemRefs({&*I, &*Paired})
       .setMIFlags(I->mergeFlagsWith(*Paired));
 
-  (void)MIB;
+  MachineInstr *NewInstr = MIB;
+
+  // Lambda for debug substitutions: transfer any instruction number on an old
+  // instruction to the specified opernad on the new instruction.
+  auto &MF = *NewInstr->getMF();
+  auto DebugSub = [&MF,NewInstr](const MachineInstr *OldInstr, unsigned NewOp) {
+    if (unsigned OldInstrNum = OldInstr->peekDebugInstrNum()) {
+      unsigned NewInstrNum = NewInstr->getDebugInstrNum();
+      MF.makeDebugValueSubstitution({OldInstrNum, 0}, {NewInstrNum, NewOp}, 0);
+    }
+  };
 
   LLVM_DEBUG(
       dbgs() << "Creating pair load/store. Replacing instructions:\n    ");
   LLVM_DEBUG(I->print(dbgs()));
   LLVM_DEBUG(dbgs() << "    ");
   LLVM_DEBUG(Paired->print(dbgs()));
   LLVM_DEBUG(dbgs() << "  with instruction:\n    ");
   if (SExtIdx != -1) {
MachineInstrBuilder MIBKill =
             .addReg(DstRegX, RegState::Define);
     MIBKill->getOperand(2).setImplicit();
     // Create the sign extension.
     MachineInstrBuilder MIBSXTW =
         BuildMI(*MBB, InsertionPoint, DL, TII->get(AArch64::SBFMXri), DstRegX)
             .addReg(DstRegX)
             .addImm(0)
             .addImm(31);
-    (void)MIBSXTW;
     LLVM_DEBUG(dbgs() << "  Extend operand:\n    ");
     LLVM_DEBUG(((MachineInstr *)MIBSXTW)->print(dbgs()));
+
+    // For the non-extended value, point it to the new instruction; the extended
+    // value should point at the SBFM.
+    auto DebugSubSBFM = [&MF,&MIBSXTW](const MachineInstr *OldInstr) {
+      if (unsigned OldInstrNum = OldInstr->peekDebugInstrNum()) {
+        unsigned NewInstrNum = MIBSXTW->getDebugInstrNum();
+        MF.makeDebugValueSubstitution({OldInstrNum, 0}, {NewInstrNum, 0}, 0);
+      }
+    };
+
+    if (SExtIdx == 0) {
+      DebugSub(Rt2MI, 1);
+      DebugSubSBFM(RtMI);
+    } else {
+      DebugSub(RtMI, 0);
+      DebugSubSBFM(Rt2MI);
+    }
   } else {
     LLVM_DEBUG(((MachineInstr *)MIB)->print(dbgs()));
+
+    // Point debug variable values from RtMI to operand 0, Rt2MI to operand 1.
+    DebugSub(RtMI, 0);
+    DebugSub(Rt2MI, 1);
   }
   LLVM_DEBUG(dbgs() << "\n");
 
   if (MergeForward)
     for (const MachineOperand &MOP : phys_regs_and_masks(*I))
       if (MOP.isReg() && MOP.isKill())
         DefinedInBB.addReg(MOP.getReg());
 

llvm/test/DebugInfo/MIR/AArch64/ldst-merge-instr-ref.mir

+# RUN: llc -run-pass=aarch64-ldst-opt -o - -mtriple=aarch64-- %s | FileCheck %s
+# Copied from ldst-miflags.mir, test for cases where instruction references
+# should be preserved when loads and stores are merged.
+---
+name:            case3
+# CHECK-LABEL: name: case3
+body:             |
+  bb.0:
+
+    $x0 = frame-setup LDRXui $x2, 0, debug-instr-number 1 :: (load 8)
+    $x1 = frame-destroy LDRXui $x2, 1, debug-instr-number 2 :: (load 8)
+
+  ; CHECK:      debugValueSubstitutions:
+  ; CHECK-NEXT: - { srcinst: 1, srcop: 0, dstinst: 3, dstop: 0, subreg: 0 }
+  ; CHECK-NEXT: - { srcinst: 2, srcop: 0, dstinst: 3, dstop: 1, subreg: 0 }
+  ; CHECK: frame-setup frame-destroy LDPXi {{.*}} debug-instr-number 3
+    RET_ReallyLR
+...
+---
+name:            case4
+# CHECK-LABEL: name: case4
+body:             |
+  bb.0:
+    $x26, $x25 = frame-setup LDPXi $sp, 0, debug-instr-number 1
+    $sp = frame-destroy ADDXri $sp, 64, 0
+
+  ;; TODO: support this addressing mode optimisation.
+  ;; Ensure we don't label to LDPXpost in the meantime.
+  ; CHECK-NOT: LDPXpost {{.*}} debug-instr-number
+  ; CHECK: = frame-setup frame-destroy LDPXpost
+  ; CHECK-NOT: LDPXpost {{.*}} debug-instr-number
+    RET_ReallyLR
+
+...
+---
+name:            ldp_half_sext_res0_int
+# CHECK-LABEL: name: ldp_half_sext_res0_int
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$x0' }
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0:
+    liveins: $x0
+  
+    renamable $x8 = LDRSWui renamable $x0, 0, debug-instr-number 1 :: (load 4)
+    renamable $w9 = LDRWui killed renamable $x0, 1, implicit-def $x9, debug-instr-number 2 :: (load 4)
+    $x0 = ADDXrs killed renamable $x9, killed renamable $x8, 0
+    RET undef $lr, implicit $x0
+
+  ;; The two loads get merged, but one only has the right value after sign
+  ;; extension. Label the sign extending instruction.
+  ; CHECK:      debugValueSubstitutions:
+  ; CHECK-NEXT: - { srcinst: 1, srcop: 0, dstinst: 4, dstop: 0, subreg: 0 }
+  ; CHECK-NEXT: - { srcinst: 2, srcop: 0, dstinst: 3, dstop: 1, subreg: 0 }
+  ; 
+  ; CHECK: $w8, renamable $w9 = LDPWi renamable $x0, 0, debug-instr-number 3
+  ; CHECK: $x8 = SBFMXri $x8, 0, 31, debug-instr-number 4
+...
+---
+name:            ldp_half_sext_res0_int_flipped
+# CHECK-LABEL: name: ldp_half_sext_res0_int_flipped
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$x0' }
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0:
+    liveins: $x0
+  
+    renamable $w9 = LDRWui renamable $x0, 1, implicit-def $x9, debug-instr-number 1 :: (load 4)
+    renamable $x8 = LDRSWui killed renamable $x0, 0, debug-instr-number 2 :: (load 4)
+    $x0 = ADDXrs killed renamable $x9, killed renamable $x8, 0
+    RET undef $lr, implicit $x0
+
+  ;; Similar to previous test, but with the instruction numbers swapped.
+  ; CHECK:      debugValueSubstitutions:
+  ; CHECK-NEXT: - { srcinst: 1, srcop: 0, dstinst: 3, dstop: 1, subreg: 0 }
+  ; CHECK-NEXT: - { srcinst: 2, srcop: 0, dstinst: 4, dstop: 0, subreg: 0 }
+  ; 
+  ; CHECK: $w8, renamable $w9 = LDPWi renamable $x0, 0, debug-instr-number 3
+  ; CHECK: $x8 = SBFMXri $x8, 0, 31, debug-instr-number 4
+
+
+...