ConstraintElimination: use dry-run to upper-bound cols

llvm/include/llvm/Analysis/ConstraintSystem.h

@@ -71,11 +71,7 @@ class ConstraintSystem {
   ConstraintSystem() : Constraints(6), View(Constraints) {}
 
   // This constructor is used by ConstraintElimination, inside ConstraintInfo.
-  // Unfortunately, due to calls to addFact, that adds local variables, it is
-  // impossible to know how many local variables there are in advance.
-  // ConstraintElimination has a fixed upper-bound on the number of columns,
-  // configurable as a cl::opt, so use that number, and don't add the constraint
-  // if it exceeds that number.
+  // ConstraintElimination upper-bounds the number of columns using a dry run.
   ConstraintSystem(ArrayRef<Value *> FunctionArgs, size_t NCols)
       : Constraints(NCols), View(Constraints) {
     NumVariables += FunctionArgs.size();
@@ -127,11 +123,6 @@ class ConstraintSystem {
     if (all_of(ArrayRef(R).drop_front(1), [](int64_t C) { return C == 0; }))
       return false;
 
-    // There is no correctness issue if we don't add a constraint, for whatever
-    // reason.
-    if (R.size() > Constraints.getNumCols())
-      return false;
-
     NumVariables = std::max(R.size(), NumVariables);
     return addVariableRow(R);
   }

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

@@ -56,10 +56,6 @@ static cl::opt<unsigned>
     MaxRows("constraint-elimination-max-rows", cl::init(500), cl::Hidden,
             cl::desc("Maximum number of rows to keep in constraint system"));
 
-static cl::opt<unsigned> MaxColumns(
-    "constraint-elimination-max-cols", cl::init(16), cl::Hidden,
-    cl::desc("Maximum number of columns to keep in constraint system"));
-
 static cl::opt<bool> DumpReproducers(
     "constraint-elimination-dump-reproducers", cl::init(false), cl::Hidden,
     cl::desc("Dump IR to reproduce successful transformations."));
@@ -276,7 +272,8 @@ class ConstraintInfo {
   const DataLayout &DL;
 
 public:
-  ConstraintInfo(const DataLayout &DL, ArrayRef<Value *> FunctionArgs)
+  ConstraintInfo(const DataLayout &DL, ArrayRef<Value *> FunctionArgs,
+                 unsigned MaxColumns)
       : UnsignedCS(FunctionArgs, MaxColumns),
         SignedCS(FunctionArgs, MaxColumns), DL(DL) {
     auto &Value2Index = getValue2Index(false);
@@ -1675,21 +1672,68 @@ tryToSimplifyOverflowMath(IntrinsicInst *II, ConstraintInfo &Info,
   return Changed;
 }
 
-static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
-                                 ScalarEvolution &SE,
-                                 OptimizationRemarkEmitter &ORE) {
-  bool Changed = false;
+/// Calls Info.getConstraint, discarding the result, and only collecting the
+/// number of new variables introduced by the constraint.
+static unsigned getNumNewVars(CmpInst::Predicate Pred, Value *A, Value *B,
+                              const ConstraintInfo &Info) {
+  auto GetIncrVal = [&Info](CmpInst::Predicate Pred, Value *A, Value *B) {
+    SmallVector<Value *> NewVars;
+    Info.getConstraint(Pred, A, B, NewVars);
+    return NewVars.size();
+  };
+
+  unsigned NumNewVars = GetIncrVal(Pred, A, B);
+
+  // What follows is an inlined dry-run of transferToOtherSystem, which
+  // upper-bounds NumNewVars conservatively.
+  if (!A->getType()->isIntegerTy())
+    return NumNewVars;
+
+  switch (Pred) {
+  default:
+    break;
+  case CmpInst::ICMP_ULT:
+  case CmpInst::ICMP_ULE:
+    NumNewVars +=
+        GetIncrVal(CmpInst::ICMP_SGE, A, ConstantInt::get(B->getType(), 0));
+    NumNewVars += GetIncrVal(CmpInst::getSignedPredicate(Pred), A, B);
+    break;
+  case CmpInst::ICMP_UGE:
+  case CmpInst::ICMP_UGT:
+    //  If A is a signed positive constant, then B >=s 0 and A >s (or >=s)
+    //  B.
+    NumNewVars +=
+        GetIncrVal(CmpInst::ICMP_SGE, B, ConstantInt::get(B->getType(), 0));
+    NumNewVars += GetIncrVal(CmpInst::getSignedPredicate(Pred), A, B);
+    break;
+  case CmpInst::ICMP_SLT:
+    NumNewVars += GetIncrVal(CmpInst::ICMP_ULT, A, B);
+    break;
+  case CmpInst::ICMP_SGT: {
+    NumNewVars +=
+        GetIncrVal(CmpInst::ICMP_UGE, A, ConstantInt::get(B->getType(), 0));
+    NumNewVars += GetIncrVal(CmpInst::ICMP_UGT, A, B);
+    break;
+  }
+  case CmpInst::ICMP_SGE:
+    NumNewVars += GetIncrVal(CmpInst::ICMP_UGE, A, B);
+    break;
+  }
+  return NumNewVars;
+}
+
+/// Performs a dry run of the transform, computing a conservative estimate of
+/// the total number of columns we need in the underlying storage.
+static std::pair<State, unsigned> dryRun(Function &F, DominatorTree &DT,
+                                         LoopInfo &LI, ScalarEvolution &SE) {
   DT.updateDFSNumbers();
   SmallVector<Value *> FunctionArgs;
   for (Value &Arg : F.args())
     FunctionArgs.push_back(&Arg);
-  ConstraintInfo Info(F.getDataLayout(), FunctionArgs);
   State S(DT, LI, SE);
-  std::unique_ptr<Module> ReproducerModule(
-      DumpReproducers ? new Module(F.getName(), F.getContext()) : nullptr);
+  unsigned MaxColumns = FunctionArgs.size() + 1;
+  ConstraintInfo Info(F.getDataLayout(), FunctionArgs, MaxColumns);
 
-  // First, collect conditions implied by branches and blocks with their
-  // Dominator DFS in and out numbers.
   for (BasicBlock &BB : F) {
     if (!DT.getNode(&BB))
       continue;
@@ -1729,6 +1773,62 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
     return A.NumIn < B.NumIn;
   });
 
+  for (const FactOrCheck &CB : S.WorkList) {
+    ICmpInst::Predicate Pred;
+    Value *A, *B;
+    if (CB.isCheck()) {
+      if (auto *Cmp = dyn_cast<ICmpInst>(CB.getInstructionToSimplify())) {
+        if (!match(Cmp, m_ICmp(Pred, m_Value(A), m_Value(B))))
+          continue;
+        MaxColumns += getNumNewVars(Pred, A, B, Info);
+      }
+      continue;
+    }
+    if (!CB.isConditionFact()) {
+      Value *X;
+      if (match(CB.Inst, m_Intrinsic<Intrinsic::abs>(m_Value(X)))) {
+        if (cast<ConstantInt>(CB.Inst->getOperand(1))->isOne())
+          MaxColumns +=
+              getNumNewVars(CmpInst::ICMP_SGE, CB.Inst,
+                            ConstantInt::get(CB.Inst->getType(), 0), Info);
+        MaxColumns += getNumNewVars(CmpInst::ICMP_SGE, CB.Inst, X, Info);
+        continue;
+      }
+
+      if (auto *MinMax = dyn_cast<MinMaxIntrinsic>(CB.Inst)) {
+        Pred = ICmpInst::getNonStrictPredicate(MinMax->getPredicate());
+        MaxColumns += getNumNewVars(Pred, MinMax, MinMax->getLHS(), Info);
+        MaxColumns += getNumNewVars(Pred, MinMax, MinMax->getRHS(), Info);
+        continue;
+      }
+    }
+
+    if (CB.isConditionFact()) {
+      Pred = CB.Cond.Pred;
+      A = CB.Cond.Op0;
+      B = CB.Cond.Op1;
+    } else {
+      match(CB.Inst, m_Intrinsic<Intrinsic::assume>(
+                         m_ICmp(Pred, m_Value(A), m_Value(B))));
+    }
+    MaxColumns += getNumNewVars(Pred, A, B, Info);
+  }
+  return {S, MaxColumns};
+}
+
+static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
+                                 ScalarEvolution &SE,
+                                 OptimizationRemarkEmitter &ORE) {
+  bool Changed = false;
+  DT.updateDFSNumbers();
+  SmallVector<Value *> FunctionArgs;
+  for (Value &Arg : F.args())
+    FunctionArgs.push_back(&Arg);
+  auto [S, MaxColumns] = dryRun(F, DT, LI, SE);
+  ConstraintInfo Info(F.getDataLayout(), FunctionArgs, MaxColumns);
+  std::unique_ptr<Module> ReproducerModule(
+      DumpReproducers ? new Module(F.getName(), F.getContext()) : nullptr);
+
   SmallVector<Instruction *> ToRemove;
 
   // Finally, process ordered worklist and eliminate implied conditions.