[llvm] [TailRecElim] Introduce support for shift accumulator optimization (PR #181331)

Federico Bruzzone via llvm-commits llvm-commits at lists.llvm.org
Tue Apr 14 02:59:20 PDT 2026


https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/181331

>From adc8ae38e4d57d83c5d0f5aeedd58040644426aa Mon Sep 17 00:00:00 2001
From: FedericoBruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 17 Feb 2026 20:11:35 +0100
Subject: [PATCH 1/5] [TailRecElim] Add support for shift accumulator
 optimization

Signed-off-by: FedericoBruzzone <federico.bruzzone.i at gmail.com>
---
 .../Scalar/TailRecursionElimination.cpp       | 193 +++++++++++++++---
 .../TailCallElim/shl-accumulator-opt.ll       | 118 +++++++++++
 2 files changed, 287 insertions(+), 24 deletions(-)
 create mode 100644 llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index 3f7eca1fc7513..888d0b223463b 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -82,6 +82,7 @@
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <cmath>
+#include <variant>
 using namespace llvm;
 
 #define DEBUG_TYPE "tailcallelim"
@@ -374,28 +375,136 @@ static bool canMoveAboveCall(Instruction *I, CallInst *CI, AliasAnalysis *AA) {
   return !is_contained(I->operands(), CI);
 }
 
-static bool canTransformAccumulatorRecursion(Instruction *I, CallInst *CI) {
-  if (!I->isAssociative() || !I->isCommutative())
+// While shifts are neither associative nor commutative, a chain of shifts by a
+// constant amount C is equivalent to a single shift by the sum of the amounts:
+//     ... (Base << C) << C) ... << C == Base << (C * Iterations)
+// This relation applies to left shifts as well as arithmetic/logical right
+// shifts when the shift amount is a constant.
+static bool isPseudoAssociative(Instruction *I) {
+  switch (I->getOpcode()) {
+  case Instruction::Shl:
+  case Instruction::AShr:
+  case Instruction::LShr:
+    break;
+  default:
+    return false;
+  }
+
+  ConstantInt *CI = dyn_cast<ConstantInt>(I->getOperand(1));
+  if (!CI)
+    return false;
+
+  return true;
+}
+
+// Find the base-case return value for function F: examine all
+// return instructions and pick return values that do not depend on a
+// recursive call to F. If there is exactly one distinct such value,
+// return it. If there are none or more than one distinct value, return
+// nullptr to indicate failure.
+//
+// FIXME: There is a room for improvement here in the future, e.g., consider
+// non-constant values and multiple base cases -- e.g., we want to be able to
+// handle code like:
+// ```
+// int f(int x) {
+//  if (x == 1) return 1;
+//  if (x == 10) return 10;
+//  return f(x-1) << 1;
+// }
+// ```
+static Constant *getReturnValue(Function &F) {
+  Constant *BaseCaseVal = nullptr;
+
+  auto ValueUsesRecursiveCall = [&](Value *V) {
+    SmallVector<Value *, 8> Worklist;
+    SmallPtrSet<Value *, 8> Visited;
+    Worklist.push_back(V);
+    while (!Worklist.empty()) {
+      Value *Cur = Worklist.pop_back_val();
+      if (!Visited.insert(Cur).second)
+        continue;
+      if (Instruction *I = dyn_cast<Instruction>(Cur)) {
+        if (CallInst *CI = dyn_cast<CallInst>(I))
+          if (CI->getCalledFunction() == &F)
+            return true;
+        for (Use &U : I->operands())
+          Worklist.push_back(U.get());
+      } else if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Cur)) {
+        for (unsigned I = 0, E = CE->getNumOperands(); I != E; ++I)
+          Worklist.push_back(CE->getOperand(I));
+      }
+    }
+    return false;
+  };
+
+  for (BasicBlock &BB : F) {
+    if (ReturnInst *RI = dyn_cast<ReturnInst>(BB.getTerminator())) {
+      Value *RV = RI->getReturnValue();
+      if (!RV)
+        continue;
+      // It isn't so trivial to perform TRE if the return value depends on the
+      // result of a recursive call, because the return value will be different
+      // for different iterations of the recursion. So we ignore return values
+      // that depend on recursive calls.
+      if (ValueUsesRecursiveCall(RV))
+        continue;
+      if (!BaseCaseVal && isa<Constant>(RV))
+        BaseCaseVal = cast<Constant>(RV);
+      else if (BaseCaseVal != RV)
+        return nullptr;
+    }
+  }
+
+  return BaseCaseVal;
+}
+
+// This function checks whether the instruction I can be used
+// to perform accumulator recursion elimination for the
+// call instruction CI.
+// In the presence of pseudo-associative operations, it returns
+// the base case constant value to both indicate success and
+// provide the value needed to initialize the accumulator.
+static std::variant<bool, Constant *>
+canTransformAccumulatorRecursion(Instruction *I, CallInst *CI) {
+  auto IsPseudoAssociative = isPseudoAssociative(I);
+  if ((!I->isAssociative() || !I->isCommutative()) && !IsPseudoAssociative)
     return false;
 
   assert(I->getNumOperands() >= 2 &&
          "Associative/commutative operations should have at least 2 args!");
 
-  if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
-    // Accumulators must have an identity.
-    if (!ConstantExpr::getIntrinsicIdentity(II->getIntrinsicID(), I->getType()))
+  Constant *BaseCaseConst = nullptr;
+  if (IsPseudoAssociative) {
+    // For pseudo-associative operations, we require that the recursive call
+    // is always on the first operand.
+    if (I->getOperand(0) != CI)
       return false;
-  }
 
-  // Exactly one operand should be the result of the call instruction.
-  if ((I->getOperand(0) == CI && I->getOperand(1) == CI) ||
-      (I->getOperand(0) != CI && I->getOperand(1) != CI))
-    return false;
+    BaseCaseConst = getReturnValue(*CI->getCalledFunction());
+    if (!BaseCaseConst)
+      return false;
+  } else {
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
+      // Accumulators must have an identity.
+      if (!ConstantExpr::getIntrinsicIdentity(II->getIntrinsicID(),
+                                              I->getType()))
+        return false;
+    }
+
+    // Exactly one operand should be the result of the call instruction.
+    if ((I->getOperand(0) == CI && I->getOperand(1) == CI) ||
+        (I->getOperand(0) != CI && I->getOperand(1) != CI))
+      return false;
+  }
 
   // The only user of this instruction we allow is a single return instruction.
   if (!I->hasOneUse() || !isa<ReturnInst>(I->user_back()))
     return false;
 
