Add support for getelementptr nusw and nuw

nikic · nikic · commit b74cc982c95f · 2024-05-21T09:17:06.000+02:00
diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
@@ -11184,6 +11184,8 @@ Syntax:
 
       <result> = getelementptr <ty>, ptr <ptrval>{, <ty> <idx>}*
       <result> = getelementptr inbounds <ty>, ptr <ptrval>{, <ty> <idx>}*
+      <result> = getelementptr nusw <ty>, ptr <ptrval>{, <ty> <idx>}*
+      <result> = getelementptr nuw <ty>, ptr <ptrval>{, <ty> <idx>}*
       <result> = getelementptr inrange(S,E) <ty>, ptr <ptrval>{, <ty> <idx>}*
       <result> = getelementptr <ty>, <N x ptr> <ptrval>, <vector index type> <idx>
 
@@ -11299,27 +11301,46 @@ memory though, even if it happens to point into allocated storage. See the
 :ref:`Pointer Aliasing Rules <pointeraliasing>` section for more
 information.
 
-If the ``inbounds`` keyword is present, the result value of a
-``getelementptr`` with any non-zero indices is a
-:ref:`poison value <poisonvalues>` if one of the following rules is violated:
-
-*  The base pointer has an *in bounds* address of an allocated object, which
+The ``getelementptr`` instruction may have a number of attributes that impose
+additional rules. If any of the rules are violated, the result value is a
+:ref:`poison value <poisonvalues>`. In cases where the base is a vector of
+pointers, the attributes apply to each computation element-wise.
+
+For ``nusw`` (no unsigned signed wrap):
+
+ * If the type of an index is larger than the pointer index type, the
+   truncation to the pointer index type preserves the signed value
+   (``trunc nsw``).
+ * The multiplication of an index by the type size does not wrap the pointer
+   index type in a signed sense (``mul nsw``).
+ * The successive addition of each offset (without adding the base address)
+   does not wrap the pointer index type in a signed sense (``add nsw``).
+ * The successive addition of the current address, truncated to the index type
+   and interpreted as an unsigned number, and each offset, interpreted as
+   a signed number, does not wrap the index type.
+
+For ``nuw`` (no unsigned wrap):
+
+ * If the type of an index is larger than the pointer index type, the
+   truncation to the pointer index type preserves the unsigned value
+   (``trunc nuw``).
+ * The multiplication of an index by the type size does not wrap the pointer
+   index type in an unsigned sense (``mul nuw``).
+ * The successive addition of each offset (without adding the base address)
+   does not wrap the pointer index type in an unsigned sense (``add nuw``).
+ * The successive addition of the current address, truncated to the index type
+   and interpreted as an unsigned number, and each offset, also interpreted as
+   an unsigned number, does not wrap the index type (``add nuw``).
+
+For ``inbounds`` all rules of the ``nusw`` attribute apply. Additionally,
+if the ``getelementptr`` has any non-zero indices, the following rules apply:
+
+ * The base pointer has an *in bounds* address of an allocated object, which
    means that it points into an allocated object, or to its end. Note that the
    object does not have to be live anymore; being in-bounds of a deallocated
    object is sufficient.
-*  If the type of an index is larger than the pointer index type, the
-   truncation to the pointer index type preserves the signed value.
-*  The multiplication of an index by the type size does not wrap the pointer
-   index type in a signed sense (``nsw``).
-*  The successive addition of each offset (without adding the base address) does
-   not wrap the pointer index type in a signed sense (``nsw``).
-*  The successive addition of the current address, interpreted as an unsigned
-   number, and each offset, interpreted as a signed number, does not wrap the
-   unsigned address space and remains *in bounds* of the allocated object.
-   As a corollary, if the added offset is non-negative, the addition does not
-   wrap in an unsigned sense (``nuw``).
-*  In cases where the base is a vector of pointers, the ``inbounds`` keyword
-   applies to each of the computations element-wise.
+ * During the successive addition of offsets to the address, the resulting
+   pointer must remain *in bounds* of the allocated object at each step.
 
 Note that ``getelementptr`` with all-zero indices is always considered to be
 ``inbounds``, even if the base pointer does not point to an allocated object.
diff --git a/llvm/docs/ReleaseNotes.rst b/llvm/docs/ReleaseNotes.rst
@@ -51,6 +51,7 @@ Changes to the LLVM IR
 ----------------------
 
 * Added Memory Model Relaxation Annotations (MMRAs).
