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

[Florian Hahn 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))
----------------
fhahn wrote:

Would it make more sense to check this separately? If not,  would be good to update the name to make it clearer that this does not only check if the loop size is within a budget?

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