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

[Antonio Frighetto via llvm-commits]

================
@@ -374,29 +376,110 @@ 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));
+}
+
+// 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;
+
+  for (BasicBlock &BB : F) {
+    auto *RI = dyn_cast<ReturnInst>(BB.getTerminator());
+    if (!RI)
+      continue;
+
+    Value *RV = RI->getReturnValue();
+    Constant *Candidate = nullptr;
+
+    if (match(RV, m_Constant(Candidate))) {
+      // Direct constant return.
+    } else if (auto *PN = dyn_cast<PHINode>(RV)) {
+      for (Value *Inc : PN->incoming_values())
+        if (match(Inc, m_Constant(Candidate)))
+          break;
+    } else if (match(RV,
+                     m_Select(m_Value(), m_Constant(Candidate), m_Value())) ||
+               match(RV,
+                     m_Select(m_Value(), m_Value(), m_Constant(Candidate)))) {
+      // Select with a constant arm.
+    }
+
+    if (!Candidate)
+      return nullptr;
+
+    if (!BaseCaseVal)
+      BaseCaseVal = Candidate;
+    else if (BaseCaseVal != Candidate)
+      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.
+static Constant *canTransformAccumulatorRecursion(Instruction *I,
+                                                  CallInst *CI) {
+  auto IsPseudoAssociative = isPseudoAssociative(I);
+  if ((!I->isAssociative() || !I->isCommutative()) && !IsPseudoAssociative)
+    return nullptr;
 
   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()))
-      return false;
-  }
+  Constant *AccInitVal = 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 nullptr;
 
-  // 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;
+    AccInitVal = getReturnValue(*CI->getCalledFunction());
----------------
antoniofrighetto wrote:

Why are we passing the whole function here in getReturnValue()? We should already be in the return basic block when encountering a tail call (see the logic in eliminateCall()). I think we should only ensure to match a recursive call.

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