[AArch64] Preserve register flags when promoting a load from store.

[fhahn]

Summary

This patch updates promoteLoadFromStore to use the store MachineOperand as the
source operand of the of the new instruction instead of creating a new
register MachineOperand. This way, the existing register flags are
preserved.

This fixes PR33468 (https://bugs.llvm.org/show_bug.cgi?id=33468).

[t.p.northover]

I'm not sure that's valid. It seem like we'd still be broken in this situation:

STR %val, %base, ...
[...]
SOMETHING %val<killed>, ...
[...]
%whatever = LDR %base

findMatchingStore looks like it'd allow this (%val isn't actually changed) but the liveness is still broken after replacing the LDR with a MOV.

[fhahn]

Agreed. It seems we have to keep track of the flags for all instructions using the store operand (between the store and load) and use that when creating the promoted operand. findMachtingStore would be a good place to do that, no?

[t.p.northover]

I'm not sure findMatchingStore can do it easily. During its reverse scan it doesn't know what the value register is going to be and keeping track of the most recent kill of all registers seems a little excessive.

I think the code clearing kill flags later on (~line 877) probably has to have a reference to the MachineOperand and set its kill flag if necessary.

[junbuml]

I think the code clearing kill flags later on (~line 877) probably has to have a reference to the MachineOperand and set its kill flag if necessary.

+1

[fhahn]

Thanks, that looks like an excellent place! Updated the patch to propagate the kill flag to the promoted operand.

[junbuml]

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

885

Isn't it okay to do break after this ?

[fhahn]

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

885

Yep it should be OK, assuming there are no redundant kills later on, but then it's not the responsibility of this pass to clean that up.

[fhahn]

Added an early exit after the first kill flag has been found.

[MatzeB]

The first part of the patch looks fine to me.

The second part of adding an extra kill flag should be unnecessary and I'd prefer staying with the naive code:

• Kill flags are optional. So not adding them should never be a correctness problem.
• Missing kill flags may mean missed optimizations today, though I believe that I fixed all passes used by AArch64 in that respect so there should be no change. (Admittedly I still have to do some larger scaling testing with a debugging mode that drops all kill flags to verify it is true).
• The long term plan is to get rid of kill flags and hence having simpler code.

test/CodeGen/AArch64/ldst-opt.mir

153

Is it possible to make %w1 not live-in to make the test slightly more obvious.

[fhahn]

Ok, I'll update the patch and remove the code to propagate the kill flag. Do you think it's worth to keep the early exit at line 887 though?

test/CodeGen/AArch64/ldst-opt.mir

153

Yes I'll do that

[MatzeB]

In D34402#785706, @fhahn wrote:

Ok, I'll update the patch and remove the code to propagate the kill flag. Do you think it's worth to keep the early exit at line 887 though?

If you want to. Though may be a good idea to put it into a separate commit (without an extra review).

[fhahn]

Undo changes to promote kill flag, early exit and removed %w1 from liveins

[MatzeB]

Thanks, LGTM