+  if (BaseCaseConst)
+    return BaseCaseConst;
+
   return true;
 }
 
@@ -436,6 +545,10 @@ class TailRecursionEliminator {
   // The instruction doing the accumulating.
   Instruction *AccumulatorRecursionInstr = nullptr;
 
+  // The base case return value. It exists only if the accumulator
+  // operation is pseudo-associative.
+  Constant *BaseCaseValue = nullptr;
+
   TailRecursionEliminator(Function &F, const TargetTransformInfo *TTI,
                           AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
                           DomTreeUpdater &DTU, BlockFrequencyInfo *BFI)
@@ -594,9 +707,17 @@ void TailRecursionEliminator::insertAccumulator(Instruction *AccRecInstr) {
   for (pred_iterator PI = PB; PI != PE; ++PI) {
     BasicBlock *P = *PI;
     if (P == &F.getEntryBlock()) {
-      Constant *Identity =
+      Constant *InitialValue =
           ConstantExpr::getIdentity(AccRecInstr, AccRecInstr->getType());
-      AccPN->addIncoming(Identity, P);
+      if (!InitialValue) {
+        // We are in the presence of a pseudo-associative operation like shift.
+        // We didn't pass `AllowRHSConstant = true` in getIdentity above because
+        // the identity for shifts is zero, which is not valid for the LHS.
+        // We need the value of the base case(s) of the function to
+        // initialize the accumulator.
+        InitialValue = BaseCaseValue;
+      }
+      AccPN->addIncoming(InitialValue, P);
     } else {
       AccPN->addIncoming(AccPN, P);
     }
@@ -667,15 +788,24 @@ bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
       continue;
 
     // If we can't move the instruction above the call, it might be because it
-    // is an associative and commutative operation that could be transformed
-    // using accumulator recursion elimination.  Check to see if this is the
-    // case, and if so, remember which instruction accumulates for later.
-    if (AccPN || !canTransformAccumulatorRecursion(&*BBI, CI))
+    // is an (associative and commutative) or pseudo-associative arithmetic
+    // operation that could be transformed using accumulator recursion
+    // elimination. Check to see if this is the case, and if so, remember which
+    // instruction accumulates for later.
+    std::variant<bool, Constant *> AccRecResult =
+        canTransformAccumulatorRecursion(&*BBI, CI);
+
+    if (AccPN || (std::holds_alternative<bool>(AccRecResult) &&
+                  !std::get<bool>(AccRecResult)))
       return false; // We cannot eliminate the tail recursion!
 
     // Yes, this is accumulator recursion.  Remember which instruction
     // accumulates.
     AccRecInstr = &*BBI;
+
+    // Keep track of the base case return value if any.
+    if (std::holds_alternative<Constant *>(AccRecResult))
+      BaseCaseValue = std::get<Constant *>(AccRecResult);
   }
 
   BasicBlock *BB = Ret->getParent();
@@ -814,13 +944,26 @@ void TailRecursionEliminator::cleanupAndFinalize() {
           if (!RI)
             continue;
 
-          Instruction *AccRecInstrNew = AccRecInstr->clone();
-          AccRecInstrNew->setName("accumulator.ret.tr");
-          AccRecInstrNew->setOperand(AccRecInstr->getOperand(0) == AccPN,
-                                     RI->getOperand(0));
-          AccRecInstrNew->insertBefore(RI->getIterator());
-          AccRecInstrNew->dropLocation();
-          RI->setOperand(0, AccRecInstrNew);
+          // Again, the BaseCaseValue exists only if the accumulator
+          // operation is pseudo-associative. In that case the accumulator PHI
+          // already holds the final result (it is initialized from the base
+          // case), so we should return it directly rather than applying the
+          // accumulation instruction one more time.
+          if (BaseCaseValue) {
+            RI->setOperand(0, AccPN);
+          } else {
+            // Otherwise, since the accumulator starts with the identity value,
+            // before the return we need to apply the accumulation instruction
+            // one more time to combine the last value with the result of the
+            // recursive call.
+            Instruction *AccRecInstrNew = AccRecInstr->clone();
+            AccRecInstrNew->setName("accumulator.ret.tr");
+            AccRecInstrNew->setOperand(AccRecInstr->getOperand(0) == AccPN,
+                                       RI->getOperand(0));
+            AccRecInstrNew->insertBefore(RI->getIterator());
+            AccRecInstrNew->dropLocation();
+            RI->setOperand(0, AccRecInstrNew);
+          }
         }
       }
     } else {
@@ -887,7 +1030,9 @@ bool TailRecursionEliminator::processBlock(BasicBlock &BB) {
 
     eliminateCall(CI);
     return true;
-  } else if (isa<ReturnInst>(TI)) {
+  }
+
+  if (isa<ReturnInst>(TI)) {
     CallInst *CI = findTRECandidate(&BB);
 
     if (CI)
diff --git a/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll b/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
new file mode 100644
index 0000000000000..257a885672cf6
--- /dev/null
+++ b/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
@@ -0,0 +1,118 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes="tailcallelim" -verify-dom-info -S | FileCheck %s
+
+; NOTE: All the following test cases are generate from the underlying C code (-O1)
+;   before that the shift accumulator optimization was implemented
+
+
+
+; InstCombine strength-reduce `f(x-1) + f(x-1)` to shl:
+; int f(int x) {
+;     if (x == 1) return 7;
+;     return f(x-1) + f(x-1); // f(x-1) * 2
+; }
+define dso_local i32 @f(i32 noundef %x) local_unnamed_addr {
+; CHECK-LABEL: define dso_local i32 @f(
+; CHECK-SAME: i32 noundef [[X:%.*]]) local_unnamed_addr {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
+; CHECK:       [[TAILRECURSE]]:
+; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = phi i32 [ 7, %[[ENTRY]] ], [ [[ADD:%.*]], %[[IF_END:.*]] ]
+; CHECK-NEXT:    [[X_TR:%.*]] = phi i32 [ [[X]], %[[ENTRY]] ], [ [[SUB:%.*]], %[[IF_END]] ]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X_TR]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END]]
+; CHECK:       [[COMMON_RET]]:
+; CHECK-NEXT:    ret i32 [[ACCUMULATOR_TR]]
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[SUB]] = add nsw i32 [[X_TR]], -1
+; CHECK-NEXT:    [[ADD]] = shl nsw i32 [[ACCUMULATOR_TR]], 1
+; CHECK-NEXT:    br label %[[TAILRECURSE]]
+;
+entry:
+  %cmp = icmp eq i32 %x, 1
+  br i1 %cmp, label %common.ret, label %if.end
+
+common.ret:
+  %common.ret.op = phi i32 [ %add, %if.end ], [ 7, %entry ]
+  ret i32 %common.ret.op
+
+if.end:
+  %sub = add nsw i32 %x, -1
+  %call = tail call i32 @f(i32 noundef %sub)
+  %add = shl nsw i32 %call, 1
+  br label %common.ret
+}
+
+
+; int f2(int x) {
+;     if (x == 1) return 14;
+;     return f2(x-1) >> 1;
+; }
+define dso_local i32 @f2(i32 noundef %x) local_unnamed_addr {
+; CHECK-LABEL: define dso_local i32 @f2(
+; CHECK-SAME: i32 noundef [[X:%.*]]) local_unnamed_addr {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
+; CHECK:       [[TAILRECURSE]]:
+; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = phi i32 [ 14, %[[ENTRY]] ], [ [[SHR:%.*]], %[[IF_END:.*]] ]
+; CHECK-NEXT:    [[X_TR:%.*]] = phi i32 [ [[X]], %[[ENTRY]] ], [ [[SUB:%.*]], %[[IF_END]] ]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X_TR]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END]]
+; CHECK:       [[COMMON_RET]]:
+; CHECK-NEXT:    ret i32 [[ACCUMULATOR_TR]]
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[SUB]] = add nsw i32 [[X_TR]], -1
+; CHECK-NEXT:    [[SHR]] = ashr i32 [[ACCUMULATOR_TR]], 1
+; CHECK-NEXT:    br label %[[TAILRECURSE]]
+;
+entry:
+  %cmp = icmp eq i32 %x, 1
+  br i1 %cmp, label %common.ret, label %if.end
+
+common.ret:
+  %common.ret.op = phi i32 [ %shr, %if.end ], [ 14, %entry ]
+  ret i32 %common.ret.op
+
+if.end:
+  %sub = add nsw i32 %x, -1
+  %call = tail call i32 @f2(i32 noundef %sub)
+  %shr = ashr i32 %call, 1
+  br label %common.ret
+}
+
+
+; unsigned int f3(unsigned int x) {
+;     if (x <= 1) return 14;
+;     return f3(x - 1) >> 1;
+; }
+define dso_local i32 @f3(i32 noundef %x) local_unnamed_addr {
+; CHECK-LABEL: define dso_local i32 @f3(
+; CHECK-SAME: i32 noundef [[X:%.*]]) local_unnamed_addr {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
+; CHECK:       [[TAILRECURSE]]:
+; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = phi i32 [ 21, %[[ENTRY]] ], [ [[SHR:%.*]], %[[IF_END:.*]] ]
+; CHECK-NEXT:    [[X_TR:%.*]] = phi i32 [ [[X]], %[[ENTRY]] ], [ [[SUB:%.*]], %[[IF_END]] ]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X_TR]], 2
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END]]
+; CHECK:       [[COMMON_RET]]:
+; CHECK-NEXT:    ret i32 [[ACCUMULATOR_TR]]
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[SUB]] = add i32 [[X_TR]], -1
+; CHECK-NEXT:    [[SHR]] = lshr i32 [[ACCUMULATOR_TR]], 1
+; CHECK-NEXT:    br label %[[TAILRECURSE]]
+;
+entry:
+  %cmp = icmp ult i32 %x, 2
+  br i1 %cmp, label %common.ret, label %if.end
+
+common.ret:
+  %common.ret.op = phi i32 [ %shr, %if.end ], [ 21, %entry ]
+  ret i32 %common.ret.op
+
+if.end:
+  %sub = add i32 %x, -1
+  %call = tail call i32 @f3(i32 noundef %sub)
+  %shr = lshr i32 %call, 1
+  br label %common.ret
+}

