From b6ee831b59ef11b424c819d272216add268e6164 Mon Sep 17 00:00:00 2001
From: David Green <david.green@arm.com>
Date: Wed, 29 Nov 2023 07:41:15 +0000
Subject: [PATCH] [AArch64] Load/store optimizer fixes and cleanup.

This includes a couple of fixes after #71908 for bundles and some cleanup for
the debug output. One was an iterator type that asserted on bundles, the second
a rather subtle issue where forAllMIsUntilDef would hit the LdStLimit when
renaming registers, meaning the last instruction was not updated leaving an
invalid `ldp x6, x6` instruction.
---
 .../AArch64/AArch64LoadStoreOptimizer.cpp     |  35 +++---
 .../CodeGen/AArch64/stp-opt-with-renaming.mir | 118 +++++++++++++++++-
 2 files changed, 129 insertions(+), 24 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
index a03a7f8737be0e..dc6d5b8950c34a 100644
--- a/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
+++ b/llvm/lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp
@@ -941,11 +941,11 @@ AArch64LoadStoreOpt::mergePairedInsns(MachineBasicBlock::iterator I,
               }
             }
           }
-          LLVM_DEBUG(dbgs() << "Renamed " << MI << "\n");
+          LLVM_DEBUG(dbgs() << "Renamed " << MI);
           return true;
         };
     forAllMIsUntilDef(MergeForward ? *I : *std::prev(Paired), RegToRename, TRI,
-                      LdStLimit, UpdateMIs);
+                      UINT32_MAX, UpdateMIs);
 
 #if !defined(NDEBUG)
     // For forward merging store:
@@ -1464,7 +1464,7 @@ canRenameUpToDef(MachineInstr &FirstMI, LiveRegUnits &UsedInBetween,
                        MOP.isImplicit() && MOP.isKill() &&
                        TRI->regsOverlap(RegToRename, MOP.getReg());
               })) {
-    LLVM_DEBUG(dbgs() << "  Operand not killed at " << FirstMI << "\n");
+    LLVM_DEBUG(dbgs() << "  Operand not killed at " << FirstMI);
     return false;
   }
 