lib/Target/AArch64/AArch64LoadStoreOptimizer.cpp

 AArch64LoadStoreOpt::promoteLoadFromStore(MachineBasicBlock::iterator LoadI,
                                           MachineBasicBlock::iterator StoreI) {
   MachineBasicBlock::iterator NextI = LoadI;
   ++NextI;
 
   int LoadSize = getMemScale(*LoadI);
   int StoreSize = getMemScale(*StoreI);
   unsigned LdRt = getLdStRegOp(*LoadI).getReg();
+  const MachineOperand &StMO = getLdStRegOp(*StoreI);
   unsigned StRt = getLdStRegOp(*StoreI).getReg();
   bool IsStoreXReg = TRI->getRegClass(AArch64::GPR64RegClassID)->contains(StRt);
 
   assert((IsStoreXReg ||
           TRI->getRegClass(AArch64::GPR32RegClassID)->contains(StRt)) &&
          "Unexpected RegClass");
 
   MachineInstr *BitExtMI;
   if (LoadSize == StoreSize && (LoadSize == 4 || LoadSize == 8)) {
     // Remove the load, if the destination register of the loads is the same
     // register for stored value.
     if (StRt == LdRt && LoadSize == 8) {
       StoreI->clearRegisterKills(StRt, TRI);
       DEBUG(dbgs() << "Remove load instruction:\n    ");
       DEBUG(LoadI->print(dbgs()));
       DEBUG(dbgs() << "\n");
       LoadI->eraseFromParent();
       return NextI;
     }
     // Replace the load with a mov if the load and store are in the same size.
     BitExtMI =
         BuildMI(*LoadI->getParent(), LoadI, LoadI->getDebugLoc(),
                 TII->get(IsStoreXReg ? AArch64::ORRXrs : AArch64::ORRWrs), LdRt)
             .addReg(IsStoreXReg ? AArch64::XZR : AArch64::WZR)
-            .addReg(StRt)
+            .add(StMO)
             .addImm(AArch64_AM::getShifterImm(AArch64_AM::LSL, 0));
   } else {
     // FIXME: Currently we disable this transformation in big-endian targets as
     // performance and correctness are verified only in little-endian.
     if (!Subtarget->isLittleEndian())
       return NextI;
     bool IsUnscaled = TII->isUnscaledLdSt(*LoadI);
     assert(IsUnscaled == TII->isUnscaledLdSt(*StoreI) &&
if (UnscaledLdOffset == UnscaledStOffset) {
                                 | ((Immr) << 6)               // immr
                                 | ((Imms) << 0)               // imms
           ;
 
       BitExtMI =
           BuildMI(*LoadI->getParent(), LoadI, LoadI->getDebugLoc(),
                   TII->get(IsStoreXReg ? AArch64::ANDXri : AArch64::ANDWri),
                   DestReg)
-              .addReg(StRt)
+              .add(StMO)
               .addImm(AndMaskEncoded);
     } else {
       BitExtMI =
           BuildMI(*LoadI->getParent(), LoadI, LoadI->getDebugLoc(),
                   TII->get(IsStoreXReg ? AArch64::UBFMXri : AArch64::UBFMWri),
                   DestReg)
-              .addReg(StRt)
+              .add(StMO)
               .addImm(Immr)
               .addImm(Imms);
     }
   }
 
-  // Clear kill flags between store and load.
+  // Clear kill flags between store and load and also propagate the kill flag
+  // to the promoted operand.
   for (MachineInstr &MI : make_range(StoreI->getIterator(),
                                      BitExtMI->getIterator()))
+    if (MI.killsRegister(StRt, TRI)) {
       MI.clearRegisterKills(StRt, TRI);
+      // Load-store promotion is only legal if StRt is not redefined between
+      // the store and load instruction, so we can just propagate the kill flag.
+      BitExtMI->addRegisterKilled(StRt, TRI);
+    }
 
   DEBUG(dbgs() << "Promoting load by replacing :\n    ");
   DEBUG(StoreI->print(dbgs()));
   DEBUG(dbgs() << "    ");
   DEBUG(LoadI->print(dbgs()));
   DEBUG(dbgs() << "  with instructions:\n    ");
   DEBUG(StoreI->print(dbgs()));
   DEBUG(dbgs() << "    ");

test/CodeGen/AArch64/ldst-opt.mir

bb.0:
     CFI_INSTRUCTION 0
     %w0 = LDRHHui killed %x0, 1 :: (load 2)
     RET %lr, implicit %w0
 
 ...
 # Don't count transient instructions towards search limits.
 # CHECK-LABEL: name: promote-load-from-store
 # CHECK: STRWui %w1
-# CHECK: UBFMWri %w1
+# CHECK: UBFMWri killed %w1
 ---
 name: store-pair
 tracksRegLiveness: true
 body: |
   bb.0:
     liveins: %w1, %x0, %lr
 
     STRWui %w1, %x0, 0 :: (store 4)
bb.0:
     liveins: %w1
 
     STRWui %w1, %sp, 0 :: (store 4)
     %wzr = COPY killed %w1 ; killing use of %w1
     %w11 = LDRWui %sp, 0 :: (load 4)
     HINT 0, implicit %w11 ; some use of %w11
 ...
 # When replaceing the load of a store-load pair with a copy the kill flags
-# along the way need to be cleared.
+# along the way need to be cleared and the kill flags needs to be set on the
+# promoted operand.
 # CHECK-LABEL: name: store-load-clearkill
 # CHECK: STRWui %w1, %sp, 0 :: (store 4)
 # CHECK-NOT: COPY killed %w1
 # CHECK: %wzr = COPY %w1
-# CHECK: %w11 = ORRWrs %wzr, %w1, 0
+# CHECK: %w11 = ORRWrs %wzr, killed %w1, 0
 # CHECK: HINT 0, implicit %w11
+---
+name: promote-load-from-store-undef
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: %w1, %x0, %x2, %lr
+
+    STRWui undef %w1, %x0, 0 :: (store 4)
+    %w0 = LDRBBui %x0, 1 :: (load 2)
+    STRHHui undef %w3, %x2, 0 :: (store 4)
+    %w1 = LDRBBui %x2, 0 :: (load 4)
+    RET %lr, implicit %w0
+...
+# CHECK-LABEL: name: promote-load-from-store-undef
+# CHECK: STRWui undef %w1
+# CHECK: UBFMWri undef %w1
+# CHECK: STRHHui undef %w3
+# CHECK: ANDWri undef %w3