[ARM][LowOverheadLoops] Make some stack spills valid for tail predication

[samtebbs]

Summary

This patch makes vector spills valid for tail predication when all loads from the same stack slot are within the loop.

[dmgreen]

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

1222

Is it worth splitting this out into a new function? To keep this stack logic separate in case it needs to grow even more sophisticated in the future.

1226

I think we may need to distinguish between a load that accesses SP and a stack spill slot. A SP address could be some other (potentially escaping) stack object. I think that info is usually available through the MMO, maybe via FrameInfo::isSpillSlotObjectIndex.

I was looking around for a way to get all the MachineInstr's that access a specific stack slot, but couldn't fine one better than searching through a lot of instructions.

1229–1231

I think this comment would be better if we explain what we do do, not why we don't use RDA (or maybe both). Explaining that we are searching all blocks after the loop for another access to the same stack slot.

1236

Use I or It or Idx or something?

[samtebbs]

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

1222

Good idea. Do you think a lambda or proper function would be best?

1226

Good spot, I'll have a look into that. Regarding other ways to find instructions that use a specific stack slot, I haven't been able to find a better way either.

1229–1231

👍

1236

👍

[samtebbs]

Break-out code to function, check frame index and rename loop counter.

[dmgreen]

The searching feels gnarly, like it can end up searching a lot of instructions. I don't know of a better idea though, and it should be rare that it needs to search.

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

1087–1165

This function can be made static.

1090–1092

getPseudoValue can return nullptr on failure, which needs to be checked.
The type should be const PseudoSourceValue *, not auto (the type isn't obvious from the function call).
Then you might be able use if (const auto *FS = dyn_cast<FixedStackPseudoSourceValue>(PseudoValue))

Or you may be able to use if (const auto *FS = dyn_cast_or_null<FixedStackPseudoSourceValue>(PseudoValue)) to get the same effect?

1106

This will have to check that it has at least one memoperand.

1121

Frame indexes look like they are usually int's.

Does GetFrameIndex return the value by parameter reference because FrameIndexes can be negative? So we can't return -1 for not found?

1122

Check for isSpillSlotObjectIndex(FI) too?

1124

I think you can remove the ", 1", and let it choose a default small size. Same below.

1139

If we find another stack store, does that mean we can stop the search?

1148

Should this check that Succ is not already in Frontier too?

1222

As a function looks good. It's nice and self-contained.

[samtebbs]

Make ValidateMVEStore static, use dyn_cast, check isSpillSlotObjectIndex, make GetFrameIndex return an int, remove the size for the SmallVector and check the size of memoperands.

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

1090–1092

The dyn_cast<...> option looks good to me. 👍

1106

Ah I thought that would be guaranteed by the instruction being a V[STR|LDR]WU32. I'll add a check.

1121

I chose for it to be taken as reference so that a boolean could be returned to indicate success, but if it can just return the frame index or -1 in case of error then that would be great. From looking through the codebase (this is a great use-case for more function documentation) it looks like frame indexes will be positive so I think returning -1 in case of failure is a good option.

1124

SmallPtrSet requires an initial size but I can remove it for SmallVector.

1139

If we have found another store to the same stack slot then we stop the search as any further loads in the same BB won't invalidate the store we're validating, since they are dependent on that second store instead.

1148

Yeah that can't hurt, however I can't find a nice contains function in SmallVector. I wish container classes had the same interface.

[dmgreen]

Thanks. This looks good, minus some mostly Nitpicks.

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

1090

const PseudoSourceValue* PseudoValue =
It still needs to check it's not null too. The store could have a non-pseudo MMO, for which this would return nullptr.

1104

== 1 is probably fine. I wouldn't expect it to be higher.

1121

I think they may be able to go negative, but a negative argument means a parameter passed via the stack (I think). So treating it the same as an invalid value sounds like it should work fine.

1124

Oh of course, yeah. I hadn't realized it didn't have the same interface.

SmallPtrSet sounds good, but you may want to increase the default to 4. It won't make a large difference, but 4 pointers are not going to hurt anyone.

1131

unsigned int -> unsigned, is more common in llvm.

1139

Sure. I was thinking in that case it wouldn't have to add any Successors, so we could potentially save from looking through a lot of other BB's.

1148

There is an is_contained(Frontier, Succ) in ADT, that is a wrapper around std::find.

[samtebbs]

Check for null pseudo value, change smallptrset size, check for just one memoperand and avoid checking successors when a VSTRW is found.

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

1090

Cheers, I'll add that.

1139

Ah yes that is a very good point! I'll make it account for that.

1148

Sweet, thanks

[dmgreen]

Thanks. LGTM

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

if (!Prev->getOperand(1).isIdenticalTo(MI->getOperand(1)) ||
     LLVM_DEBUG(dbgs() << "ARM Loops: Found VCTP with a different reaching "
                          "definition from the main VCTP");
     return false;
   }
   VCTPs.push_back(MI);
   return true;
 }
 
