[llvm] [LoopFusion] Simplify the logic of checking trip count equality (PR #201446)

Ehsan Amiri via llvm-commits llvm-commits at lists.llvm.org
Thu Jun 4 08:19:28 PDT 2026


https://github.com/amehsan updated https://github.com/llvm/llvm-project/pull/201446

>From f2f6b9c20216bd7b8874bd2f99a856b65b1ca00c Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Tue, 26 May 2026 21:01:57 -0400
Subject: [PATCH 1/4] [LoopFusion] Simplify the logic of checking trip count
 equality (NFCI)

---
 llvm/lib/Transforms/Scalar/LoopFuse.cpp       | 76 ++++++-------------
 .../test/Transforms/LoopFusion/cannot_fuse.ll |  3 +-
 2 files changed, 25 insertions(+), 54 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
index bce9a04ddd0b2..df91ab2191c97 100644
--- a/llvm/lib/Transforms/Scalar/LoopFuse.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
@@ -590,29 +590,26 @@ struct LoopFuser {
     return true;
   }
 
-  /// Determine if two fusion candidates have the same trip count (i.e., they
-  /// execute the same number of iterations).
+  /// Computes the integer difference in trip counts:
+  /// TripCount(FC0) - TripCount(FC1).
   ///
-  /// This function will return a pair of values. The first is a boolean,
-  /// stating whether or not the two candidates are known at compile time to
-  /// have the same TripCount. The second is the difference in the two
-  /// TripCounts. This information can be used later to determine whether or not
-  /// peeling can be performed on either one of the candidates.
-  std::pair<bool, std::optional<unsigned>>
-  haveIdenticalTripCounts(const FusionCandidate &FC0,
-                          const FusionCandidate &FC1) const {
+  /// \returns The integer difference, or std::nullopt if it
+  ///          cannot be determined.
+  std::optional<int>
+  calculateTripCountDiff(const FusionCandidate &FC0,
+                         const FusionCandidate &FC1) const {
     const SCEV *TripCount0 = SE.getBackedgeTakenCount(FC0.L);
     if (isa<SCEVCouldNotCompute>(TripCount0)) {
       UncomputableTripCount++;
       LLVM_DEBUG(dbgs() << "Trip count of first loop could not be computed!");
-      return {false, std::nullopt};
+      return std::nullopt;
     }
 
     const SCEV *TripCount1 = SE.getBackedgeTakenCount(FC1.L);
     if (isa<SCEVCouldNotCompute>(TripCount1)) {
       UncomputableTripCount++;
       LLVM_DEBUG(dbgs() << "Trip count of second loop could not be computed!");
-      return {false, std::nullopt};
+      return std::nullopt;
     }
 
     LLVM_DEBUG(dbgs() << "\tTrip counts: " << *TripCount0 << " & "
@@ -621,7 +618,7 @@ struct LoopFuser {
                       << "\n");
 
     if (TripCount0 == TripCount1)
-      return {true, 0};
+      return 0;
 
     LLVM_DEBUG(dbgs() << "The loops do not have the same tripcount, "
                          "determining the difference between trip counts\n");
@@ -637,24 +634,10 @@ struct LoopFuser {
       LLVM_DEBUG(dbgs() << "Loop(s) do not have a single exit point or do not "
                            "have a constant number of iterations. Peeling "
                            "is not benefical\n");
-      return {false, std::nullopt};
-    }
-
-    std::optional<unsigned> Difference;
-    int Diff = TC0 - TC1;
-
-    if (Diff > 0)
-      Difference = Diff;
-    else {
-      LLVM_DEBUG(
-          dbgs() << "Difference is less than 0. FC1 (second loop) has more "
-                    "iterations than the first one. Currently not supported\n");
+      return std::nullopt;
     }
 
-    LLVM_DEBUG(dbgs() << "Difference in loop trip count is: " << Difference
-                      << "\n");
-
-    return {false, Difference};
+    return TC0 - TC1;
   }
 
   void peelFusionCandidate(FusionCandidate &FC0, const FusionCandidate &FC1,
@@ -670,9 +653,9 @@ struct LoopFuser {
     LLVM_DEBUG(dbgs() << "Done Peeling\n");
 
 #ifndef NDEBUG
-    auto IdenticalTripCount = haveIdenticalTripCounts(FC0, FC1);
+    auto TCDiff = calculateTripCountDiff(FC0, FC1);
 
-    assert(IdenticalTripCount.first && *IdenticalTripCount.second == 0 &&
+    assert(TCDiff == 0 &&
            "Loops should have identical trip counts after peeling");
 #endif
 
@@ -755,34 +738,21 @@ struct LoopFuser {
         FC0.verify();
         FC1.verify();
 
-        // Check if the candidates have identical tripcounts (first value of
-        // pair), and if not check the difference in the tripcounts between
-        // the loops (second value of pair). The difference is not equal to
-        // std::nullopt iff the loops iterate a constant number of times, and
-        // have a single exit.
-        std::pair<bool, std::optional<unsigned>> IdenticalTripCountRes =
-            haveIdenticalTripCounts(FC0, FC1);
-        bool SameTripCount = IdenticalTripCountRes.first;
-        std::optional<unsigned> TCDifference = IdenticalTripCountRes.second;
-
+        std::optional<int> TCDifference = calculateTripCountDiff(FC0, FC1);
         // Here we are checking that FC0 (the first loop) can be peeled, and
-        // both loops have different tripcounts.
-        if (FC0.AbleToPeel && !SameTripCount && TCDifference) {
-          if (*TCDifference > FusionPeelMaxCount) {
-            LLVM_DEBUG(dbgs()
-                       << "Difference in loop trip counts: " << *TCDifference
-                       << " is greater than maximum peel count specificed: "
-                       << FusionPeelMaxCount << "\n");
-          } else {
+        // the first loop has a larger trip count.
+        bool WillPeel = false;
+        if (FC0.AbleToPeel && TCDifference > 0 &&
+            TCDifference <= FusionPeelMaxCount) {
             // Dependent on peeling being performed on the first loop, and
             // assuming all other conditions for fusion return true.
-            SameTripCount = true;
+            WillPeel = true;
           }
-        }
 
-        if (!SameTripCount) {
+        if (!WillPeel && TCDifference != 0) {
           LLVM_DEBUG(dbgs() << "Fusion candidates do not have identical trip "
-                               "counts. Not fusing.\n");
+                               "counts and peeling is not supported for this "
+                               "case. Not fusing.\n");
           reportLoopFusion<OptimizationRemarkMissed>(FC0, FC1,
                                                      NonEqualTripCount);
           continue;
diff --git a/llvm/test/Transforms/LoopFusion/cannot_fuse.ll b/llvm/test/Transforms/LoopFusion/cannot_fuse.ll
index 0eea73940904d..619b5a1ddb04d 100644
--- a/llvm/test/Transforms/LoopFusion/cannot_fuse.ll
+++ b/llvm/test/Transforms/LoopFusion/cannot_fuse.ll
@@ -148,7 +148,8 @@ bb25:                                             ; preds = %bb15
 ; CHECK: Attempting fusion on Candidate List:
 ; CHECK-NEXT: [[LOOP1PREHEADER]]
 ; CHECK-NEXT: [[LOOP2PREHEADER]]
-; CHECK: Fusion candidates do not have identical trip counts. Not fusing.
+; CHECK: Fusion candidates do not have identical trip counts
+; CHECK-SAME: and peeling is not supported for this
 ; CHECK: Loop Fusion complete
 define void @different_bounds(ptr noalias %arg) {
 bb:

>From 4e8b74ba1297e06bb8186ee35d99eef79522653f Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Wed, 3 Jun 2026 22:30:08 -0400
Subject: [PATCH 2/4] addressing a couple of issues

---
 llvm/lib/Transforms/Scalar/LoopFuse.cpp | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
index df91ab2191c97..b6207d111242d 100644
--- a/llvm/lib/Transforms/Scalar/LoopFuse.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
@@ -96,7 +96,7 @@ STATISTIC(NumHoistedInsts, "Number of hoisted preheader instructions.");
 STATISTIC(NumSunkInsts, "Number of hoisted preheader instructions.");
 STATISTIC(NumDA, "DA checks passed");
 
-static cl::opt<unsigned> FusionPeelMaxCount(
+static cl::opt<uint32_t> FusionPeelMaxCount(
     "loop-fusion-peel-max-count", cl::init(0), cl::Hidden,
     cl::desc("Max number of iterations to be peeled from a loop, such that "
              "fusion can take place"));
@@ -595,7 +595,7 @@ struct LoopFuser {
   ///
   /// \returns The integer difference, or std::nullopt if it
   ///          cannot be determined.
-  std::optional<int>
+  std::optional<int64_t>
   calculateTripCountDiff(const FusionCandidate &FC0,
                          const FusionCandidate &FC1) const {
     const SCEV *TripCount0 = SE.getBackedgeTakenCount(FC0.L);
@@ -624,9 +624,11 @@ struct LoopFuser {
                          "determining the difference between trip counts\n");
 
     // Currently only considering loops with a single exit point
-    // and a non-constant trip count.
-    const unsigned TC0 = SE.getSmallConstantTripCount(FC0.L);
-    const unsigned TC1 = SE.getSmallConstantTripCount(FC1.L);
+    // and a non-constant trip count. Note that the return value 
+    // of getSmallConstantTripCount is a 32 bit number, based on
+    // the existing implementation.
+    const int64_t TC0 = static_cast<int64_t>(SE.getSmallConstantTripCount(FC0.L));
+    const int64_t TC1 = static_cast<int64_t>(SE.getSmallConstantTripCount(FC1.L));
 
     // If any of the tripcounts are zero that means that loop(s) do not have
     // a single exit or a constant tripcount.
@@ -738,18 +740,18 @@ struct LoopFuser {
         FC0.verify();
         FC1.verify();
 
-        std::optional<int> TCDifference = calculateTripCountDiff(FC0, FC1);
+        std::optional<int64_t> TCDifference = calculateTripCountDiff(FC0, FC1);
         // Here we are checking that FC0 (the first loop) can be peeled, and
         // the first loop has a larger trip count.
         bool WillPeel = false;
         if (FC0.AbleToPeel && TCDifference > 0 &&
-            TCDifference <= FusionPeelMaxCount) {
+            TCDifference <= static_cast<int64_t>(FusionPeelMaxCount)) {
             // Dependent on peeling being performed on the first loop, and
             // assuming all other conditions for fusion return true.
             WillPeel = true;
           }
 
-        if (!WillPeel && TCDifference != 0) {
+        if (!WillPeel && (!TCDifference || *TCDifference != 0)) {
           LLVM_DEBUG(dbgs() << "Fusion candidates do not have identical trip "
                                "counts and peeling is not supported for this "
                                "case. Not fusing.\n");

>From 688dc95fb1940a57fc36b543ca35407982d63b6d Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Wed, 3 Jun 2026 23:33:56 -0400
Subject: [PATCH 3/4] address formatting issues

---
 llvm/lib/Transforms/Scalar/LoopFuse.cpp | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
index b6207d111242d..03bfcf2c8eb0f 100644
--- a/llvm/lib/Transforms/Scalar/LoopFuse.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
@@ -624,11 +624,13 @@ struct LoopFuser {
                          "determining the difference between trip counts\n");
 
     // Currently only considering loops with a single exit point
-    // and a non-constant trip count. Note that the return value 
+    // and a non-constant trip count. Note that the return value
     // of getSmallConstantTripCount is a 32 bit number, based on
     // the existing implementation.
-    const int64_t TC0 = static_cast<int64_t>(SE.getSmallConstantTripCount(FC0.L));
-    const int64_t TC1 = static_cast<int64_t>(SE.getSmallConstantTripCount(FC1.L));
+    const int64_t TC0 =
+        static_cast<int64_t>(SE.getSmallConstantTripCount(FC0.L));
+    const int64_t TC1 =
+        static_cast<int64_t>(SE.getSmallConstantTripCount(FC1.L));
 
     // If any of the tripcounts are zero that means that loop(s) do not have
     // a single exit or a constant tripcount.
@@ -746,10 +748,10 @@ struct LoopFuser {
         bool WillPeel = false;
         if (FC0.AbleToPeel && TCDifference > 0 &&
             TCDifference <= static_cast<int64_t>(FusionPeelMaxCount)) {
-            // Dependent on peeling being performed on the first loop, and
-            // assuming all other conditions for fusion return true.
-            WillPeel = true;
-          }
+          // Dependent on peeling being performed on the first loop, and
+          // assuming all other conditions for fusion return true.
+          WillPeel = true;
+        }
 
         if (!WillPeel && (!TCDifference || *TCDifference != 0)) {
           LLVM_DEBUG(dbgs() << "Fusion candidates do not have identical trip "

>From 4d5f9f90b5e83ebd1997fe548e4f72ba290c0763 Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Thu, 4 Jun 2026 11:16:49 -0400
Subject: [PATCH 4/4] consistent pattern of using optional

---
 llvm/lib/Transforms/Scalar/LoopFuse.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopFuse.cpp b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
index 03bfcf2c8eb0f..2279d75699088 100644
--- a/llvm/lib/Transforms/Scalar/LoopFuse.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopFuse.cpp
@@ -659,7 +659,7 @@ struct LoopFuser {
 #ifndef NDEBUG
     auto TCDiff = calculateTripCountDiff(FC0, FC1);
 
-    assert(TCDiff == 0 &&
+    assert(TCDiff && *TCDiff == 0 &&
            "Loops should have identical trip counts after peeling");
 #endif
 
@@ -744,12 +744,12 @@ struct LoopFuser {
 
         std::optional<int64_t> TCDifference = calculateTripCountDiff(FC0, FC1);
         // Here we are checking that FC0 (the first loop) can be peeled, and
-        // the first loop has a larger trip count.
+        // the first loop has a larger trip count. In this case it is possible
+        // that the first loop is peeled to expose the fusion opportunity.
+        // Peeling the second loop is not currently supported.
         bool WillPeel = false;
-        if (FC0.AbleToPeel && TCDifference > 0 &&
-            TCDifference <= static_cast<int64_t>(FusionPeelMaxCount)) {
-          // Dependent on peeling being performed on the first loop, and
-          // assuming all other conditions for fusion return true.
+        if (FC0.AbleToPeel && TCDifference && *TCDifference > 0 &&
+            *TCDifference <= static_cast<int64_t>(FusionPeelMaxCount)) {
           WillPeel = true;
         }
 



More information about the llvm-commits mailing list