>From 81af21e5cf97328d74a7fd2cf924c56fa06adf15 Mon Sep 17 00:00:00 2001
From: FedericoBruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 18 Feb 2026 17:21:00 +0100
Subject: [PATCH 2/5] Add negative tests and remove useless attrs

Signed-off-by: FedericoBruzzone <federico.bruzzone.i at gmail.com>
---
 .../TailCallElim/shl-accumulator-opt.ll       | 88 +++++++++++++++++--
 1 file changed, 79 insertions(+), 9 deletions(-)

diff --git a/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll b/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
index 257a885672cf6..de615b8664d18 100644
--- a/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
+++ b/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
@@ -11,9 +11,9 @@
 ;     if (x == 1) return 7;
 ;     return f(x-1) + f(x-1); // f(x-1) * 2
 ; }
-define dso_local i32 @f(i32 noundef %x) local_unnamed_addr {
-; CHECK-LABEL: define dso_local i32 @f(
-; CHECK-SAME: i32 noundef [[X:%.*]]) local_unnamed_addr {
+define i32 @f(i32 %x) {
+; CHECK-LABEL: define i32 @f(
+; CHECK-SAME: i32 [[X:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
 ; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
 ; CHECK:       [[TAILRECURSE]]:
@@ -48,9 +48,9 @@ if.end:
 ;     if (x == 1) return 14;
 ;     return f2(x-1) >> 1;
 ; }
-define dso_local i32 @f2(i32 noundef %x) local_unnamed_addr {
-; CHECK-LABEL: define dso_local i32 @f2(
-; CHECK-SAME: i32 noundef [[X:%.*]]) local_unnamed_addr {
+define i32 @f2(i32 %x) {
+; CHECK-LABEL: define i32 @f2(
+; CHECK-SAME: i32 [[X:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
 ; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
 ; CHECK:       [[TAILRECURSE]]:
@@ -85,9 +85,9 @@ if.end:
 ;     if (x <= 1) return 14;
 ;     return f3(x - 1) >> 1;
 ; }
-define dso_local i32 @f3(i32 noundef %x) local_unnamed_addr {
-; CHECK-LABEL: define dso_local i32 @f3(
-; CHECK-SAME: i32 noundef [[X:%.*]]) local_unnamed_addr {
+define i32 @f3(i32 %x) {
+; CHECK-LABEL: define i32 @f3(
+; CHECK-SAME: i32 [[X:%.*]]) {
 ; CHECK-NEXT:  [[ENTRY:.*]]:
 ; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
 ; CHECK:       [[TAILRECURSE]]:
@@ -116,3 +116,73 @@ if.end:
   %shr = lshr i32 %call, 1
   br label %common.ret
 }
+
+; Negative test: non-constant shift amount prevents shift-accumulator optimization
+; int f4(int x, int k) {
+;   if (x == 1) return 7;
+;   return f4(x-1, k) << k; // variable shift amount
+; }
+define i32 @f4(i32 %x, i32 %k) {
+; CHECK-LABEL: define i32 @f4(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[K:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END:.*]]
+; CHECK:       [[COMMON_RET]]:
+; CHECK-NEXT:    ret i32 7
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[SUB:%.*]] = add nsw i32 [[X]], -1
+; CHECK-NEXT:    [[CALL:%.*]] = tail call i32 @f4(i32 [[SUB]], i32 [[K]])
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[CALL]], [[K]]
+; CHECK-NEXT:    ret i32 [[SHL]]
+;
+entry:
+  %cmp = icmp eq i32 %x, 1
+  br i1 %cmp, label %common.ret, label %if.end
+
+common.ret:
+  %common.ret.op = phi i32 [ %shl, %if.end ], [ 7, %entry ]
+  ret i32 %common.ret.op
+
+if.end:
+  %sub = add nsw i32 %x, -1
+  %call = tail call i32 @f4(i32 %sub, i32 %k)
+  %shl = shl i32 %call, %k
+  br label %common.ret
+}
+
+; Negative test: intervening operation on result prevents forming clean accumulator
+; int f5(int x) {
+;   if (x == 1) return 7;
+;   return (f5(x-1) + 1) << 1; // extra add breaks accumulator invariant
+; }
+define i32 @f5(i32 %x) {
+; CHECK-LABEL: define i32 @f5(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END:.*]]
+; CHECK:       [[COMMON_RET]]:
+; CHECK-NEXT:    ret i32 7
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[SUB:%.*]] = add nsw i32 [[X]], -1
+; CHECK-NEXT:    [[CALL:%.*]] = tail call i32 @f5(i32 [[SUB]])
+; CHECK-NEXT:    [[TMP:%.*]] = add i32 [[CALL]], 1
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[TMP]], 1
+; CHECK-NEXT:    ret i32 [[SHL]]
+;
+entry:
+  %cmp = icmp eq i32 %x, 1
+  br i1 %cmp, label %common.ret, label %if.end
+
+common.ret:
+  %common.ret.op = phi i32 [ %shl, %if.end ], [ 7, %entry ]
+  ret i32 %common.ret.op
+
+if.end:
+  %sub = add nsw i32 %x, -1
+  %call = tail call i32 @f5(i32 %sub)
+  %tmp = add i32 %call, 1
+  %shl = shl i32 %tmp, 1
+  br label %common.ret
+}

>From a8ed483eb9fa3ef7db8755928c0252c9c5c61b4f Mon Sep 17 00:00:00 2001
From: FedericoBruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 18 Feb 2026 18:01:53 +0100
Subject: [PATCH 3/5] Remove `ConstantExprs` and simplify `isPseudoAssociative`

Signed-off-by: FedericoBruzzone <federico.bruzzone.i at gmail.com>
---
 .../Scalar/TailRecursionElimination.cpp       | 20 +++----------------
 1 file changed, 3 insertions(+), 17 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index 888d0b223463b..d862c52b2334b 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -381,20 +381,9 @@ static bool canMoveAboveCall(Instruction *I, CallInst *CI, AliasAnalysis *AA) {
 // This relation applies to left shifts as well as arithmetic/logical right
 // shifts when the shift amount is a constant.
 static bool isPseudoAssociative(Instruction *I) {
-  switch (I->getOpcode()) {
-  case Instruction::Shl:
-  case Instruction::AShr:
-  case Instruction::LShr:
-    break;
-  default:
-    return false;
-  }
-
-  ConstantInt *CI = dyn_cast<ConstantInt>(I->getOperand(1));
-  if (!CI)
-    return false;
-
-  return true;
+    if (!I->isShift())
+      return false;
+    return isa<ConstantInt>(I->getOperand(1));
 }
 
 // Find the base-case return value for function F: examine all
@@ -430,9 +419,6 @@ static Constant *getReturnValue(Function &F) {
             return true;
         for (Use &U : I->operands())
           Worklist.push_back(U.get());
-      } else if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Cur)) {
-        for (unsigned I = 0, E = CE->getNumOperands(); I != E; ++I)
-          Worklist.push_back(CE->getOperand(I));
       }
     }
     return false;

>From 01dc728e94ec86ff6c710b2bf0f3c5fc8799109e Mon Sep 17 00:00:00 2001
From: FedericoBruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 18 Feb 2026 18:42:42 +0100
Subject: [PATCH 4/5] Replace the use of use-def chains with simpler matches

Signed-off-by: FedericoBruzzone <federico.bruzzone.i at gmail.com>
---
 .../Scalar/TailRecursionElimination.cpp       | 69 +++++++++++--------
 .../TailCallElim/shl-accumulator-opt.ll       | 63 ++++++++---------
 2 files changed, 66 insertions(+), 66 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index d862c52b2334b..105ab8337c2b6 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -381,9 +381,9 @@ static bool canMoveAboveCall(Instruction *I, CallInst *CI, AliasAnalysis *AA) {
 // This relation applies to left shifts as well as arithmetic/logical right
 // shifts when the shift amount is a constant.
 static bool isPseudoAssociative(Instruction *I) {
-    if (!I->isShift())
-      return false;
-    return isa<ConstantInt>(I->getOperand(1));
+  if (!I->isShift())
+    return false;
+  return isa<ConstantInt>(I->getOperand(1));
 }
 
 // Find the base-case return value for function F: examine all
@@ -405,39 +405,48 @@ static bool isPseudoAssociative(Instruction *I) {
 static Constant *getReturnValue(Function &F) {
   Constant *BaseCaseVal = nullptr;
 
-  auto ValueUsesRecursiveCall = [&](Value *V) {
-    SmallVector<Value *, 8> Worklist;
-    SmallPtrSet<Value *, 8> Visited;
-    Worklist.push_back(V);
-    while (!Worklist.empty()) {
-      Value *Cur = Worklist.pop_back_val();
-      if (!Visited.insert(Cur).second)
-        continue;
-      if (Instruction *I = dyn_cast<Instruction>(Cur)) {
-        if (CallInst *CI = dyn_cast<CallInst>(I))
-          if (CI->getCalledFunction() == &F)
-            return true;
-        for (Use &U : I->operands())
-          Worklist.push_back(U.get());
+  // Local lambda with conservative bail-out: It isn't so trivial to perform TRE
+  // if the return value depends on the result of a recursive call, because the
+  // return value will be different for different iterations of the recursion.
+  // So we ignore return values that depend on recursive calls.
+  auto TryGetConstantBaseCase = [&](Value *RV) -> Constant * {
+    // Case 1: direct constant return.
+    if (auto *C = dyn_cast<Constant>(RV))
+      return C;
+
+    // Case 2: PHI with one constant incoming.
+    if (auto *PN = dyn_cast<PHINode>(RV)) {
+      if (PN->getNumIncomingValues() != 2)
+        return nullptr;
+      for (unsigned I = 0; I < 2; ++I) {
+        if (auto *C = dyn_cast<Constant>(PN->getIncomingValue(I)))
+          return C;
       }
+      return nullptr;
     }
-    return false;
+
+    // Case 3: select with a constant arm.
+    if (auto *SI = dyn_cast<SelectInst>(RV)) {
+      if (auto *CTrue = dyn_cast<Constant>(SI->getTrueValue()))
+        return CTrue;
+      if (auto *CFalse = dyn_cast<Constant>(SI->getFalseValue()))
+        return CFalse;
+      return nullptr;
+    }
+
+    return nullptr;
   };
 
   for (BasicBlock &BB : F) {
-    if (ReturnInst *RI = dyn_cast<ReturnInst>(BB.getTerminator())) {
+    if (auto *RI = dyn_cast<ReturnInst>(BB.getTerminator())) {
       Value *RV = RI->getReturnValue();
-      if (!RV)
-        continue;
-      // It isn't so trivial to perform TRE if the return value depends on the
-      // result of a recursive call, because the return value will be different
-      // for different iterations of the recursion. So we ignore return values
-      // that depend on recursive calls.
-      if (ValueUsesRecursiveCall(RV))
-        continue;
-      if (!BaseCaseVal && isa<Constant>(RV))
-        BaseCaseVal = cast<Constant>(RV);
-      else if (BaseCaseVal != RV)
+      Constant *Candidate = TryGetConstantBaseCase(RV);
+      if (!Candidate)
+        return nullptr;
+
+      if (!BaseCaseVal)
+        BaseCaseVal = Candidate;
+      else if (BaseCaseVal != Candidate)
         return nullptr;
     }
   }
diff --git a/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll b/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
index de615b8664d18..81ba00e073b85 100644
--- a/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
+++ b/llvm/test/Transforms/TailCallElim/shl-accumulator-opt.ll
@@ -14,19 +14,16 @@
 define i32 @f(i32 %x) {
 ; CHECK-LABEL: define i32 @f(
 ; CHECK-SAME: i32 [[X:%.*]]) {
-; CHECK-NEXT:  [[ENTRY:.*]]:
-; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
-; CHECK:       [[TAILRECURSE]]:
-; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = phi i32 [ 7, %[[ENTRY]] ], [ [[ADD:%.*]], %[[IF_END:.*]] ]
-; CHECK-NEXT:    [[X_TR:%.*]] = phi i32 [ [[X]], %[[ENTRY]] ], [ [[SUB:%.*]], %[[IF_END]] ]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X_TR]], 1
-; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END]]
+; CHECK-NEXT:  [[TAILRECURSE:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END:.*]]
 ; CHECK:       [[COMMON_RET]]:
-; CHECK-NEXT:    ret i32 [[ACCUMULATOR_TR]]
+; CHECK-NEXT:    ret i32 7
 ; CHECK:       [[IF_END]]:
-; CHECK-NEXT:    [[SUB]] = add nsw i32 [[X_TR]], -1
-; CHECK-NEXT:    [[ADD]] = shl nsw i32 [[ACCUMULATOR_TR]], 1
-; CHECK-NEXT:    br label %[[TAILRECURSE]]
+; CHECK-NEXT:    [[SUB:%.*]] = add nsw i32 [[X]], -1
+; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = tail call i32 @f(i32 [[SUB]])
+; CHECK-NEXT:    [[ADD:%.*]] = shl nsw i32 [[ACCUMULATOR_TR]], 1
+; CHECK-NEXT:    ret i32 [[ADD]]
 ;
 entry:
   %cmp = icmp eq i32 %x, 1
@@ -38,7 +35,7 @@ common.ret:
 
 if.end:
   %sub = add nsw i32 %x, -1
-  %call = tail call i32 @f(i32 noundef %sub)
+  %call = tail call i32 @f(i32 %sub)
   %add = shl nsw i32 %call, 1
   br label %common.ret
 }
@@ -51,19 +48,16 @@ if.end:
 define i32 @f2(i32 %x) {
 ; CHECK-LABEL: define i32 @f2(
 ; CHECK-SAME: i32 [[X:%.*]]) {
-; CHECK-NEXT:  [[ENTRY:.*]]:
-; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
-; CHECK:       [[TAILRECURSE]]:
-; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = phi i32 [ 14, %[[ENTRY]] ], [ [[SHR:%.*]], %[[IF_END:.*]] ]
-; CHECK-NEXT:    [[X_TR:%.*]] = phi i32 [ [[X]], %[[ENTRY]] ], [ [[SUB:%.*]], %[[IF_END]] ]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X_TR]], 1
-; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END]]
+; CHECK-NEXT:  [[TAILRECURSE:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X]], 1
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END:.*]]
 ; CHECK:       [[COMMON_RET]]:
-; CHECK-NEXT:    ret i32 [[ACCUMULATOR_TR]]
+; CHECK-NEXT:    ret i32 14
 ; CHECK:       [[IF_END]]:
-; CHECK-NEXT:    [[SUB]] = add nsw i32 [[X_TR]], -1
-; CHECK-NEXT:    [[SHR]] = ashr i32 [[ACCUMULATOR_TR]], 1
-; CHECK-NEXT:    br label %[[TAILRECURSE]]
+; CHECK-NEXT:    [[SUB:%.*]] = add nsw i32 [[X]], -1
+; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = tail call i32 @f2(i32 [[SUB]])
+; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[ACCUMULATOR_TR]], 1
+; CHECK-NEXT:    ret i32 [[SHR]]
 ;
 entry:
   %cmp = icmp eq i32 %x, 1
@@ -75,7 +69,7 @@ common.ret:
 
 if.end:
   %sub = add nsw i32 %x, -1
-  %call = tail call i32 @f2(i32 noundef %sub)
+  %call = tail call i32 @f2(i32 %sub)
   %shr = ashr i32 %call, 1
   br label %common.ret
 }
@@ -88,19 +82,16 @@ if.end:
 define i32 @f3(i32 %x) {
 ; CHECK-LABEL: define i32 @f3(
 ; CHECK-SAME: i32 [[X:%.*]]) {
-; CHECK-NEXT:  [[ENTRY:.*]]:
-; CHECK-NEXT:    br label %[[TAILRECURSE:.*]]
-; CHECK:       [[TAILRECURSE]]:
-; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = phi i32 [ 21, %[[ENTRY]] ], [ [[SHR:%.*]], %[[IF_END:.*]] ]
-; CHECK-NEXT:    [[X_TR:%.*]] = phi i32 [ [[X]], %[[ENTRY]] ], [ [[SUB:%.*]], %[[IF_END]] ]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X_TR]], 2
-; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END]]
+; CHECK-NEXT:  [[TAILRECURSE:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X]], 2
+; CHECK-NEXT:    br i1 [[CMP]], label %[[COMMON_RET:.*]], label %[[IF_END:.*]]
 ; CHECK:       [[COMMON_RET]]:
-; CHECK-NEXT:    ret i32 [[ACCUMULATOR_TR]]
+; CHECK-NEXT:    ret i32 21
 ; CHECK:       [[IF_END]]:
-; CHECK-NEXT:    [[SUB]] = add i32 [[X_TR]], -1
-; CHECK-NEXT:    [[SHR]] = lshr i32 [[ACCUMULATOR_TR]], 1
-; CHECK-NEXT:    br label %[[TAILRECURSE]]
+; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[X]], -1
+; CHECK-NEXT:    [[ACCUMULATOR_TR:%.*]] = tail call i32 @f3(i32 [[SUB]])
+; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 [[ACCUMULATOR_TR]], 1
+; CHECK-NEXT:    ret i32 [[SHR]]
 ;
 entry:
   %cmp = icmp ult i32 %x, 2
@@ -112,7 +103,7 @@ common.ret:
 
 if.end:
   %sub = add i32 %x, -1
-  %call = tail call i32 @f3(i32 noundef %sub)
+  %call = tail call i32 @f3(i32 %sub)
   %shr = lshr i32 %call, 1
   br label %common.ret
 }

