[llvm] [AArch64] Enable unrolling for small multi-exit loops (PR #131998)

[David Sherwood via llvm-commits]

================
@@ -4528,6 +4528,95 @@ getFalkorUnrollingPreferences(Loop *L, ScalarEvolution &SE,
   }
 }
 
+static bool shouldUnrollLoopWithInstruction(Instruction &I,
+                                            AArch64TTIImpl &TTI) {
+  // Don't unroll vectorised loop.
+  if (I.getType()->isVectorTy())
+    return false;
+
+  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;
+}
+
+// This function returns true if the loop:
+//  1. Contains only those instructions that should be unrolled,
+//  2. Has a valid cost,
+//  3. Has a cost within the supplied budget.
+// Otherwise it returns false.
+static bool isLoopSizeWithinBudget(Loop *L, AArch64TTIImpl &TTI,
+                                   InstructionCost Budget,
+                                   unsigned *FinalSize) {
+  // Estimate the size of the loop.
+  InstructionCost LoopCost = 0;
+
+  if (findStringMetadataForLoop(L, "llvm.loop.isvectorized"))
+    return false;
+
+  for (auto *BB : L->getBlocks()) {
+    for (auto &I : *BB) {
+      if (!shouldUnrollLoopWithInstruction(I, TTI))
----------------
david-arm wrote:

I've kept it as part of the loop, because even the previous loop checking the size for the apple preferences bailed out on call instructions. It seems a bit inefficient to walk the loop twice, so I've renamed it to `canUnrollLoopWithinBudget`. I've also move the `findStringMetadataForLoop` out into `shouldUnrollMultiExitLoop`, since there's no need to check here.

https://github.com/llvm/llvm-project/pull/131998