[ConstraintElimination] Add Decomposition class (NFCI)

[nikic]

Summary

Replace the vector of DecompEntry with a class that stores the constant offset separately. I think this is cleaner than giving the first element special handling.

This probably also fixes some potential ubsan errors by more consistently using addWithOverflow/multiplyWithOverflow.

Depends on D137847.

[fhahn]

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

196

Could this be a class instead of a struct to ensure the offsets are only adjusted using the safe methods?

[nikic]

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

196

What would we do with the code in ConstraintInfo::getConstraint, which actually uses the decomposition in that case? Read-only getters?

[fhahn]

LGTM, thanks! This is a nice improvement, we can iterate over making it a class as follow-up. It might be good to update title/description to say struct instead of class.

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

196

Yeah that would be an option, but there might be a better API. I can play around a bit once this lands as struct.

[nikic]

Merged as https://github.com/llvm/llvm-project/commit/e82b5b5bbd1a97c289988a9bf2007511d35eb5cf, forgot the differential revision line.

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

 
 STATISTIC(NumCondsRemoved, "Number of instructions removed");
 DEBUG_COUNTER(EliminatedCounter, "conds-eliminated",
               "Controls which conditions are eliminated");
 
 static int64_t MaxConstraintValue = std::numeric_limits<int64_t>::max();
 static int64_t MinSignedConstraintValue = std::numeric_limits<int64_t>::min();
 