-bool LowOverheadLoop::ValidateMVEInst(MachineInstr* MI) {
+static bool ValidateMVEStore(MachineInstr *MI, MachineLoop *ML) {
+
+  auto GetFrameIndex = [](MachineMemOperand *Operand) {
+    const PseudoSourceValue *PseudoValue = Operand->getPseudoValue();
+    if (PseudoValue && PseudoValue->kind() == PseudoSourceValue::FixedStack) {
+      if (const auto *FS = dyn_cast<FixedStackPseudoSourceValue>(PseudoValue)) {
+        return FS->getFrameIndex();
+      }
+    }
+    return -1;
+  };
+
+  auto IsStackOp = [GetFrameIndex](MachineInstr *I) {
+    switch (I->getOpcode()) {
+    case ARM::MVE_VSTRWU32:
+    case ARM::MVE_VLDRWU32: {
+      return I->getOperand(1).getReg() == ARM::SP &&
+             I->memoperands().size() == 1 &&
+             GetFrameIndex(I->memoperands().front()) >= 0;
+    }
+    default:
+      return false;
+    }
+  };
+
+  // An unpredicated vector register spill is allowed if all of the uses of the
+  // stack slot are within the loop
+  if (MI->getOpcode() != ARM::MVE_VSTRWU32 || !IsStackOp(MI))
+    return false;
+
+  // Search all blocks after the loop for accesses to the same stack slot.
+  // ReachingDefAnalysis doesn't work for sp as it relies on registers being
+  // live-out (which sp never is) to know what blocks to look in
+  if (MI->memoperands().size() == 0)
+    return false;
+  int FI = GetFrameIndex(MI->memoperands().front());
+
+  MachineFrameInfo FrameInfo = MI->getParent()->getParent()->getFrameInfo();
+  if (FI == -1 || !FrameInfo.isSpillSlotObjectIndex(FI))
+    return false;
+
+  SmallVector<MachineBasicBlock *> Frontier;
+  ML->getExitBlocks(Frontier);
+  SmallPtrSet<MachineBasicBlock *, 4> Visited{MI->getParent()};
+  unsigned Idx = 0;
+  while (Idx < Frontier.size()) {
+    MachineBasicBlock *BB = Frontier[Idx];
+    bool LookAtSuccessors = true;
+    for (auto &I : *BB) {
+      if (!IsStackOp(&I) || I.memoperands().size() == 0)
+        continue;
+      if (GetFrameIndex(I.memoperands().front()) != FI)
+        continue;
+      // If this block has a store to the stack slot before any loads then we
+      // can ignore the block
+      if (I.getOpcode() == ARM::MVE_VSTRWU32) {
+        LookAtSuccessors = false;
+        break;
+      }
+      // If the store and the load are using the same stack slot then the
+      // store isn't valid for tail predication
+      if (I.getOpcode() == ARM::MVE_VLDRWU32)
+        return false;
+    }
+
+    if (LookAtSuccessors) {
+      for (auto Succ : BB->successors()) {
+        if (!Visited.contains(Succ) && !is_contained(Frontier, Succ))
+          Frontier.push_back(Succ);
+      }
+    }
+    Visited.insert(BB);
+    Idx++;
+  }
+
+  return true;
+}
+
+bool LowOverheadLoop::ValidateMVEInst(MachineInstr *MI) {
   if (CannotTailPredicate)
     return false;
 
   if (!shouldInspect(*MI))
     return true;
 
   if (MI->getOpcode() == ARM::MVE_VPSEL ||
       MI->getOpcode() == ARM::MVE_VPNOT) {
bool LowOverheadLoop::ValidateMVEInst(MachineInstr *MI) {
   if (isDomainMVE(MI) && RequiresExplicitPredication) {
     LLVM_DEBUG(if (!IsUse)
                dbgs() << "ARM Loops: Can't tail predicate: " << *MI);
     return IsUse;
   }
 
   // If the instruction is already explicitly predicated, then the conversion
   // will be fine, but ensure that all store operations are predicated.
-  if (MI->mayStore())
+  if (MI->mayStore() && !ValidateMVEStore(MI, &ML))
     return IsUse;
 
   // If this instruction defines the VPR, update the predicate for the
   // proceeding instructions.
   if (isVectorPredicate(MI)) {
     // Clear the existing predicate when we're not in VPT Active state,
     // otherwise we add to it.
     if (!isVectorPredicated(MI))
       VPTState::resetPredicate(MI);
     else
       VPTState::addPredicate(MI);
   }
 
   // Finally once the predicate has been modified, we can start a new VPT
   // block if necessary.
   if (isVPTOpcode(MI->getOpcode()))
     VPTState::CreateVPTBlock(MI);
 
   return true;
 }
 
 bool ARMLowOverheadLoops::runOnMachineFunction(MachineFunction &mf) {
   const ARMSubtarget &ST = static_cast<const ARMSubtarget&>(mf.getSubtarget());

llvm/test/CodeGen/Thumb2/LowOverheadLoops/vector_spill_in_loop.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=thumbv8.1m.main -mattr=+mve -run-pass=arm-low-overhead-loops %s -o - | FileCheck %s
+--- |
+  define hidden void @vector_spill_in_loop() {
+  entry:
+    ret void
+  }
+
+  define hidden void @vector_spill_load_outside() {
+  entry:
+    ret void
+  }
+...
+---
+name:            vector_spill_in_loop
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: -120, size: 16, alignment: 8,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+body:             |
+  ; CHECK-LABEL: name: vector_spill_in_loop
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   liveins: $q0, $r1, $r2, $r3, $r4, $r5, $r7, $r8, $r10, $r11, $r12
+  ; CHECK:   $r6 = tMOVr $r2, 14 /* CC::al */, $noreg
+  ; CHECK:   $r0 = tMOVr $r12, 14 /* CC::al */, $noreg
+  ; CHECK:   $lr = MVE_DLSTP_16 renamable $r3
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.1(0x7c000000), %bb.2(0x04000000)
+  ; CHECK:   liveins: $lr, $q0, $r0, $r1, $r2, $r3, $r4, $r5, $r6, $r7, $r8, $r10, $r11, $r12
+  ; CHECK:   renamable $r0, renamable $q6 = MVE_VLDRHU16_post killed renamable $r0, 16, 0, $noreg
+  ; CHECK:   renamable $q3 = MVE_VLDRHU16 renamable $r6, 0, 0, $noreg
+  ; CHECK:   renamable $q5 = MVE_VSHR_immu16 killed renamable $q3, 11, 0, $noreg, undef renamable $q5
+  ; CHECK:   MVE_VSTRWU32 killed renamable $q5, $sp, 80, 0, $noreg :: (store (s128) into %stack.0, align 8)
+  ; CHECK:   dead renamable $q7 = MVE_VLDRWU32 $sp, 80, 0, $noreg :: (load (s128) from %stack.0, align 8)
+  ; CHECK:   dead renamable $vpr = MVE_VCMPi16r killed renamable $q6, renamable $r8, 1, 0, killed $noreg
+  ; CHECK:   $lr = MVE_LETP killed renamable $lr, %bb.1
+  bb.0:
+    successors: %bb.1(0x80000000)
+    liveins: $q0, $r1, $r2, $r3, $r4, $r5, $r7, $r8, $r10, $r11, $r12
+
+    $r6 = tMOVr $r2, 14 /* CC::al */, $noreg
+    $r0 = tMOVr $r12, 14 /* CC::al */, $noreg
+    $r9 = tMOVr $r3, 14 /* CC::al */, $noreg
+    renamable $lr = t2DoLoopStartTP renamable $r1, renamable $r3
+
+  bb.1:
+    successors: %bb.1(0x7c000000), %bb.2(0x04000000)
+    liveins: $lr, $q0, $r0, $r1, $r2, $r3, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11, $r12
+
+    renamable $vpr = MVE_VCTP16 renamable $r9, 0, $noreg
+    MVE_VPST 8, implicit $vpr
+    renamable $r0, renamable $q6 = MVE_VLDRHU16_post killed renamable $r0, 16, 1, renamable $vpr
+    renamable $q3 = MVE_VLDRHU16 renamable $r6, 0, 1, renamable $vpr
+    MVE_VPST 2, implicit $vpr
+    renamable $q5 = MVE_VSHR_immu16 renamable $q3, 11, 1, renamable $vpr, undef renamable $q5
+    renamable $r9 = nsw t2SUBri killed renamable $r9, 8, 14 /* CC::al */, $noreg, $noreg
+    MVE_VSTRWU32 killed renamable $q5, $sp, 80, 0, $noreg :: (store (s128) into %stack.0, align 8)
+    MVE_VPST 8, implicit $vpr
+    renamable $q7 = MVE_VLDRWU32 $sp, 80, 0, $noreg :: (load (s128) from %stack.0, align 8)
+    MVE_VPST 1, implicit $vpr
+    renamable $vpr = MVE_VCMPi16r killed renamable $q6, renamable $r8, 1, 1, killed renamable $vpr
+    renamable $lr = t2LoopEndDec killed renamable $lr, %bb.1, implicit-def dead $cpsr
+    tB %bb.2, 14 /* CC::al */, $noreg
+
+  bb.2:
+    successors: %bb.3(0x04000000), %bb.0(0x7c000000)
+    liveins: $q0, $r1, $r2, $r3, $r4, $r5, $r7, $r8, $r10, $r11, $r12
+
+    renamable $r0 = tLDRspi $sp, 1, 14 /* CC::al */, $noreg
+    renamable $r10 = nuw t2ADDri killed renamable $r10, 1, 14 /* CC::al */, $noreg, $noreg
+    renamable $r12 = t2ADDrs killed renamable $r12, killed renamable $r0, 10, 14 /* CC::al */, $noreg, $noreg
+    renamable $r0 = tLDRspi $sp, 3, 14 /* CC::al */, $noreg
+    renamable $r2 = t2ADDrs killed renamable $r2, killed renamable $r0, 10, 14 /* CC::al */, $noreg, $noreg
+    renamable $r0 = tLDRspi $sp, 2, 14 /* CC::al */, $noreg
+    tCMPhir renamable $r10, killed renamable $r0, 14 /* CC::al */, $noreg, implicit-def $cpsr
+    tBcc %bb.0, 1 /* CC::ne */, killed $cpsr
+
+  bb.3:
+    $sp = frame-destroy tADDspi $sp, 24, 14 /* CC::al */, $noreg
+    $sp = frame-destroy VLDMDIA_UPD $sp, 14 /* CC::al */, $noreg, def $d8, def $d9, def $d10, def $d11, def $d12, def $d13, def $d14, def $d15
+    $sp = frame-destroy tADDspi $sp, 1, 14 /* CC::al */, $noreg
+    $sp = frame-destroy t2LDMIA_RET $sp, 14 /* CC::al */, $noreg, def $r4, def $r5, def $r6, def $r7, def $r8, def $r9, def $r10, def $r11, def $pc
+...
+---
+name:            vector_spill_load_outside
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: -120, size: 16, alignment: 8,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: vector_spill_load_outside
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   liveins: $q0, $r1, $r2, $r3, $r4, $r5, $r7, $r8, $r10, $r11, $r12
+  ; CHECK:   $r6 = tMOVr $r2, 14 /* CC::al */, $noreg
+  ; CHECK:   $r0 = tMOVr $r12, 14 /* CC::al */, $noreg
+  ; CHECK:   $r9 = tMOVr $r3, 14 /* CC::al */, $noreg
+  ; CHECK:   $lr = t2DLS renamable $r1
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.1(0x7c000000), %bb.2(0x04000000)
+  ; CHECK:   liveins: $lr, $q0, $r0, $r1, $r2, $r3, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11, $r12
+  ; CHECK:   renamable $vpr = MVE_VCTP16 renamable $r9, 0, $noreg
+  ; CHECK:   MVE_VPST 8, implicit $vpr
+  ; CHECK:   renamable $r0, renamable $q6 = MVE_VLDRHU16_post killed renamable $r0, 16, 1, renamable $vpr
+  ; CHECK:   renamable $q3 = MVE_VLDRHU16 renamable $r6, 0, 1, renamable $vpr
+  ; CHECK:   MVE_VPST 2, implicit $vpr
+  ; CHECK:   renamable $q5 = MVE_VSHR_immu16 killed renamable $q3, 11, 1, renamable $vpr, undef renamable $q5
+  ; CHECK:   renamable $r9 = nsw t2SUBri killed renamable $r9, 8, 14 /* CC::al */, $noreg, $noreg
+  ; CHECK:   MVE_VSTRWU32 killed renamable $q5, $sp, 80, 0, $noreg :: (store (s128) into %stack.0, align 8)
+  ; CHECK:   MVE_VPST 8, implicit $vpr
+  ; CHECK:   dead renamable $q7 = MVE_VLDRWU32 $sp, 80, 0, $noreg :: (load (s128) from %stack.0, align 8)
+  ; CHECK:   MVE_VPST 1, implicit $vpr
+  ; CHECK:   dead renamable $vpr = MVE_VCMPi16r killed renamable $q6, renamable $r8, 1, 1, killed renamable $vpr
+  ; CHECK:   $lr = t2LEUpdate killed renamable $lr, %bb.1
+  bb.0:
+    successors: %bb.1(0x80000000)
+    liveins: $q0, $r1, $r2, $r3, $r4, $r5, $r7, $r8, $r10, $r11, $r12
+
+    $r6 = tMOVr $r2, 14 /* CC::al */, $noreg
+    $r0 = tMOVr $r12, 14 /* CC::al */, $noreg
+    $r9 = tMOVr $r3, 14 /* CC::al */, $noreg
+    renamable $lr = t2DoLoopStartTP renamable $r1, renamable $r3
+
+  bb.1:
+    successors: %bb.1(0x7c000000), %bb.2(0x04000000)
+    liveins: $lr, $q0, $r0, $r1, $r2, $r3, $r4, $r5, $r6, $r7, $r8, $r9, $r10, $r11, $r12
+
+    renamable $vpr = MVE_VCTP16 renamable $r9, 0, $noreg
+    MVE_VPST 8, implicit $vpr
+    renamable $r0, renamable $q6 = MVE_VLDRHU16_post killed renamable $r0, 16, 1, renamable $vpr
+    renamable $q3 = MVE_VLDRHU16 renamable $r6, 0, 1, renamable $vpr
+    MVE_VPST 2, implicit $vpr
+    renamable $q5 = MVE_VSHR_immu16 renamable $q3, 11, 1, renamable $vpr, undef renamable $q5
+    renamable $r9 = nsw t2SUBri killed renamable $r9, 8, 14 /* CC::al */, $noreg, $noreg
+    MVE_VSTRWU32 killed renamable $q5, $sp, 80, 0, $noreg :: (store (s128) into %stack.0, align 8)
+    MVE_VPST 8, implicit $vpr
+    renamable $q7 = MVE_VLDRWU32 $sp, 80, 0, $noreg :: (load (s128) from %stack.0, align 8)
+    MVE_VPST 1, implicit $vpr
+    renamable $vpr = MVE_VCMPi16r killed renamable $q6, renamable $r8, 1, 1, killed renamable $vpr
+    renamable $lr = t2LoopEndDec killed renamable $lr, %bb.1, implicit-def dead $cpsr
+    tB %bb.2, 14 /* CC::al */, $noreg
+
+  bb.2:
+    successors: %bb.3(0x04000000), %bb.0(0x7c000000)
+    liveins: $q0, $r1, $r2, $r3, $r4, $r5, $r7, $r8, $r10, $r11, $r12
+
+    renamable $q7 = MVE_VLDRWU32 $sp, 80, 0, $noreg :: (load (s128) from %stack.0, align 8)
+    renamable $r0 = tLDRspi $sp, 1, 14 /* CC::al */, $noreg
+    renamable $r10 = nuw t2ADDri killed renamable $r10, 1, 14 /* CC::al */, $noreg, $noreg
+    renamable $r12 = t2ADDrs killed renamable $r12, killed renamable $r0, 10, 14 /* CC::al */, $noreg, $noreg
+    renamable $r0 = tLDRspi $sp, 3, 14 /* CC::al */, $noreg
+    renamable $r2 = t2ADDrs killed renamable $r2, killed renamable $r0, 10, 14 /* CC::al */, $noreg, $noreg
+    renamable $r0 = tLDRspi $sp, 2, 14 /* CC::al */, $noreg
+    tCMPhir renamable $r10, killed renamable $r0, 14 /* CC::al */, $noreg, implicit-def $cpsr
+    tBcc %bb.0, 1 /* CC::ne */, killed $cpsr
+
+  bb.3:
+    $sp = frame-destroy tADDspi $sp, 24, 14 /* CC::al */, $noreg
+    $sp = frame-destroy VLDMDIA_UPD $sp, 14 /* CC::al */, $noreg, def $d8, def $d9, def $d10, def $d11, def $d12, def $d13, def $d14, def $d15
+    $sp = frame-destroy tADDspi $sp, 1, 14 /* CC::al */, $noreg
+    $sp = frame-destroy t2LDMIA_RET $sp, 14 /* CC::al */, $noreg, def $r4, def $r5, def $r6, def $r7, def $r8, def $r9, def $r10, def $r11, def $pc
+
+...