>From ed0f0ba7c824fffb2f5b319f20a9dffc5480734b Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 14 Apr 2026 11:58:02 +0200
Subject: [PATCH 5/5] Replace `std::variant` with a structured return type

---
 .../Scalar/TailRecursionElimination.cpp       | 40 ++++++++++---------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index 105ab8337c2b6..c7abb67143ce7 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -82,7 +82,6 @@
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <cmath>
-#include <variant>
 using namespace llvm;
 
 #define DEBUG_TYPE "tailcallelim"
@@ -454,17 +453,24 @@ static Constant *getReturnValue(Function &F) {
   return BaseCaseVal;
 }
 
+namespace {
+struct AccumulatorRecursionInfo {
+  bool CanTransform = false;
+  Constant *BaseCaseConst = nullptr;
+};
+} // namespace
+
 // This function checks whether the instruction I can be used
 // to perform accumulator recursion elimination for the
 // call instruction CI.
 // In the presence of pseudo-associative operations, it returns
-// the base case constant value to both indicate success and
+// the base case constant value in BaseCaseConst to both indicate success and
 // provide the value needed to initialize the accumulator.
-static std::variant<bool, Constant *>
-canTransformAccumulatorRecursion(Instruction *I, CallInst *CI) {
+static AccumulatorRecursionInfo canTransformAccumulatorRecursion(Instruction *I,
+                                                                 CallInst *CI) {
   auto IsPseudoAssociative = isPseudoAssociative(I);
   if ((!I->isAssociative() || !I->isCommutative()) && !IsPseudoAssociative)
-    return false;
+    return {};
 
   assert(I->getNumOperands() >= 2 &&
          "Associative/commutative operations should have at least 2 args!");
@@ -474,33 +480,30 @@ canTransformAccumulatorRecursion(Instruction *I, CallInst *CI) {
     // For pseudo-associative operations, we require that the recursive call
     // is always on the first operand.
     if (I->getOperand(0) != CI)
-      return false;
+      return {};
 
     BaseCaseConst = getReturnValue(*CI->getCalledFunction());
     if (!BaseCaseConst)
-      return false;
+      return {};
   } else {
     if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
       // Accumulators must have an identity.
       if (!ConstantExpr::getIntrinsicIdentity(II->getIntrinsicID(),
                                               I->getType()))
-        return false;
+        return {};
     }
 
     // Exactly one operand should be the result of the call instruction.
     if ((I->getOperand(0) == CI && I->getOperand(1) == CI) ||
         (I->getOperand(0) != CI && I->getOperand(1) != CI))
-      return false;
+      return {};
   }
 
   // The only user of this instruction we allow is a single return instruction.
   if (!I->hasOneUse() || !isa<ReturnInst>(I->user_back()))
-    return false;
-
-  if (BaseCaseConst)
-    return BaseCaseConst;
+    return {};
 
-  return true;
+  return {true, BaseCaseConst};
 }
 
 namespace {
@@ -787,11 +790,10 @@ bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
     // operation that could be transformed using accumulator recursion
     // elimination. Check to see if this is the case, and if so, remember which
     // instruction accumulates for later.
-    std::variant<bool, Constant *> AccRecResult =
+    AccumulatorRecursionInfo AccRecInfo =
         canTransformAccumulatorRecursion(&*BBI, CI);
 
-    if (AccPN || (std::holds_alternative<bool>(AccRecResult) &&
-                  !std::get<bool>(AccRecResult)))
+    if (AccPN || !AccRecInfo.CanTransform)
       return false; // We cannot eliminate the tail recursion!
 
     // Yes, this is accumulator recursion.  Remember which instruction
@@ -799,8 +801,8 @@ bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
     AccRecInstr = &*BBI;
 
     // Keep track of the base case return value if any.
-    if (std::holds_alternative<Constant *>(AccRecResult))
-      BaseCaseValue = std::get<Constant *>(AccRecResult);
+    if (AccRecInfo.BaseCaseConst)
+      BaseCaseValue = AccRecInfo.BaseCaseConst;
   }
 
   BasicBlock *BB = Ret->getParent();



More information about the llvm-commits mailing list