+// A helper to multiply 2 signed integers where overflowing is allowed.
+static int64_t multiplyWithOverflow(int64_t A, int64_t B) {
+  int64_t Result;
+  MulOverflow(A, B, Result);
+  return Result;
+}
+
+// A helper to add 2 signed integers where overflowing is allowed.
+static int64_t addWithOverflow(int64_t A, int64_t B) {
+  int64_t Result;
+  AddOverflow(A, B, Result);
+  return Result;
+}
+
 namespace {
 
 class ConstraintInfo;
 
 struct StackEntry {
   unsigned NumIn;
   unsigned NumOut;
   bool IsSigned = false;
struct DecompEntry {
   bool IsKnownPositive;
 
   DecompEntry(int64_t Coefficient, Value *Variable,
               bool IsKnownPositive = false)
       : Coefficient(Coefficient), Variable(Variable),
         IsKnownPositive(IsKnownPositive) {}
 };
 
+/// Represents an Offset + Coefficient1 * Variable1 + ... decomposition.
+struct Decomposition {
+  int64_t Offset = 0;
+  SmallVector<DecompEntry, 3> Vars;
+
+  Decomposition(int64_t Offset) : Offset(Offset) {}
+  Decomposition(Value *V, bool IsKnownPositive = false) {
+    Vars.emplace_back(1, V, IsKnownPositive);
+  }
+  Decomposition(int64_t Offset, ArrayRef<DecompEntry> Vars)
+      : Offset(Offset), Vars(Vars) {}
+
+  void add(int64_t OtherOffset) {
+    Offset = addWithOverflow(Offset, OtherOffset);
+  }
+
+  void add(const Decomposition &Other) {
+    add(Other.Offset);
+    append_range(Vars, Other.Vars);
+  }
+
+  void mul(int64_t Factor) {
+    Offset = multiplyWithOverflow(Offset, Factor);
+    for (auto &Var : Vars)
+      Var.Coefficient = multiplyWithOverflow(Var.Coefficient, Factor);
+  }
+};
+
 } // namespace
 
-static SmallVector<DecompEntry, 4>
-decompose(Value *V, SmallVector<PreconditionTy, 4> &Preconditions,
+static Decomposition decompose(Value *V,
+                               SmallVector<PreconditionTy, 4> &Preconditions,
                                bool IsSigned, const DataLayout &DL);
 
 static bool canUseSExt(ConstantInt *CI) {
   const APInt &Val = CI->getValue();
   return Val.sgt(MinSignedConstraintValue) && Val.slt(MaxConstraintValue);
 }
 
-// A helper to multiply 2 signed integers where overflowing is allowed.
-static int64_t multiplyWithOverflow(int64_t A, int64_t B) {
-  int64_t Result;
-  MulOverflow(A, B, Result);
-  return Result;
-}
-
-// A helper to add 2 signed integers where overflowing is allowed.
-static int64_t addWithOverflow(int64_t A, int64_t B) {
-  int64_t Result;
-  AddOverflow(A, B, Result);
-  return Result;
-}
-
-static SmallVector<DecompEntry, 4>
-decomposeGEP(GetElementPtrInst &GEP,
-             SmallVector<PreconditionTy, 4> &Preconditions, bool IsSigned,
-             const DataLayout &DL) {
+static Decomposition decomposeGEP(GetElementPtrInst &GEP,
+                                  SmallVector<PreconditionTy, 4> &Preconditions,
+                                  bool IsSigned, const DataLayout &DL) {
   // Do not reason about pointers where the index size is larger than 64 bits,
   // as the coefficients used to encode constraints are 64 bit integers.
   if (DL.getIndexTypeSizeInBits(GEP.getPointerOperand()->getType()) > 64)
-    return {{0, nullptr}, {1, &GEP}};
+    return &GEP;
 
   if (!GEP.isInBounds())
-    return {{0, nullptr}, {1, &GEP}};
+    return &GEP;
 
   // Handle the (gep (gep ....), C) case by incrementing the constant
   // coefficient of the inner GEP, if C is a constant.
   auto *InnerGEP = dyn_cast<GetElementPtrInst>(GEP.getPointerOperand());
   if (InnerGEP && GEP.getNumOperands() == 2 &&
       isa<ConstantInt>(GEP.getOperand(1))) {
     APInt Offset = cast<ConstantInt>(GEP.getOperand(1))->getValue();
     auto Result = decompose(InnerGEP, Preconditions, IsSigned, DL);
 
     auto GTI = gep_type_begin(GEP);
     // Bail out for scalable vectors for now.
     if (isa<ScalableVectorType>(GTI.getIndexedType()))
-      return {{0, nullptr}, {1, &GEP}};
+      return &GEP;
     int64_t Scale = static_cast<int64_t>(
         DL.getTypeAllocSize(GTI.getIndexedType()).getFixedSize());
 
-    Result[0].Coefficient =
-        addWithOverflow(Result[0].Coefficient,
-                        multiplyWithOverflow(Scale, Offset.getSExtValue()));
+    Result.add(multiplyWithOverflow(Scale, Offset.getSExtValue()));
     if (Offset.isNegative()) {
       // Add pre-condition ensuring the GEP is increasing monotonically and
       // can be de-composed.
       Preconditions.emplace_back(
           CmpInst::ICMP_SGE, InnerGEP->getOperand(1),
           ConstantInt::get(InnerGEP->getOperand(1)->getType(),
                            -1 * Offset.getSExtValue()));
     }
     return Result;
   }
 
-  SmallVector<DecompEntry, 4> Result = {{0, nullptr},
-                                        {1, GEP.getPointerOperand()}};
+  Decomposition Result = GEP.getPointerOperand();
   gep_type_iterator GTI = gep_type_begin(GEP);
   for (User::const_op_iterator I = GEP.op_begin() + 1, E = GEP.op_end(); I != E;
        ++I, ++GTI) {
     Value *Index = *I;
 
     // Bail out for scalable vectors for now.
     if (isa<ScalableVectorType>(GTI.getIndexedType()))
-      return {{0, nullptr}, {1, &GEP}};
+      return &GEP;
 
     // Struct indices must be constants (and reference an existing field). Add
     // them to the constant factor.
     if (StructType *STy = GTI.getStructTypeOrNull()) {
       // For a struct, add the member offset.
       unsigned FieldNo = cast<ConstantInt>(Index)->getZExtValue();
       if (FieldNo == 0)
         continue;
 
       // Add offset to constant factor.
-      Result[0].Coefficient = addWithOverflow(
-          Result[0].Coefficient,
-          int64_t(DL.getStructLayout(STy)->getElementOffset(FieldNo)));
+      Result.add(int64_t(DL.getStructLayout(STy)->getElementOffset(FieldNo)));
       continue;
     }
 
     // For an array/pointer, add the element offset, explicitly scaled.
     unsigned Scale = DL.getTypeAllocSize(GTI.getIndexedType()).getFixedSize();
 
     auto IdxResult = decompose(Index, Preconditions, IsSigned, DL);
-    for (auto &KV : IdxResult)
-      KV.Coefficient = multiplyWithOverflow(KV.Coefficient, Scale);
-    Result[0].Coefficient += IdxResult[0].Coefficient;
-    append_range(Result, ArrayRef<DecompEntry>(IdxResult).drop_front());
+    IdxResult.mul(Scale);
+    Result.add(IdxResult);
 
     // If Op0 is signed non-negative, the GEP is increasing monotonically and
     // can be de-composed.
     if (!isKnownNonNegative(Index, DL, /*Depth=*/MaxAnalysisRecursionDepth - 1))
       Preconditions.emplace_back(CmpInst::ICMP_SGE, Index,
                                  ConstantInt::get(Index->getType(), 0));
   }
   return Result;
 }
 
 // Decomposes \p V into a vector of entries of the form { Coefficient, Variable
 // } where Coefficient * Variable. The sum of the pairs equals \p V.  The first
 // pair is the constant-factor and X must be nullptr. If the expression cannot
 // be decomposed, returns an empty vector.
-static SmallVector<DecompEntry, 4>
-decompose(Value *V, SmallVector<PreconditionTy, 4> &Preconditions,
+static Decomposition decompose(Value *V,
+                               SmallVector<PreconditionTy, 4> &Preconditions,
                                bool IsSigned, const DataLayout &DL) {
 
-  auto MergeResults = [&Preconditions, IsSigned,
-                       &DL](Value *A, Value *B,
-                            bool IsSignedB) -> SmallVector<DecompEntry, 4> {
+  auto MergeResults = [&Preconditions, IsSigned, &DL](Value *A, Value *B,
+                                                      bool IsSignedB) {
     auto ResA = decompose(A, Preconditions, IsSigned, DL);
     auto ResB = decompose(B, Preconditions, IsSignedB, DL);
-    ResA[0].Coefficient += ResB[0].Coefficient;
-    append_range(ResA, drop_begin(ResB));
+    ResA.add(ResB);
     return ResA;
   };
 
   // Decompose \p V used with a signed predicate.
   if (IsSigned) {
     if (auto *CI = dyn_cast<ConstantInt>(V)) {
       if (canUseSExt(CI))
-        return {{CI->getSExtValue(), nullptr}};
+        return CI->getSExtValue();
     }
     Value *Op0;
     Value *Op1;
     if (match(V, m_NSWAdd(m_Value(Op0), m_Value(Op1))))
       return MergeResults(Op0, Op1, IsSigned);
 
-    return {{0, nullptr}, {1, V}};
+    return V;
   }
 
   if (auto *CI = dyn_cast<ConstantInt>(V)) {
     if (CI->uge(MaxConstraintValue))
-      return {{0, nullptr}, {1, V}};
-    return {{int64_t(CI->getZExtValue()), nullptr}};
+      return V;
+    return int64_t(CI->getZExtValue());
   }
 
   if (auto *GEP = dyn_cast<GetElementPtrInst>(V))
     return decomposeGEP(*GEP, Preconditions, IsSigned, DL);
 
   Value *Op0;
   bool IsKnownPositive = false;
   if (match(V, m_ZExt(m_Value(Op0)))) {
Preconditions.emplace_back(
         CmpInst::ICMP_UGE, Op0,
         ConstantInt::get(Op0->getType(), CI->getSExtValue() * -1));
     return MergeResults(Op0, CI, true);
   }
 
   if (match(V, m_NUWShl(m_Value(Op1), m_ConstantInt(CI))) && canUseSExt(CI)) {
     int64_t Mult = int64_t(std::pow(int64_t(2), CI->getSExtValue()));
     auto Result = decompose(Op1, Preconditions, IsSigned, DL);
-    for (auto &KV : Result)
-      KV.Coefficient *= Mult;
+    Result.mul(Mult);
     return Result;
   }
 
   if (match(V, m_NUWMul(m_Value(Op1), m_ConstantInt(CI))) && canUseSExt(CI) &&
       (!CI->isNegative())) {
     auto Result = decompose(Op1, Preconditions, IsSigned, DL);
-    for (auto &KV : Result)
-      KV.Coefficient *= CI->getSExtValue();
+    Result.mul(CI->getSExtValue());
     return Result;
   }
 
   if (match(V, m_NUWSub(m_Value(Op0), m_ConstantInt(CI))) && canUseSExt(CI))
-    return {{-1 * CI->getSExtValue(), nullptr}, {1, Op0}};
+    return {-1 * CI->getSExtValue(), {{1, Op0}}};
   if (match(V, m_NUWSub(m_Value(Op0), m_Value(Op1))))
-    return {{0, nullptr}, {1, Op0}, {-1, Op1}};
+    return {0, {{1, Op0}, {-1, Op1}}};
 
-  return {{0, nullptr}, {1, V, IsKnownPositive}};
+  return {V, IsKnownPositive};
 }
 
 ConstraintTy
 ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
                               SmallVectorImpl<Value *> &NewVariables) const {
   assert(NewVariables.empty() && "NewVariables must be empty when passed in");
   bool IsEq = false;
   // Try to convert Pred to one of ULE/SLT/SLE/SLT.
ConstraintInfo::getConstraint(CmpInst::Predicate Pred, Value *Op0, Value *Op1,
 
   SmallVector<PreconditionTy, 4> Preconditions;
   bool IsSigned = CmpInst::isSigned(Pred);
   auto &Value2Index = getValue2Index(IsSigned);
   auto ADec = decompose(Op0->stripPointerCastsSameRepresentation(),
                         Preconditions, IsSigned, DL);
   auto BDec = decompose(Op1->stripPointerCastsSameRepresentation(),
                         Preconditions, IsSigned, DL);
-  int64_t Offset1 = ADec[0].Coefficient;
-  int64_t Offset2 = BDec[0].Coefficient;
+  int64_t Offset1 = ADec.Offset;
+  int64_t Offset2 = BDec.Offset;
   Offset1 *= -1;
 
-  // Create iterator ranges that skip the constant-factor.
-  auto VariablesA = llvm::drop_begin(ADec);
-  auto VariablesB = llvm::drop_begin(BDec);
+  auto &VariablesA = ADec.Vars;
+  auto &VariablesB = BDec.Vars;
 
   // First try to look up \p V in Value2Index and NewVariables. Otherwise add a
   // new entry to NewVariables.
   DenseMap<Value *, unsigned> NewIndexMap;
   auto GetOrAddIndex = [&Value2Index, &NewVariables,
                         &NewIndexMap](Value *V) -> unsigned {
     auto V2I = Value2Index.find(V);
     if (V2I != Value2Index.end())