[HardwareLoops] Loop guard intrinsic to recognise zext

[sherwin-dc]

Summary

I am working on a downstream custom backend that uses the loop guard intrinsic.

If a loop count was initially represented by a 32b unsigned int in C then the hardware-loop pass can recognise the loop guard and insert the llvm.test.set.loop.iterations intrinsic. However if this was instead a unsigned short/char then clang inserts a zext instruction to expand the loop count to an i32.

This means the loop guard intrinsic would not be used as shown here because %x was being compared with 0 in the loop guard and the hardware loop pass was expecting %0 to be compared.

entry:
  %tobool.not5 = icmp eq i8 %x, 0
  %0 = zext i8 %x to i32
  br i1 %tobool.not5, label %while.end, label %while.body.preheader

while.body.preheader:                             ; preds = %entry
  %1 = zext i8 %x to i32
  %scevgep = getelementptr i32, i32* %arr, i32 %1
  call void @llvm.set.loop.iterations.i32(i32 %0)
  br label %while.body

This patch would benefit other backends as well that make use of the loop guard intrinsic. For example, if we have a loop:

void loop(unsigned *a, unsigned *b, unsigned length)
{
    for(int i = 0; i < length; i++)
    {
        a[i] = b[i];
    }
}

and compile it to IR we get int_loop.ll. After changing unsigned length to unsigned short length and compiling to IR we have short_loop.ll.

int_loop.ll800 BDownload

short_loop.ll841 BDownload

Using ARM as an example (since it makes use of the hardware loop guard intrinsic), compiling as llc -mtriple=thumbv8.1m.main-none-none-eabi -mattr=+lob,+mve.fp int_loop.ll -o int_loop.asm we get int_loop.asm and short_loop.asm respectively.

int_loop.asm1 KBDownload

short_loop.asm1 KBDownload

The int_loop.ll example compiles to a hardware loop with ARM's WLS instruction, while short_loop.ll does not and includes a loop guard. With this patch the short_loop.ll example also compiles to a WLS loop since the loop guard intrinsic is used.

[dmgreen]

The int_loop.ll example compiles to a hardware loop with ARM's WLS instruction, while short_loop.ll does not and includes a loop guard. With this patch the short_loop.ll example also compiles to a WLS loop since the loop guard intrinsic is used.

Sounds good, Can you add a test then?

[samparker]

llvm/lib/CodeGen/HardwareLoops.cpp

369

Can we use SCEV isLoopEntryGuardedByCond instead of doing the pattern matching? It feels like we should be able to replace IsCompareZero with it which would remove the need to do an explicit zext check.

[sherwin-dc]

llvm/lib/CodeGen/HardwareLoops.cpp

369

It seems that isLoopEntryGuardedByCond is already used just before this (see below comment). This was added at the same time as IsCompareZero in D63809, so I'm not sure if IsCompareZero was needed in some way to explicitly check the instructions.

395–396

isLoopEntryGuardedByCond used here

[sherwin-dc]

• Add loop guard test

[samparker]

llvm/lib/CodeGen/HardwareLoops.cpp

369

Ah, yes... those unfortunate cases where pattern matching is more useful than SCEV.

[sherwin-dc]

@dmgreen would this test case suffice or should I add a ARM test that compiles to a WLS loop?

[dmgreen]

Yeah thanks. The test looks good. @samparker - are you happy with this?

[samparker]

are you happy with this?

Yep!
Thanks @sherwin-dc

[sherwin-dc]

Thanks for reviewing. I dont have commit access, so I'll need help with landing the change.

[samparker]

Ok, will do.

llvm/lib/CodeGen/HardwareLoops.cpp

if (!ICmp->isEquality())
     return false;
 
   auto IsCompareZero = [](ICmpInst *ICmp, Value *Count, unsigned OpIdx) {
     if (auto *Const = dyn_cast<ConstantInt>(ICmp->getOperand(OpIdx)))
       return Const->isZero() && ICmp->getOperand(OpIdx ^ 1) == Count;
     return false;
   };
 
-  if (!IsCompareZero(ICmp, Count, 0) && !IsCompareZero(ICmp, Count, 1))
+  // Check if Count is a zext.
+  Value *CountBefZext =
+      isa<ZExtInst>(Count) ? cast<ZExtInst>(Count)->getOperand(0) : nullptr;
+
+  if (!IsCompareZero(ICmp, Count, 0) && !IsCompareZero(ICmp, Count, 1) &&
+      !IsCompareZero(ICmp, CountBefZext, 0) &&
+      !IsCompareZero(ICmp, CountBefZext, 1))
     return false;
 
   unsigned SuccIdx = ICmp->getPredicate() == ICmpInst::ICMP_NE ? 0 : 1;
   if (BI->getSuccessor(SuccIdx) != Preheader)
     return false;
 
   return true;
 }
 
 Value *HardwareLoop::InitLoopCount() {
   LLVM_DEBUG(dbgs() << "HWLoops: Initialising loop counter value:\n");
   // Can we replace a conditional branch with an intrinsic that sets the
   // loop counter and tests that is not zero?
 
   SCEVExpander SCEVE(SE, DL, "loopcnt");
 
   // If we're trying to use the 'test and set' form of the intrinsic, we need
   // to replace a conditional branch that is controlling entry to the loop. It
   // is likely (guaranteed?) that the preheader has an unconditional branch to
   // the loop header, so also check if it has a single predecessor.
   if (SE.isLoopEntryGuardedByCond(L, ICmpInst::ICMP_NE, TripCount,
                                   SE.getZero(TripCount->getType()))) {
     LLVM_DEBUG(dbgs() << " - Attempting to use test.set counter.\n");
     UseLoopGuard |= ForceGuardLoopEntry;
   } else
     UseLoopGuard = false;
 
   BasicBlock *BB = L->getLoopPreheader();
   if (UseLoopGuard && BB->getSinglePredecessor() &&
       cast<BranchInst>(BB->getTerminator())->isUnconditional()) {