[LoopUnroll] Unroll small loops with 'pragma unroll' when its trip count is destroyed by preious optimization

Author: xiangzh1

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

@@ -755,7 +755,7 @@ static unsigned getFullUnrollBoostingFactor(const EstimatedUnrollCost &Cost,
 static std::optional<unsigned>
 shouldPragmaUnroll(Loop *L, const PragmaInfo &PInfo,
                    const unsigned TripMultiple, const unsigned TripCount,
-                   const UnrollCostEstimator UCE,
+                   unsigned MaxTripCount, const UnrollCostEstimator UCE,
                    const TargetTransformInfo::UnrollingPreferences &UP) {
 
   // Using unroll pragma
@@ -776,6 +776,10 @@ shouldPragmaUnroll(Loop *L, const PragmaInfo &PInfo,
   if (PInfo.PragmaFullUnroll && TripCount != 0)
     return TripCount;
 
+  if (PInfo.PragmaEnableUnroll && !TripCount && MaxTripCount &&
+      MaxTripCount <= UnrollMaxUpperBound)
+    return MaxTripCount;
+
   // if didn't return until here, should continue to other priorties
   return std::nullopt;
 }
@@ -902,7 +906,7 @@ bool llvm::computeUnrollCount(
   // 1st priority is unroll count set by "unroll-count" option.
   // 2nd priority is unroll count set by pragma.
   if (auto UnrollFactor = shouldPragmaUnroll(L, PInfo, TripMultiple, TripCount,
-                                             UCE, UP)) {
+                                             MaxTripCount, UCE, UP)) {
     UP.Count = *UnrollFactor;
 
     if (UserUnrollCount || (PragmaCount > 0)) {