@@ -1476,11 +1476,11 @@ canRenameUpToDef(MachineInstr &FirstMI, LiveRegUnits &UsedInBetween,
   // * collect the registers used and required register classes for RegToRename.
   std::function<bool(MachineInstr &, bool)> CheckMIs = [&](MachineInstr &MI,
                                                            bool IsDef) {
-    LLVM_DEBUG(dbgs() << "Checking " << MI << "\n");
+    LLVM_DEBUG(dbgs() << "Checking " << MI);
     // Currently we do not try to rename across frame-setup instructions.
     if (MI.getFlag(MachineInstr::FrameSetup)) {
-      LLVM_DEBUG(dbgs() << "  Cannot rename framesetup instructions currently ("
-                        << MI << ")\n");
+      LLVM_DEBUG(dbgs() << "  Cannot rename framesetup instructions "
+                        << "currently\n");
       return false;
     }
 
@@ -1500,8 +1500,7 @@ canRenameUpToDef(MachineInstr &FirstMI, LiveRegUnits &UsedInBetween,
       //       1. Insert an extra copy, to materialize the def.
       //       2. Skip pseudo-defs until we find an non-pseudo def.
       if (MI.isPseudo()) {
-        LLVM_DEBUG(dbgs() << "  Cannot rename pseudo instruction " << MI
-                          << "\n");
+        LLVM_DEBUG(dbgs() << "  Cannot rename pseudo/bundle instruction\n");
         return false;
       }
 
@@ -1510,8 +1509,7 @@ canRenameUpToDef(MachineInstr &FirstMI, LiveRegUnits &UsedInBetween,
             !TRI->regsOverlap(MOP.getReg(), RegToRename))
           continue;
         if (!canRenameMOP(MOP, TRI)) {
-          LLVM_DEBUG(dbgs()
-                     << "  Cannot rename " << MOP << " in " << MI << "\n");
+          LLVM_DEBUG(dbgs() << "  Cannot rename " << MOP << " in " << MI);
           return false;
         }
         RequiredClasses.insert(TRI->getMinimalPhysRegClass(MOP.getReg()));
@@ -1524,8 +1522,7 @@ canRenameUpToDef(MachineInstr &FirstMI, LiveRegUnits &UsedInBetween,
           continue;
 
         if (!canRenameMOP(MOP, TRI)) {
-          LLVM_DEBUG(dbgs()
-                     << "  Cannot rename " << MOP << " in " << MI << "\n");
+          LLVM_DEBUG(dbgs() << "  Cannot rename " << MOP << " in " << MI);
           return false;
         }
         RequiredClasses.insert(TRI->getMinimalPhysRegClass(MOP.getReg()));
@@ -1565,15 +1562,12 @@ static bool canRenameUntilSecondLoad(
   auto RegToRename = getLdStRegOp(FirstLoad).getReg();
   bool Success = std::all_of(
       FirstLoad.getIterator(), SecondLoad.getIterator(),
-      [&](MachineBasicBlock::iterator MBBI) {
-        MachineInstr &MI = *MBBI;
-
-        LLVM_DEBUG(dbgs() << "Checking " << MI << "\n");
+      [&](MachineInstr &MI) {
+        LLVM_DEBUG(dbgs() << "Checking " << MI);
         // Currently we do not try to rename across frame-setup instructions.
         if (MI.getFlag(MachineInstr::FrameSetup)) {
-          LLVM_DEBUG(dbgs()
-                     << "  Cannot rename framesetup instructions currently ("
-                     << MI << ")\n");
+          LLVM_DEBUG(dbgs() << "  Cannot rename framesetup instructions "
+                            << "currently\n");
           return false;
         }
 
@@ -1582,8 +1576,7 @@ static bool canRenameUntilSecondLoad(
               !TRI->regsOverlap(MOP.getReg(), RegToRename))
             continue;
           if (!canRenameMOP(MOP, TRI)) {
-            LLVM_DEBUG(dbgs()
-                       << "  Cannot rename " << MOP << " in " << MI << "\n");
+            LLVM_DEBUG(dbgs() << "  Cannot rename " << MOP << " in " << MI);
             return false;
           }
           RequiredClasses.insert(TRI->getMinimalPhysRegClass(MOP.getReg()));
diff --git a/llvm/test/CodeGen/AArch64/stp-opt-with-renaming.mir b/llvm/test/CodeGen/AArch64/stp-opt-with-renaming.mir
index 6168a02b3d9574..cbb7f09513aa15 100644
--- a/llvm/test/CodeGen/AArch64/stp-opt-with-renaming.mir
+++ b/llvm/test/CodeGen/AArch64/stp-opt-with-renaming.mir
@@ -622,7 +622,7 @@ body:             |
     STRQui killed renamable $q0, $sp, 2 :: (store 16, align 128)
     RET undef $lr
 
-... 
+...
 ---
 #
 # CHECK-LABEL: name: ldst-no-reg-available
@@ -650,7 +650,7 @@ body:             |
     STRQui killed renamable $q0, $sp, 2 :: (store 16, align 32)
     RET undef $lr
 
-... 
+...
 ---
 #
 # CHECK-LABEL: name: ldst-basereg-modified
@@ -677,7 +677,7 @@ body:             |
     STRQui killed renamable $q0, $sp, 2 :: (store 16, align 32)
     RET undef $lr
 
-... 
+...
 ---
 #
 # CHECK-LABEL: name: ldr-dest-reg-implicit-killed
@@ -700,3 +700,115 @@ body:             |
     renamable $q0 = LDRQui $sp, 0 :: (load 16, align 32)
     STRSui renamable $s0, $sp, 10, implicit killed $q0 :: (store (s32))
     RET undef $lr
+
+...
+---
+#
+# CHECK-LABEL: name: bundled
+# CHECK: renamable $q0, $q1 = LDPQi $sp, 0 :: (load (s128)), (load (s128), align 32)
+# CHECK-NEXT: STPSi $s1, renamable $s0, $sp, 9 :: (store (s32))
+# CHECK-NEXT: BUNDLE implicit-def $z3
+# CHECK-NEXT:  $z3 = MOVPRFX_ZZ $z19
+# CHECK-NEXT:  $z3 = FMUL_ZPmZ_S renamable $p0, killed $z3, renamable $z16
+# CHECK-NEXT: }
+# CHECK-NEXT: RET undef $lr
+#
+name:            bundled
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0.entry:
+    liveins: $z3, $z19, $p0, $z16
+    renamable $q0 = LDRQui $sp, 1 :: (load 16)
+    STRSui renamable $s0, $sp, 9, implicit killed $q0 :: (store (s32))
+    BUNDLE implicit-def $z3, implicit-def $q3, implicit-def $d3, implicit-def $s3, implicit-def $h3, implicit-def $b3, implicit $z19, implicit $p0, implicit $z16 {
+      $z3 = MOVPRFX_ZZ $z19
+      $z3 = FMUL_ZPmZ_S renamable $p0, killed $z3, renamable $z16
+    }
+    renamable $q0 = LDRQui $sp, 0 :: (load 16, align 32)
+    STRSui renamable $s0, $sp, 10, implicit killed $q0 :: (store (s32))
+    RET undef $lr
+...
+---
+#
+# CHECK-LABEL: name: bundled_use
+# CHECK: renamable $q0 = LDRQui $sp, 1
+# CHECK-NEXT: STRQui renamable $q0, $sp, 3, implicit-def $z0
+# CHECK-NEXT: BUNDLE implicit-def $z3
+# CHECK-NEXT:  $z3 = MOVPRFX_ZZ $z19
+# CHECK-NEXT:  $z3 = FMUL_ZPmZ_S renamable $p0, killed $z3, renamable $z0
+# CHECK-NEXT: }
+# CHECK-NEXT: renamable $q0 = LDRQui $sp, 0
+# CHECK-NEXT: STRQui killed renamable $q0, $sp, 2
+# CHECK-NEXT: RET undef $lr
+#
+name:            bundled_use
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0.entry:
+    liveins: $z3, $z19, $p0, $z16
+    renamable $q0 = LDRQui $sp, 1 :: (load 16)
+    STRQui renamable $q0, $sp, 3, implicit-def $z0 :: (store 16, basealign 128)
+    BUNDLE implicit-def $z3, implicit-def $q3, implicit-def $d3, implicit-def $s3, implicit-def $h3, implicit-def $b3, implicit $z19, implicit $p0, implicit $z0, implicit $q0 {
+      $z3 = MOVPRFX_ZZ $z19
+      $z3 = FMUL_ZPmZ_S renamable $p0, killed $z3, renamable $z0
+    }
+    renamable $q0 = LDRQui $sp, 0 :: (load 16, align 128)
+    STRQui killed renamable $q0, $sp, 2 :: (store 16, align 128)
+    RET undef $lr
+
+...
+---
+#
+# CHECK-LABEL: name: bundle_at_limit
+# CHECK: renamable $x3, $x4 = LDPXi $sp, 1 :: (load (s64))
+#
+name:            bundle_at_limit
+alignment:       4
+tracksRegLiveness: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0.entry:
+    liveins: $x0, $x1, $x2, $z0, $p0, $z3, $z19
+    renamable $x3 = LDRXui $sp, 2 :: (load (s64))
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x3, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x2, 1, 0
+    $x2 = ADDXri renamable $x3, 1, 0
+    BUNDLE implicit-def $z3, implicit-def $q3, implicit-def $d3, implicit-def $s3, implicit-def $h3, implicit-def $b3, implicit $z19, implicit $p0, implicit $z0 {
+      $z3 = MOVPRFX_ZZ $z19
+      $z3 = FMUL_ZPmZ_S renamable $p0, killed $z3, renamable $z0
+    }
+    $x2 = ADDXri renamable $x3, 1, 0
+    renamable $x3 = LDRXui $sp, 1 :: (load (s64))
+    RET undef $lr
+
+...