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

[Antonio Frighetto via llvm-commits]

================
@@ -374,29 +374,117 @@ 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) {
+  if (!I->isShift())
     return false;
+  return isa<ConstantInt>(I->getOperand(1));
+}
+
+// Return true if V is a recursive call to F or an instruction directly using
+// the result of one. A depth-1 check is enough here: the value feeding a
+// return either uses the recursive call as an immediate operand (the
+// accumulator instruction, or the PHI merging it with the base case), or it
+// is rejected by getReturnValue below as a non-constant anyway.
+static bool usesRecursiveCall(Value *V, Function &F) {
+  auto IsRecursiveCall = [&F](Value *V) {
+    auto *CI = dyn_cast<CallInst>(V);
+    return CI && CI->getCalledFunction() == &F;
+  };
+  if (IsRecursiveCall(V))
+    return true;
+  auto *I = dyn_cast<Instruction>(V);
+  return I && llvm::any_of(I->operands(), IsRecursiveCall);
+}
+
+// Find the base-case return value for function F: examine all return
+// instructions, skipping those whose return value depends on a recursive call
+// to F (that value differs for each iteration of the recursion). If the
+// remaining returns yield exactly one distinct constant, return it; otherwise
+// 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;
+
+  for (BasicBlock &BB : F) {
+    auto *RI = dyn_cast<ReturnInst>(BB.getTerminator());
+    if (!RI)
+      continue;
+
+    Value *RV = RI->getReturnValue();
+    if (usesRecursiveCall(RV, F))
+      continue;
+
+    auto *C = dyn_cast<Constant>(RV);
+    if (!C)
+      return nullptr;
+
+    if (!BaseCaseVal)
+      BaseCaseVal = C;
+    else if (BaseCaseVal != C)
+      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 in BaseCaseConst to both indicate success and
+// provide the value needed to initialize the accumulator.
----------------
antoniofrighetto wrote:

```suggestion
// This function checks whether the instruction I can be used
// to perform accumulator recursion elimination for the
// call instruction CI.
```

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