[IR] Add NUW/NSW to GetElementPtr. (WIP)

[fhahn]

Summary

Prototype for
https://discourse.llvm.org/t/retain-fact-that-index-offsets-for-gep-are-non-negative-for-size-t-uint64-t-offsets

clang/lib/CodeGen/CGBuilder.h

public:
   /// inbounds operation.
   using CGBuilderBaseTy::CreateStructGEP;
   Address CreateStructGEP(Address Addr, unsigned Index,
                           const llvm::Twine &Name = "") {
     llvm::StructType *ElTy = cast<llvm::StructType>(Addr.getElementType());
     const llvm::DataLayout &DL = BB->getParent()->getParent()->getDataLayout();
     const llvm::StructLayout *Layout = DL.getStructLayout(ElTy);
     auto Offset = CharUnits::fromQuantity(Layout->getElementOffset(Index));
-
-    return Address(
-        CreateStructGEP(Addr.getElementType(), Addr.getPointer(), Index, Name),
-        ElTy->getElementType(Index),
+    auto *GEP =
+        CreateStructGEP(Addr.getElementType(), Addr.getPointer(), Index, Name);
+    cast<llvm::GetElementPtrInst>(GEP)->setHasNoUnsignedWrap(true);
+    return Address(GEP, ElTy->getElementType(Index),
                    Addr.getAlignment().alignmentAtOffset(Offset));
   }
 
   /// Given
   ///   %addr = [n x T]* ...
   /// produce
   ///   %name = getelementptr inbounds %addr, i64 0, i64 index
   /// where i64 is actually the target word size.
   ///

clang/lib/CodeGen/CGExprScalar.cpp

 
 Value *
 CodeGenFunction::EmitCheckedInBoundsGEP(llvm::Type *ElemTy, Value *Ptr,
                                         ArrayRef<Value *> IdxList,
                                         bool SignedIndices, bool IsSubtraction,
                                         SourceLocation Loc, const Twine &Name) {
   llvm::Type *PtrTy = Ptr->getType();
   Value *GEPVal = Builder.CreateInBoundsGEP(ElemTy, Ptr, IdxList, Name);
+  if (!SignedIndices && !IsSubtraction)
+    cast<llvm::GetElementPtrInst>(GEPVal)->setHasNoUnsignedWrap(true);
 
   // If the pointer overflow sanitizer isn't enabled, do nothing.
   if (!SanOpts.has(SanitizerKind::PointerOverflow))
     return GEPVal;
 
   // Perform nullptr-and-offset check unless the nullptr is defined.
   bool PerformNullCheck = !NullPointerIsDefined(
       Builder.GetInsertBlock()->getParent(), PtrTy->getPointerAddressSpace());

llvm/include/llvm/IR/Instructions.h

public:
 
   /// Set or clear the inbounds flag on this GEP instruction.
   /// See LangRef.html for the meaning of inbounds on a getelementptr.
   void setIsInBounds(bool b = true);
 
   /// Determine whether the GEP has the inbounds flag.
   bool isInBounds() const;
 
+  void setHasNoUnsignedWrap(bool B);
+  bool hasNoUnsignedWrap() const;
+
+  void setHasNoSignedWrap(bool B);
+  bool hasNoSignedWrap() const;
+
   /// Accumulate the constant address offset of this GEP if possible.
   ///
   /// This routine accepts an APInt into which it will accumulate the constant
   /// offset of this GEP if the GEP is in fact constant. If the GEP is not
   /// all-constant, it returns false and the value of the offset APInt is
   /// undefined (it is *not* preserved!). The APInt passed into this routine
   /// must be at least as wide as the IntPtr type for the address space of
   /// the base GEP pointer.

llvm/include/llvm/IR/Operator.h

class GEPOperator
   : public ConcreteOperator<Operator, Instruction::GetElementPtr> {
   friend class GetElementPtrInst;
   friend class ConstantExpr;
 
   enum {
     IsInBounds = (1 << 0),
     // InRangeIndex: bits 1-6
   };
+  enum {
+    AnyWrap = 0,
+    HasNoUnsignedWrap = (1 << 0),
+    HasNoSignedWrap = (1 << 1),
+  };
 
   void setIsInBounds(bool B) {
     SubclassOptionalData =
       (SubclassOptionalData & ~IsInBounds) | (B * IsInBounds);
+    setHasNoSignedWrap(B);
+  }
+
+  void setHasNoUnsignedWrap(bool B) {
+    short D = getSubclassDataFromValue();
+    D = (D & ~HasNoUnsignedWrap) | (B * HasNoUnsignedWrap);
+    setValueSubclassData(D);
+  }
+
+  void setHasNoSignedWrap(bool B) {
+    short D = getSubclassDataFromValue();
+    D = (D & ~HasNoSignedWrap) | (B * HasNoSignedWrap);
+    setValueSubclassData(D);
   }
 
 public:
   /// Test whether this is an inbounds GEP, as defined by LangRef.html.
   bool isInBounds() const {
     return SubclassOptionalData & IsInBounds;
   }
 
   /// Returns the offset of the index with an inrange attachment, or None if
   /// none.
   Optional<unsigned> getInRangeIndex() const {
     if (SubclassOptionalData >> 1 == 0) return None;
     return (SubclassOptionalData >> 1) - 1;
   }
 
+  bool hasNoUnsignedWrap() const {
+    short D = getSubclassDataFromValue();
+    return D & HasNoUnsignedWrap;
+  }
+
+  bool hasNoSignedWrap() const {
+    short D = getSubclassDataFromValue();
+    return D & HasNoSignedWrap;
+  }
+
   inline op_iterator       idx_begin()       { return op_begin()+1; }
   inline const_op_iterator idx_begin() const { return op_begin()+1; }
   inline op_iterator       idx_end()         { return op_end(); }
   inline const_op_iterator idx_end()   const { return op_end(); }
 
   inline iterator_range<op_iterator> indices() {
     return make_range(idx_begin(), idx_end());
   }

llvm/lib/AsmParser/LLParser.cpp

bool LLParser::parseValID(ValID &ID, PerFunctionState *PFS, Type *ExpectedTy) {
   case lltok::kw_getelementptr:
   case lltok::kw_shufflevector:
   case lltok::kw_insertelement:
   case lltok::kw_extractelement:
   case lltok::kw_select: {
     unsigned Opc = Lex.getUIntVal();
     SmallVector<Constant*, 16> Elts;
     bool InBounds = false;
+    bool HasNoUnsignedWrap = false;
+    bool HasNoSignedWrap = false;
     Type *Ty;
     Lex.Lex();
 
-    if (Opc == Instruction::GetElementPtr)
+    if (Opc == Instruction::GetElementPtr) {
       InBounds = EatIfPresent(lltok::kw_inbounds);
+      HasNoUnsignedWrap = EatIfPresent(lltok::kw_nuw);
+      HasNoSignedWrap = EatIfPresent(lltok::kw_nsw);
+    }
 
     if (parseToken(lltok::lparen, "expected '(' in constantexpr"))
       return true;
 
     LocTy ExplicitTypeLoc = Lex.getLoc();
     if (Opc == Instruction::GetElementPtr) {
       if (parseType(Ty) ||
           parseToken(lltok::comma, "expected comma after getelementptr's type"))

 /// parseGetElementPtr
 ///   ::= 'getelementptr' 'inbounds'? TypeAndValue (',' TypeAndValue)*
 int LLParser::parseGetElementPtr(Instruction *&Inst, PerFunctionState &PFS) {
   Value *Ptr = nullptr;
   Value *Val = nullptr;
   LocTy Loc, EltLoc;
 
   bool InBounds = EatIfPresent(lltok::kw_inbounds);
+  bool HasNoUnsignedWrap = EatIfPresent(lltok::kw_nuw);
+  bool HasNoSignedWrap = EatIfPresent(lltok::kw_nsw);
 
   Type *Ty = nullptr;
   LocTy ExplicitTypeLoc = Lex.getLoc();
   if (parseType(Ty) ||
       parseToken(lltok::comma, "expected comma after getelementptr's type") ||
       parseTypeAndValue(Ptr, Loc, PFS))
     return true;
 
int LLParser::parseGetElementPtr(Instruction *&Inst, PerFunctionState &PFS) {
   if (!Indices.empty() && !Ty->isSized(&Visited))
     return error(Loc, "base element of getelementptr must be sized");
 
   if (!GetElementPtrInst::getIndexedType(Ty, Indices))
     return error(Loc, "invalid getelementptr indices");
   Inst = GetElementPtrInst::Create(Ty, Ptr, Indices);
   if (InBounds)
     cast<GetElementPtrInst>(Inst)->setIsInBounds(true);
+  if (HasNoUnsignedWrap)
+    cast<GetElementPtrInst>(Inst)->setHasNoUnsignedWrap(true);
+  if (HasNoSignedWrap)
+    cast<GetElementPtrInst>(Inst)->setHasNoSignedWrap(true);
+
   return AteExtraComma ? InstExtraComma : InstNormal;
 }
 
 /// parseExtractValue
 ///   ::= 'extractvalue' TypeAndValue (',' uint32)+
 int LLParser::parseExtractValue(Instruction *&Inst, PerFunctionState &PFS) {
   Value *Val; LocTy Loc;
   SmallVector<unsigned, 4> Indices;

llvm/lib/IR/AsmWriter.cpp

if (OBO->hasNoSignedWrap())
       Out << " nsw";
   } else if (const PossiblyExactOperator *Div =
                dyn_cast<PossiblyExactOperator>(U)) {
     if (Div->isExact())
       Out << " exact";
   } else if (const GEPOperator *GEP = dyn_cast<GEPOperator>(U)) {
     if (GEP->isInBounds())
       Out << " inbounds";
+    if (GEP->hasNoUnsignedWrap())
+      Out << " nuw";
+    /*   if (GEP->hasNoSignedWrap())*/
+    /*Out << " nsw";*/
   }
 }
 
 static void WriteConstantInternal(raw_ostream &Out, const Constant *CV,
                                   AsmWriterContext &WriterCtx) {
   if (const ConstantInt *CI = dyn_cast<ConstantInt>(CV)) {
     if (CI->getType()->isIntegerTy(1)) {
       Out << (CI->getZExtValue() ? "true" : "false");

llvm/lib/IR/Instructions.cpp

 void GetElementPtrInst::setIsInBounds(bool B) {
   cast<GEPOperator>(this)->setIsInBounds(B);
 }
 
 bool GetElementPtrInst::isInBounds() const {
   return cast<GEPOperator>(this)->isInBounds();
 }
 
+void GetElementPtrInst::setHasNoUnsignedWrap(bool B) {
+  cast<GEPOperator>(this)->setHasNoUnsignedWrap(B);
+}
+
+bool GetElementPtrInst::hasNoUnsignedWrap() const {
+  return cast<GEPOperator>(this)->hasNoUnsignedWrap();
+}
+
+void GetElementPtrInst::setHasNoSignedWrap(bool B) {
+  cast<GEPOperator>(this)->setHasNoSignedWrap(B);
+}
+
+bool GetElementPtrInst::hasNoSignedWrap() const {
+  return cast<GEPOperator>(this)->hasNoSignedWrap();
+}
+
 bool GetElementPtrInst::accumulateConstantOffset(const DataLayout &DL,
                                                  APInt &Offset) const {
   // Delegate to the generic GEPOperator implementation.
   return cast<GEPOperator>(this)->accumulateConstantOffset(DL, Offset);
 }
 
 bool GetElementPtrInst::collectOffset(
     const DataLayout &DL, unsigned BitWidth,

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

decomposeGEP(GetElementPtrInst &GEP,
              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.
   unsigned AS =
       cast<PointerType>(GEP.getPointerOperand()->getType())->getAddressSpace();
   if (DL.getIndexSizeInBits(AS) > 64)
     return {};
 
-  if (!GEP.isInBounds())
+  if (!GEP.hasNoUnsignedWrap() && !GEP.hasNoSignedWrap())
     return {{0, nullptr}, {1, &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 && InnerGEP->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 {};
     int64_t Scale = static_cast<int64_t>(
         DL.getTypeAllocSize(GTI.getIndexedType()).getFixedSize());
 
     Result[0].Coefficient =
         addWithOverflow(Result[0].Coefficient,
                         multiplyWithOverflow(Scale, Offset.getSExtValue()));
     if (Offset.isNegative()) {
+      if (InnerGEP->hasNoUnsignedWrap()) {
         // Add pre-condition ensuring the GEP is increasing monotonically and
         // can be de-composed.
         Preconditions.emplace_back(
+            CmpInst::ICMP_UGE, InnerGEP->getOperand(1),
+            ConstantInt::get(InnerGEP->getOperand(1)->getType(),
+                             -1 * Offset.getSExtValue()));
+      } else {
+        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()}};
   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) {
for (User::const_op_iterator I = GEP.op_begin() + 1, E = GEP.op_end(); I != E;
     } else {
       for (auto &KV : IdxResult)
         KV.Coefficient = multiplyWithOverflow(KV.Coefficient, Scale);
       Result[0].Coefficient += IdxResult[0].Coefficient;
       append_range(Result, ArrayRef<DecompEntry>(IdxResult).drop_front());
     }
     // If Op0 is signed non-negative, the GEP is increasing monotonically and
     // can be de-composed.
-    if (!isKnownNonNegative(Index, DL, /*Depth=*/MaxAnalysisRecursionDepth - 1))
+    if (!GEP.hasNoUnsignedWrap() &&
+        !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