+* Added ``nusw`` and ``nuw`` flags to ``getelementptr`` instruction.
 * Renamed ``llvm.experimental.vector.reverse`` intrinsic to ``llvm.vector.reverse``.
 * Renamed ``llvm.experimental.vector.splice`` intrinsic to ``llvm.vector.splice``.
 * Renamed ``llvm.experimental.vector.interleave2`` intrinsic to ``llvm.vector.interleave2``.
diff --git a/llvm/include/llvm/AsmParser/LLToken.h b/llvm/include/llvm/AsmParser/LLToken.h
@@ -109,6 +109,7 @@ enum Kind {
   kw_fast,
   kw_nuw,
   kw_nsw,
+  kw_nusw,
   kw_exact,
   kw_disjoint,
   kw_inbounds,
diff --git a/llvm/include/llvm/Bitcode/LLVMBitCodes.h b/llvm/include/llvm/Bitcode/LLVMBitCodes.h
@@ -524,6 +524,14 @@ enum PossiblyExactOperatorOptionalFlags { PEO_EXACT = 0 };
 /// PossiblyDisjointInst's SubclassOptionalData contents.
 enum PossiblyDisjointInstOptionalFlags { PDI_DISJOINT = 0 };
 
+/// GetElementPtrOptionalFlags - Flags for serializing
+/// GEPOperator's SubclassOptionalData contents.
+enum GetElementPtrOptionalFlags {
+  GEP_INBOUNDS = 0,
+  GEP_NUSW = 1,
+  GEP_NUW = 2,
+};
+
 /// Encoded AtomicOrdering values.
 enum AtomicOrderingCodes {
   ORDERING_NOTATOMIC = 0,
diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
@@ -1171,9 +1171,23 @@ class GetElementPtrInst : public Instruction {
   /// See LangRef.html for the meaning of inbounds on a getelementptr.
   void setIsInBounds(bool b = true);
 
+  /// Set or clear the nusw flag on this GEP instruction.
+  /// See LangRef.html for the meaning of nusw on a getelementptr.
+  void setHasNoUnsignedSignedWrap(bool B = true);
+
+  /// Set or clear the nuw flag on this GEP instruction.
+  /// See LangRef.html for the meaning of nuw on a getelementptr.
+  void setHasNoUnsignedWrap(bool B = true);
+
   /// Determine whether the GEP has the inbounds flag.
   bool isInBounds() const;
 
+  /// Determine whether the GEP has the nusw flag.
+  bool hasNoUnsignedSignedWrap() const;
+
+  /// Determine whether the GEP has the nuw flag.
+  bool hasNoUnsignedWrap() const;
+
   /// Accumulate the constant address offset of this GEP if possible.
   ///
   /// This routine accepts an APInt into which it will accumulate the constant
diff --git a/llvm/include/llvm/IR/Operator.h b/llvm/include/llvm/IR/Operator.h
@@ -405,11 +405,27 @@ class GEPOperator
 
   enum {
     IsInBounds = (1 << 0),
+    HasNoUnsignedSignedWrap = (1 << 1),
+    HasNoUnsignedWrap = (1 << 2),
   };
 
   void setIsInBounds(bool B) {
+    // Also set nusw when inbounds is set.
+    SubclassOptionalData = (SubclassOptionalData & ~IsInBounds) |
+                           (B * (IsInBounds | HasNoUnsignedSignedWrap));
+  }
+
+  void setHasNoUnsignedSignedWrap(bool B) {
+    // Also unset inbounds when nusw is unset.
+    if (B)
+      SubclassOptionalData |= HasNoUnsignedSignedWrap;
+    else
+      SubclassOptionalData &= ~(IsInBounds | HasNoUnsignedSignedWrap);
+  }
+
+  void setHasNoUnsignedWrap(bool B) {
     SubclassOptionalData =
-      (SubclassOptionalData & ~IsInBounds) | (B * IsInBounds);
+        (SubclassOptionalData & ~HasNoUnsignedWrap) | (B * HasNoUnsignedWrap);
   }
 
 public:
@@ -421,6 +437,14 @@ class GEPOperator
     return SubclassOptionalData & IsInBounds;
   }
 
+  bool hasNoUnsignedSignedWrap() const {
+    return SubclassOptionalData & HasNoUnsignedSignedWrap;
+  }
+
+  bool hasNoUnsignedWrap() const {
+    return SubclassOptionalData & HasNoUnsignedWrap;
+  }
+
   /// Returns the offset of the index with an inrange attachment, or
   /// std::nullopt if none.
   std::optional<ConstantRange> getInRange() const;
diff --git a/llvm/lib/AsmParser/LLLexer.cpp b/llvm/lib/AsmParser/LLLexer.cpp
@@ -566,6 +566,7 @@ lltok::Kind LLLexer::LexIdentifier() {
   KEYWORD(fast);
   KEYWORD(nuw);
   KEYWORD(nsw);
+  KEYWORD(nusw);
   KEYWORD(exact);
   KEYWORD(disjoint);
   KEYWORD(inbounds);
diff --git a/llvm/lib/AsmParser/LLParser.cpp b/llvm/lib/AsmParser/LLParser.cpp
@@ -8340,7 +8340,17 @@ int LLParser::parseGetElementPtr(Instruction *&Inst, PerFunctionState &PFS) {
   Value *Val = nullptr;
   LocTy Loc, EltLoc;
 
-  bool InBounds = EatIfPresent(lltok::kw_inbounds);
+  bool InBounds = false, NUSW = false, NUW = false;
+  while (true) {
+    if (EatIfPresent(lltok::kw_inbounds))
+      InBounds = true;
+    else if (EatIfPresent(lltok::kw_nusw))
+      NUSW = true;
+    else if (EatIfPresent(lltok::kw_nuw))
+      NUW = true;
+    else
+      break;
+  }
 
   Type *Ty = nullptr;
   if (parseType(Ty) ||
@@ -8393,9 +8403,14 @@ int LLParser::parseGetElementPtr(Instruction *&Inst, PerFunctionState &PFS) {
 
   if (!GetElementPtrInst::getIndexedType(Ty, Indices))
     return error(Loc, "invalid getelementptr indices");
-  Inst = GetElementPtrInst::Create(Ty, Ptr, Indices);
+  GetElementPtrInst *GEP = GetElementPtrInst::Create(Ty, Ptr, Indices);
+  Inst = GEP;
   if (InBounds)
-    cast<GetElementPtrInst>(Inst)->setIsInBounds(true);
+    GEP->setIsInBounds(true);
+  if (NUSW)
+    GEP->setHasNoUnsignedSignedWrap(true);
+  if (NUW)
+    GEP->setHasNoUnsignedWrap(true);
   return AteExtraComma ? InstExtraComma : InstNormal;
 }
 
diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -5062,10 +5062,17 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
 
       unsigned TyID;
       Type *Ty;
-      bool InBounds;
+      bool InBounds = false, NUSW = false, NUW = false;
 
       if (BitCode == bitc::FUNC_CODE_INST_GEP) {
-        InBounds = Record[OpNum++];
+        uint64_t Flags = Record[OpNum++];
+        if (Flags & (1 << bitc::GEP_INBOUNDS))
+          InBounds = true;
+        if (Flags & (1 << bitc::GEP_NUSW))
+          NUSW = true;
+        if (Flags & (1 << bitc::GEP_NUW))
+          NUW = true;
+
         TyID = Record[OpNum++];
         Ty = getTypeByID(TyID);
       } else {
@@ -5096,7 +5103,8 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
         GEPIdx.push_back(Op);
       }
 
-      I = GetElementPtrInst::Create(Ty, BasePtr, GEPIdx);
+      auto *GEP = GetElementPtrInst::Create(Ty, BasePtr, GEPIdx);
+      I = GEP;
 
       ResTypeID = TyID;
       if (cast<GEPOperator>(I)->getNumIndices() != 0) {
@@ -5123,7 +5131,11 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
 
       InstructionList.push_back(I);
       if (InBounds)
-        cast<GetElementPtrInst>(I)->setIsInBounds(true);
+        GEP->setIsInBounds(true);
+      if (NUSW)
+        GEP->setHasNoUnsignedSignedWrap(true);
+      if (NUW)
+        GEP->setHasNoUnsignedWrap(true);
       break;
     }
 
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -2961,7 +2961,14 @@ void ModuleBitcodeWriter::writeInstruction(const Instruction &I,
     Code = bitc::FUNC_CODE_INST_GEP;
     AbbrevToUse = FUNCTION_INST_GEP_ABBREV;
     auto &GEPInst = cast<GetElementPtrInst>(I);
-    Vals.push_back(GEPInst.isInBounds());
+    uint64_t Flags = 0;
+    if (GEPInst.isInBounds())
+      Flags |= 1 << bitc::GEP_INBOUNDS;
+    if (GEPInst.hasNoUnsignedSignedWrap())
+      Flags |= 1 << bitc::GEP_NUSW;
+    if (GEPInst.hasNoUnsignedWrap())
+      Flags |= 1 << bitc::GEP_NUW;
+    Vals.push_back(Flags);
     Vals.push_back(VE.getTypeID(GEPInst.getSourceElementType()));
     for (unsigned i = 0, e = I.getNumOperands(); i != e; ++i)
       pushValueAndType(I.getOperand(i), InstID, Vals);
@@ -3859,7 +3866,7 @@ void ModuleBitcodeWriter::writeBlockInfo() {
   {
     auto Abbv = std::make_shared<BitCodeAbbrev>();
     Abbv->Add(BitCodeAbbrevOp(bitc::FUNC_CODE_INST_GEP));
-    Abbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Fixed, 1));
+    Abbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Fixed, 3));
     Abbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Fixed, // dest ty
                               Log2_32_Ceil(VE.getTypes().size() + 1)));
     Abbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Array));
