[SCEV] Support modelling of same base pointer `select`s in more complex than most trivial cases (when there is a base variable offset)

[lebedev.ri]

Summary

umin_seq select modelling requires that one of the hands of the select must be a constant.
After we remove pointer base, in all but most trivial cases,
there will likely be some other offset between the base and the selected-between pointers,
and if both of them happen to have some variable offset, we fail,
even if they have *the same* variable offset.

Instead, we should try to factor it out, perhaps that leaves us with a constant in one hand.

[mkazantsev]

Looks nice, but I'd suggest some restructuring.

llvm/lib/Analysis/ScalarEvolution.cpp

6001

Should it be a method of ScalarEvolution? I looks like it can be reused somewhere.

Should we make it bool and return false if there are no common expressions?

6036

I'd rather bail here if we detected that no common ops were found just to save CT on SCEV simplification.

6040

Thinking more of it, as a follow-up maybe: if both old LHS and RHS had no-wrap flags, then all new expressions and Common expr should also preserve them. Need to think more, but maybe it's worth doing.

llvm/lib/Analysis/ScalarEvolution.cpp

 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ConstantFolding.h"
case ICmpInst::ICMP_EQ:
     break;
   default:
     break;
   }
 
   return getUnknown(I);
 }
 
+static void factorCommonOperandsOutOfAddExprs(
+    ScalarEvolution *SE, const SCEV *LHSExpr, const SCEV *RHSExpr,
+    const SCEV *&CommonExpr, const SCEV *&NewLHSExpr, const SCEV *&NewRHSExpr) {
+  Type *Ty = LHSExpr->getType();
+  assert(RHSExpr->getType() == Ty && "Unexpected types!");
+
+  auto getOps = [](const SCEV **Expr) -> SCEVNAryExpr::op_range {
+    if (auto *AddExpr = dyn_cast<SCEVAddExpr>(*Expr))
+      return AddExpr->operands();
+    return make_range(Expr, Expr + 1);
+  };
+
+  struct NumOccurrences {
+    unsigned LHS = 0, RHS = 0;
+  };
+  SmallDenseMap<const SCEV *, NumOccurrences, 16> Ops;
+  for (const SCEV *Op : getOps(&LHSExpr))
+    ++Ops[Op].LHS;
+  for (const SCEV *Op : getOps(&RHSExpr))
+    ++Ops[Op].RHS;
+
+  SmallVector<const SCEV *, 16> CommonOps, NewLHSOps, NewRHSOps;
+  for (auto I : Ops) {
+    const SCEV *Op = I.first;
+    const NumOccurrences Times = I.second;
+
+    unsigned NumCommonOccurrences = std::min(Times.LHS, Times.RHS);
+
+    CommonOps.emplace_back(
+        SE->getMulExpr(SE->getConstant(Ty, NumCommonOccurrences), Op));
+    NewLHSOps.emplace_back(SE->getMulExpr(
+        SE->getConstant(Ty, Times.LHS - NumCommonOccurrences), Op));
+    NewRHSOps.emplace_back(SE->getMulExpr(
+        SE->getConstant(Ty, Times.RHS - NumCommonOccurrences), Op));
+  }
+
+  CommonExpr = SE->getAddExpr(CommonOps);
+  NewLHSExpr = SE->getAddExpr(NewLHSOps);
+  NewRHSExpr = SE->getAddExpr(NewRHSOps);
+}
+
 static Optional<const SCEV *>
 createNodeForSelectViaUMinSeq(ScalarEvolution *SE, const SCEV *CondExpr,
                               const SCEV *TrueExpr, const SCEV *FalseExpr) {
   assert(CondExpr->getType()->isIntegerTy(1) &&
          TrueExpr->getType() == FalseExpr->getType() &&
          TrueExpr->getType()->isIntegerTy() &&
          "Unexpected operands of a select.");
 
+  // i1 cond ? (iN base + iN X + iN C0) : (iN base + iN C1)
+  //   -->
+  // iN base + (i1 cond ? (iN X + iN C0) : iN C1)
+
   // i1 cond ? iN x : iN C  -->  C + (i1  cond ? (iN x - iN C) : iN 0)
   //                        -->  C + (umin_seq (sext  cond), x - C)
   //
   // i1 cond ? iN C : iN x  -->  C + (i1  cond ? iN 0 : (iN x - iN C))
   //                        -->  C + (i1 ~cond ? (iN x - iN C) : iN 0)
   //                        -->  C + (umin_seq (sext ~cond), x - C)
 
-  // FIXME: while we can't legally model the case where both of the hands
-  // are fully variable, we only require that the *difference* is constant.
+  const SCEV *Base = SE->getZero(TrueExpr->getType());
+
+  // We require that at least one hand of the select be a constant.
+  if (!isa<SCEVConstant>(TrueExpr) && !isa<SCEVConstant>(FalseExpr)) {
+    factorCommonOperandsOutOfAddExprs(SE, TrueExpr, FalseExpr, Base, TrueExpr,
+                                      FalseExpr);
+
     if (!isa<SCEVConstant>(TrueExpr) && !isa<SCEVConstant>(FalseExpr))
       return None;
+  }
 
   const SCEV *X, *C;
   if (SE->getTypeSizeInBits(TrueExpr->getType()) >
       SE->getTypeSizeInBits(CondExpr->getType()))
     CondExpr = SE->getSignExtendExpr(CondExpr, TrueExpr->getType());
   if (isa<SCEVConstant>(TrueExpr)) {
     CondExpr = SE->getNotSCEV(CondExpr);
     X = FalseExpr;
     C = TrueExpr;
   } else {
     X = TrueExpr;
     C = FalseExpr;
   }
-  return SE->getAddExpr(C, SE->getUMinExpr(CondExpr, SE->getMinusSCEV(X, C),
-                                           /*Sequential=*/true));
+  return SE->getAddExpr(
+      Base, C,
+      SE->getUMinExpr(CondExpr, SE->getMinusSCEV(X, C), /*Sequential=*/true));
 }
 
 static Optional<const SCEV *> createNodeForSelectViaUMinSeq(ScalarEvolution *SE,
                                                             Value *Cond,
                                                             Value *TrueVal,
                                                             Value *FalseVal) {
   if (!isa<ConstantInt>(TrueVal) && !isa<ConstantInt>(FalseVal))
     return None;

llvm/test/Analysis/ScalarEvolution/pointer-select.ll

 ; CHECK-NEXT:  Classifying expressions for: @pointer_select_same_object_with_variable_base_offset__constant_offsets
 ; CHECK-NEXT:    %obj = getelementptr i8, ptr %obj.base, i64 %base_offset
 ; CHECK-NEXT:    --> (%base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %true_ptr = getelementptr i8, ptr %obj, i64 42
 ; CHECK-NEXT:    --> (42 + %base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %false_ptr = getelementptr i8, ptr %obj, i64 24
 ; CHECK-NEXT:    --> (24 + %base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %r = select i1 %cond, ptr %true_ptr, ptr %false_ptr
-; CHECK-NEXT:    --> %r U: full-set S: full-set
+; CHECK-NEXT:    --> (42 + (-18 umin (-1 + (-1 * (sext i1 %cond to i64))<nsw>)<nsw>) + %base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:  Determining loop execution counts for: @pointer_select_same_object_with_variable_base_offset__constant_offsets
 ;
   %obj = getelementptr i8, ptr %obj.base, i64 %base_offset
   %true_ptr = getelementptr i8, ptr %obj, i64 42
   %false_ptr = getelementptr i8, ptr %obj, i64 24
   %r = select i1 %cond, ptr %true_ptr, ptr %false_ptr
   ret ptr %r
 }

 ; CHECK-NEXT:  Classifying expressions for: @pointer_select_same_object_with_variable_base_offset__constant_offset_vs_variable_offset
 ; CHECK-NEXT:    %obj = getelementptr i8, ptr %obj.base, i64 %base_offset
 ; CHECK-NEXT:    --> (%base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %true_ptr = getelementptr i8, ptr %obj, i64 42
 ; CHECK-NEXT:    --> (42 + %base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %false_ptr = getelementptr i8, ptr %obj, i64 %false_off
 ; CHECK-NEXT:    --> (%base_offset + %false_off + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %r = select i1 %cond, ptr %true_ptr, ptr %false_ptr
-; CHECK-NEXT:    --> %r U: full-set S: full-set
+; CHECK-NEXT:    --> (42 + ((-1 + (-1 * (sext i1 %cond to i64))<nsw>)<nsw> umin_seq (-42 + %false_off)) + %base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:  Determining loop execution counts for: @pointer_select_same_object_with_variable_base_offset__constant_offset_vs_variable_offset
 ;
   %obj = getelementptr i8, ptr %obj.base, i64 %base_offset
   %true_ptr = getelementptr i8, ptr %obj, i64 42
   %false_ptr = getelementptr i8, ptr %obj, i64 %false_off
   %r = select i1 %cond, ptr %true_ptr, ptr %false_ptr
   ret ptr %r
 }
 
 define ptr @pointer_select_same_object_with_variable_base_offset__variable_offset_vs_constant_offset(i1 %cond, ptr %obj.base, i64 %base_offset, i64 %true_off) {
 ; CHECK-LABEL: 'pointer_select_same_object_with_variable_base_offset__variable_offset_vs_constant_offset'
 ; CHECK-NEXT:  Classifying expressions for: @pointer_select_same_object_with_variable_base_offset__variable_offset_vs_constant_offset
 ; CHECK-NEXT:    %obj = getelementptr i8, ptr %obj.base, i64 %base_offset
 ; CHECK-NEXT:    --> (%base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %true_ptr = getelementptr i8, ptr %obj, i64 %true_off
 ; CHECK-NEXT:    --> (%base_offset + %true_off + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %false_ptr = getelementptr i8, ptr %obj, i64 42
 ; CHECK-NEXT:    --> (42 + %base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:    %r = select i1 %cond, ptr %true_ptr, ptr %false_ptr
-; CHECK-NEXT:    --> %r U: full-set S: full-set
+; CHECK-NEXT:    --> (42 + ((sext i1 %cond to i64) umin_seq (-42 + %true_off)) + %base_offset + %obj.base) U: full-set S: full-set
 ; CHECK-NEXT:  Determining loop execution counts for: @pointer_select_same_object_with_variable_base_offset__variable_offset_vs_constant_offset
 ;
   %obj = getelementptr i8, ptr %obj.base, i64 %base_offset
   %true_ptr = getelementptr i8, ptr %obj, i64 %true_off
   %false_ptr = getelementptr i8, ptr %obj, i64 42
   %r = select i1 %cond, ptr %true_ptr, ptr %false_ptr
   ret ptr %r
 }