Address review comments

Author: david-arm

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

@@ -69,7 +69,7 @@ static cl::opt<bool> EnableLSRCostOpt("enable-aarch64-lsr-cost-opt",
                                       cl::init(true), cl::Hidden);
 
 static cl::opt<unsigned> SmallMultiExitLoopUF(
-    "small-multi-exit-loop-unroll-factor", cl::init(0), cl::Hidden,
+    "aarch64-small-multi-exit-loop-unroll", cl::init(0), cl::Hidden,
     cl::desc(
         "Force unrolling of small multi-exit loops with given unroll factor"));
 
@@ -4386,35 +4386,45 @@ static bool shouldUnrollLoopWithInstruction(Instruction &I,
   if (I.getType()->isVectorTy())
     return false;
 
-  if (isa<CallInst>(I) || isa<InvokeInst>(I)) {
-    if (const Function *F = cast<CallBase>(I).getCalledFunction())
-      if (!TTI.isLoweredToCall(F))
-        return true;
+  if (isa<CallBase>(I)) {
+    if (isa<CallInst>(I) || isa<InvokeInst>(I))
+      if (const Function *F = cast<CallBase>(I).getCalledFunction())
+        if (!TTI.isLoweredToCall(F))
+          return true;
     return false;
   }
 
   return true;
 }
 
-static InstructionCost getSizeOfLoop(Loop *L, AArch64TTIImpl &TTI) {
+static unsigned getLoopSize(Loop *L, AArch64TTIImpl &TTI,
+                            InstructionCost Budget) {
   // Estimate the size of the loop.
-  InstructionCost Size = 0;
+  InstructionCost LoopCost = 0;
+
+  if (findStringMetadataForLoop(L, "llvm.loop.isvectorized"))
+    return 0;
+
   for (auto *BB : L->getBlocks()) {
     for (auto &I : *BB) {
       if (!shouldUnrollLoopWithInstruction(I, TTI))
-        return InstructionCost::getInvalid();
+        return 0;
 
       SmallVector<const Value *, 4> Operands(I.operand_values());
       InstructionCost Cost =
           TTI.getInstructionCost(&I, Operands, TTI::TCK_CodeSize);
       // This can happen with intrinsics that don't currently have a cost model
       // or for some operations that require SVE.
       if (!Cost.isValid())
-        return InstructionCost::getInvalid();
-      Size += *Cost.getValue();
+        return 0;
+
+      LoopCost += Cost;
+      if (LoopCost > Budget)
+        return 0;
     }
   }
-  return Size;
+
+  return *LoopCost.getValue();
 }
 
 static bool shouldUnrollMultiExitLoop(Loop *L, ScalarEvolution &SE,
@@ -4432,12 +4442,8 @@ static bool shouldUnrollMultiExitLoop(Loop *L, ScalarEvolution &SE,
   if (MaxTC > 0 && MaxTC <= 32)
     return false;
 
-  if (findStringMetadataForLoop(L, "llvm.loop.isvectorized"))
-    return false;
-
   // Estimate the size of the loop.
-  InstructionCost Size = getSizeOfLoop(L, TTI);
-  if (!Size.isValid())
+  if (!getLoopSize(L, TTI, 5))
     return false;
 
   // Small search loops with multiple exits can be highly beneficial to unroll.
@@ -4452,7 +4458,7 @@ static bool shouldUnrollMultiExitLoop(Loop *L, ScalarEvolution &SE,
       }))
     return false;
 
-  return *Size.getValue() < 6;
+  return true;
 }
 
 /// For Apple CPUs, we want to runtime-unroll loops to make better use if the
@@ -4469,28 +4475,15 @@ getAppleRuntimeUnrollPreferences(Loop *L, ScalarEvolution &SE,
   if (!L->isInnermost() || L->getNumBlocks() > 8)
     return;
 
+  // This is handled by common code.
+  if (!L->getExitBlock())
+    return;
+
   const SCEV *BTC = SE.getSymbolicMaxBackedgeTakenCount(L);
   if (isa<SCEVConstant>(BTC) || isa<SCEVCouldNotCompute>(BTC) ||
       (SE.getSmallConstantMaxTripCount(L) > 0 &&
        SE.getSmallConstantMaxTripCount(L) <= 32))
     return;
-  if (findStringMetadataForLoop(L, "llvm.loop.isvectorized"))
-    return;
-
-  int64_t Size = 0;
-  for (auto *BB : L->getBlocks()) {
-    for (auto &I : *BB) {
-      if (!isa<IntrinsicInst>(&I) && isa<CallBase>(&I))
-        return;
-      SmallVector<const Value *, 4> Operands(I.operand_values());
-      Size +=
-          *TTI.getInstructionCost(&I, Operands, TTI::TCK_CodeSize).getValue();
-    }
-  }
-
-  // This is handled by common code.
-  if (!L->getExitBlock())
-    return;
 
   if (SE.getSymbolicMaxBackedgeTakenCount(L) != SE.getBackedgeTakenCount(L))
     return;
@@ -4502,7 +4495,9 @@ getAppleRuntimeUnrollPreferences(Loop *L, ScalarEvolution &SE,
   // dependencies, to expose more parallel memory access streams.
   BasicBlock *Header = L->getHeader();
   if (Header == L->getLoopLatch()) {
-    if (Size > 8)
+    // Estimate the size of the loop.
+    unsigned Size = getLoopSize(L, TTI, 8);
+    if (!Size)
       return;
 
     SmallPtrSet<Value *, 8> LoadedValues;

llvm/test/Transforms/LoopUnroll/AArch64/apple-unrolling-multi-exit.ll

@@ -3,7 +3,7 @@
 ; RUN: opt -p loop-unroll -mcpu=apple-m2 -S %s | FileCheck --check-prefix=APPLE %s
 ; RUN: opt -p loop-unroll -mcpu=apple-m3 -S %s | FileCheck --check-prefix=APPLE %s
 ; RUN: opt -p loop-unroll -mcpu=apple-m4 -S %s | FileCheck --check-prefix=APPLE %s
-; RUN: opt -p loop-unroll -mcpu=apple-m4 -small-multi-exit-loop-unroll-factor=2 -S %s | FileCheck --check-prefix=UNROLL2 %s
+; RUN: opt -p loop-unroll -mcpu=apple-m4 -aarch64-small-multi-exit-loop-unroll=2 -S %s | FileCheck --check-prefix=UNROLL2 %s
 ; RUN: opt -p loop-unroll -mcpu=cortex-a57 -S %s | FileCheck --check-prefix=OTHER %s
 
 target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-n32:64-S128-Fn32"

llvm/test/Transforms/LoopUnroll/AArch64/unrolling-multi-exit.ll

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -p loop-unroll -mcpu=generic -small-multi-exit-loop-unroll-factor=2 -S %s | FileCheck --check-prefixes=COMMON,UNROLL2 %s
+; RUN: opt -p loop-unroll -mcpu=generic -aarch64-small-multi-exit-loop-unroll=2 -S %s | FileCheck --check-prefixes=COMMON,UNROLL2 %s
 ; RUN: opt -p loop-unroll -mcpu=generic -S %s | FileCheck --check-prefixes=COMMON,GENERIC %s
 
 target triple = "aarch64-linux-gnu"