diff --git a/llvm/lib/IR/AsmWriter.cpp b/llvm/lib/IR/AsmWriter.cpp
@@ -1417,6 +1417,10 @@ static void WriteOptimizationInfo(raw_ostream &Out, const User *U) {
   } else if (const GEPOperator *GEP = dyn_cast<GEPOperator>(U)) {
     if (GEP->isInBounds())
       Out << " inbounds";
+    else if (GEP->hasNoUnsignedSignedWrap())
+      Out << " nusw";
+    if (GEP->hasNoUnsignedWrap())
+      Out << " nuw";
     if (auto InRange = GEP->getInRange()) {
       Out << " inrange(" << InRange->getLower() << ", " << InRange->getUpper()
           << ")";
diff --git a/llvm/lib/IR/Instruction.cpp b/llvm/lib/IR/Instruction.cpp
@@ -442,6 +442,8 @@ void Instruction::dropPoisonGeneratingFlags() {
 
   case Instruction::GetElementPtr:
     cast<GetElementPtrInst>(this)->setIsInBounds(false);
+    cast<GetElementPtrInst>(this)->setHasNoUnsignedSignedWrap(false);
+    cast<GetElementPtrInst>(this)->setHasNoUnsignedWrap(false);
     break;
 
   case Instruction::UIToFP:
@@ -658,9 +660,15 @@ void Instruction::copyIRFlags(const Value *V, bool IncludeWrapFlags) {
     if (isa<FPMathOperator>(this))
       copyFastMathFlags(FP->getFastMathFlags());
 
-  if (auto *SrcGEP = dyn_cast<GetElementPtrInst>(V))
-    if (auto *DestGEP = dyn_cast<GetElementPtrInst>(this))
+  if (auto *SrcGEP = dyn_cast<GetElementPtrInst>(V)) {
+    if (auto *DestGEP = dyn_cast<GetElementPtrInst>(this)) {
       DestGEP->setIsInBounds(SrcGEP->isInBounds() || DestGEP->isInBounds());
+      DestGEP->setHasNoUnsignedSignedWrap(SrcGEP->hasNoUnsignedSignedWrap() ||
+                                          DestGEP->hasNoUnsignedSignedWrap());
+      DestGEP->setHasNoUnsignedWrap(SrcGEP->hasNoUnsignedWrap() ||
+                                    DestGEP->hasNoUnsignedWrap());
+    }
+  }
 
   if (auto *NNI = dyn_cast<PossiblyNonNegInst>(V))
     if (isa<PossiblyNonNegInst>(this))
@@ -698,9 +706,15 @@ void Instruction::andIRFlags(const Value *V) {
     }
   }
 
-  if (auto *SrcGEP = dyn_cast<GetElementPtrInst>(V))
-    if (auto *DestGEP = dyn_cast<GetElementPtrInst>(this))
+  if (auto *SrcGEP = dyn_cast<GetElementPtrInst>(V)) {
+    if (auto *DestGEP = dyn_cast<GetElementPtrInst>(this)) {
       DestGEP->setIsInBounds(SrcGEP->isInBounds() && DestGEP->isInBounds());
+      DestGEP->setHasNoUnsignedSignedWrap(SrcGEP->hasNoUnsignedSignedWrap() &&
+                                          DestGEP->hasNoUnsignedSignedWrap());
+      DestGEP->setHasNoUnsignedWrap(SrcGEP->hasNoUnsignedWrap() &&
+                                    DestGEP->hasNoUnsignedWrap());
+    }
+  }
 
   if (auto *NNI = dyn_cast<PossiblyNonNegInst>(V))
     if (isa<PossiblyNonNegInst>(this))
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
@@ -2047,10 +2047,26 @@ void GetElementPtrInst::setIsInBounds(bool B) {
   cast<GEPOperator>(this)->setIsInBounds(B);
 }
 
+void GetElementPtrInst::setHasNoUnsignedSignedWrap(bool B) {
+  cast<GEPOperator>(this)->setHasNoUnsignedSignedWrap(B);
+}
+
+void GetElementPtrInst::setHasNoUnsignedWrap(bool B) {
+  cast<GEPOperator>(this)->setHasNoUnsignedWrap(B);
+}
+
 bool GetElementPtrInst::isInBounds() const {
   return cast<GEPOperator>(this)->isInBounds();
 }
 
+bool GetElementPtrInst::hasNoUnsignedSignedWrap() const {
+  return cast<GEPOperator>(this)->hasNoUnsignedSignedWrap();
+}
+
+bool GetElementPtrInst::hasNoUnsignedWrap() const {
+  return cast<GEPOperator>(this)->hasNoUnsignedWrap();
+}
+
 bool GetElementPtrInst::accumulateConstantOffset(const DataLayout &DL,
                                                  APInt &Offset) const {
   // Delegate to the generic GEPOperator implementation.
diff --git a/llvm/lib/IR/Operator.cpp b/llvm/lib/IR/Operator.cpp
@@ -42,7 +42,8 @@ bool Operator::hasPoisonGeneratingFlags() const {
   case Instruction::GetElementPtr: {
     auto *GEP = cast<GEPOperator>(this);
     // Note: inrange exists on constexpr only
-    return GEP->isInBounds() || GEP->getInRange() != std::nullopt;
+    return GEP->isInBounds() || GEP->hasNoUnsignedSignedWrap() ||
+           GEP->hasNoUnsignedWrap() || GEP->getInRange() != std::nullopt;
   }
   case Instruction::UIToFP:
   case Instruction::ZExt:
diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -1111,6 +1111,9 @@ bool SeparateConstOffsetFromGEP::splitGEP(GetElementPtrInst *GEP) {
   // possible. GEPs with inbounds are more friendly to alias analysis.
   bool GEPWasInBounds = GEP->isInBounds();
   GEP->setIsInBounds(false);
+  // TODO(gep_nowrap): Try to preserve these.
+  GEP->setHasNoUnsignedSignedWrap(false);
+  GEP->setHasNoUnsignedWrap(false);
 
   // Lowers a GEP to either GEPs with a single index or arithmetic operations.
   if (LowerGEP) {
@@ -1386,6 +1389,11 @@ void SeparateConstOffsetFromGEP::swapGEPOperand(GetElementPtrInst *First,
      Offset.ugt(ObjectSize)) {
     First->setIsInBounds(false);
     Second->setIsInBounds(false);
+    // TODO(gep_nowrap): Make flag preservation more precise.
+    First->setHasNoUnsignedSignedWrap(false);
+    Second->setHasNoUnsignedSignedWrap(false);
+    First->setHasNoUnsignedWrap(false);
+    Second->setHasNoUnsignedWrap(false);
   } else
     First->setIsInBounds(true);
 }
diff --git a/llvm/lib/Transforms/Utils/FunctionComparator.cpp b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
diff --git a/llvm/test/Assembler/flags.ll b/llvm/test/Assembler/flags.ll
diff --git a/llvm/test/Transforms/InstCombine/freeze.ll b/llvm/test/Transforms/InstCombine/freeze.ll
diff --git a/llvm/test/Transforms/SimplifyCFG/HoistCode.ll b/llvm/test/Transforms/SimplifyCFG/HoistCode.ll
diff --git a/llvm/test/tools/llvm-reduce/reduce-flags.ll b/llvm/test/tools/llvm-reduce/reduce-flags.ll
diff --git a/llvm/tools/llvm-reduce/deltas/ReduceInstructionFlags.cpp b/llvm/tools/llvm-reduce/deltas/ReduceInstructionFlags.cpp
