[IR] Add getelementptr nusw and nuw flags (#90824)

This implements the `nusw` and `nuw` flags for `getelementptr` as
proposed at
https://discourse.llvm.org/t/rfc-add-nusw-and-nuw-flags-for-getelementptr/78672.

The three possible flags are encapsulated in the new `GEPNoWrapFlags`
class. Currently this class has a ctor from bool, interpreted as the
InBounds flag. This ctor should be removed in the future, as code gets
migrated to handle all flags.

There are a few places annotated with `TODO(gep_nowrap)`, where I've had
to touch code but opted to not infer or precisely preserve the new
flags, so as to keep this as NFC as possible and make sure any changes
of that kind get test coverage when they are made.

Author: nikic

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,47 @@ 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 pointer
+   index type and interpreted as an unsigned number, and each offset,
+   interpreted as a signed number, does not wrap the pointer 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 pointer
+   index type and interpreted as an unsigned number, and each offset, also
+   interpreted as an unsigned number, does not wrap the pointer 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.
@@ -11330,6 +11352,10 @@ These rules are based on the assumption that no allocated object may cross
 the unsigned address space boundary, and no allocated object may be larger
 than half the pointer index type space.
 
+If ``inbounds`` is present on a ``getelementptr`` instruction, the ``nusw``
+attribute will be automatically set as well. For this reason, the ``nusw``
+will also not be printed in textual IR if ``inbounds`` is already present.
+
 If the ``inrange(Start, End)`` attribute is present, loading from or
 storing to any pointer derived from the ``getelementptr`` has undefined
 behavior if the load or store would access memory outside the half-open range

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``.

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,

llvm/include/llvm/Bitcode/LLVMBitCodes.h

@@ -385,7 +385,7 @@ enum ConstantsCodes {
   CST_CODE_CSTRING = 9,          // CSTRING:       [values]
   CST_CODE_CE_BINOP = 10,        // CE_BINOP:      [opcode, opval, opval]
   CST_CODE_CE_CAST = 11,         // CE_CAST:       [opcode, opty, opval]
-  CST_CODE_CE_GEP = 12,          // CE_GEP:        [n x operands]
+  CST_CODE_CE_GEP_OLD = 12,      // CE_GEP:        [n x operands]
   CST_CODE_CE_SELECT = 13,       // CE_SELECT:     [opval, opval, opval]
   CST_CODE_CE_EXTRACTELT = 14,   // CE_EXTRACTELT: [opty, opval, opval]
   CST_CODE_CE_INSERTELT = 15,    // CE_INSERTELT:  [opval, opval, opval]
@@ -412,6 +412,7 @@ enum ConstantsCodes {
                                       //                 asmdialect|unwind,
                                       //                 asmstr,conststr]
   CST_CODE_CE_GEP_WITH_INRANGE = 31,  // [opty, flags, range, n x operands]
+  CST_CODE_CE_GEP = 32,               // [opty, flags, n x operands]
 };
 
 /// CastOpcodes - These are values used in the bitcode files to encode which
@@ -524,6 +525,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,

llvm/include/llvm/IR/Constants.h

@@ -28,6 +28,7 @@
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/GEPNoWrapFlags.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/OperandTraits.h"
 #include "llvm/IR/User.h"
@@ -1198,45 +1199,46 @@ class ConstantExpr : public Constant {
   /// \param OnlyIfReducedTy see \a getWithOperands() docs.
   static Constant *
   getGetElementPtr(Type *Ty, Constant *C, ArrayRef<Constant *> IdxList,
-                   bool InBounds = false,
+                   GEPNoWrapFlags NW = GEPNoWrapFlags::none(),
                    std::optional<ConstantRange> InRange = std::nullopt,
                    Type *OnlyIfReducedTy = nullptr) {
     return getGetElementPtr(
-        Ty, C, ArrayRef((Value *const *)IdxList.data(), IdxList.size()),
-        InBounds, InRange, OnlyIfReducedTy);
+        Ty, C, ArrayRef((Value *const *)IdxList.data(), IdxList.size()), NW,
+        InRange, OnlyIfReducedTy);
   }
   static Constant *
-  getGetElementPtr(Type *Ty, Constant *C, Constant *Idx, bool InBounds = false,
+  getGetElementPtr(Type *Ty, Constant *C, Constant *Idx,
+                   GEPNoWrapFlags NW = GEPNoWrapFlags::none(),
                    std::optional<ConstantRange> InRange = std::nullopt,
                    Type *OnlyIfReducedTy = nullptr) {
     // This form of the function only exists to avoid ambiguous overload
     // warnings about whether to convert Idx to ArrayRef<Constant *> or
     // ArrayRef<Value *>.
-    return getGetElementPtr(Ty, C, cast<Value>(Idx), InBounds, InRange,
+    return getGetElementPtr(Ty, C, cast<Value>(Idx), NW, InRange,
                             OnlyIfReducedTy);
   }
   static Constant *
   getGetElementPtr(Type *Ty, Constant *C, ArrayRef<Value *> IdxList,
-                   bool InBounds = false,
+                   GEPNoWrapFlags NW = GEPNoWrapFlags::none(),
                    std::optional<ConstantRange> InRange = std::nullopt,
                    Type *OnlyIfReducedTy = nullptr);
 
   /// Create an "inbounds" getelementptr. See the documentation for the
   /// "inbounds" flag in LangRef.html for details.
   static Constant *getInBoundsGetElementPtr(Type *Ty, Constant *C,
                                             ArrayRef<Constant *> IdxList) {
-    return getGetElementPtr(Ty, C, IdxList, true);
+    return getGetElementPtr(Ty, C, IdxList, GEPNoWrapFlags::inBounds());
   }
   static Constant *getInBoundsGetElementPtr(Type *Ty, Constant *C,
                                             Constant *Idx) {
     // This form of the function only exists to avoid ambiguous overload
     // warnings about whether to convert Idx to ArrayRef<Constant *> or
     // ArrayRef<Value *>.
-    return getGetElementPtr(Ty, C, Idx, true);
+    return getGetElementPtr(Ty, C, Idx, GEPNoWrapFlags::inBounds());
   }
   static Constant *getInBoundsGetElementPtr(Type *Ty, Constant *C,
                                             ArrayRef<Value *> IdxList) {
-    return getGetElementPtr(Ty, C, IdxList, true);
+    return getGetElementPtr(Ty, C, IdxList, GEPNoWrapFlags::inBounds());
   }
 
   static Constant *getExtractElement(Constant *Vec, Constant *Idx,

llvm/include/llvm/IR/GEPNoWrapFlags.h

@@ -0,0 +1,93 @@
+//===-- llvm/GEPNoWrapFlags.h - NoWrap flags for GEPs -----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the nowrap flags for getelementptr operators.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_IR_GEPNOWRAPFLAGS_H
+#define LLVM_IR_GEPNOWRAPFLAGS_H
+
+namespace llvm {
+
+/// Represents flags for the getelementptr instruction/expression.
+/// The following flags are supported:
+///  * inbounds (implies nusw)
+///  * nusw (no unsigned signed wrap)
+///  * nuw (no unsigned wrap)
+/// See LangRef for a description of their semantics.
+class GEPNoWrapFlags {
+  enum : unsigned {
+    InBoundsFlag = (1 << 0),
+    NUSWFlag = (1 << 1),
+    NUWFlag = (1 << 2),
+  };
+
+  unsigned Flags;
+  GEPNoWrapFlags(unsigned Flags) : Flags(Flags) {
+    assert((!isInBounds() || hasNoUnsignedSignedWrap()) &&
+           "inbounds implies nusw");
+  }
+
+public:
+  GEPNoWrapFlags() : Flags(0) {}
+  // For historical reasons, interpret plain boolean as InBounds.
+  // TODO: Migrate users to pass explicit GEPNoWrapFlags and remove this ctor.
+  GEPNoWrapFlags(bool IsInBounds)
+      : Flags(IsInBounds ? (InBoundsFlag | NUSWFlag) : 0) {}
+
+  static GEPNoWrapFlags none() { return GEPNoWrapFlags(); }
+  static GEPNoWrapFlags inBounds() {
+    return GEPNoWrapFlags(InBoundsFlag | NUSWFlag);
+  }
+  static GEPNoWrapFlags noUnsignedSignedWrap() {
+    return GEPNoWrapFlags(NUSWFlag);
+  }
+  static GEPNoWrapFlags noUnsignedWrap() { return GEPNoWrapFlags(NUWFlag); }
+
+  static GEPNoWrapFlags fromRaw(unsigned Flags) {
+    return GEPNoWrapFlags(Flags);
+  }
+  unsigned getRaw() const { return Flags; }
+
+  bool isInBounds() const { return Flags & InBoundsFlag; }
+  bool hasNoUnsignedSignedWrap() const { return Flags & NUSWFlag; }
+  bool hasNoUnsignedWrap() const { return Flags & NUWFlag; }
+
+  GEPNoWrapFlags withoutInBounds() const {
+    return GEPNoWrapFlags(Flags & ~InBoundsFlag);
+  }
+  GEPNoWrapFlags withoutNoUnsignedSignedWrap() const {
+    return GEPNoWrapFlags(Flags & ~(InBoundsFlag | NUSWFlag));
+  }
+  GEPNoWrapFlags withoutNoUnsignedWrap() const {
+    return GEPNoWrapFlags(Flags & ~NUWFlag);
+  }
+
+  bool operator==(GEPNoWrapFlags Other) const { return Flags == Other.Flags; }
+  bool operator!=(GEPNoWrapFlags Other) const { return !(*this == Other); }
+
+  GEPNoWrapFlags operator&(GEPNoWrapFlags Other) const {
+    return GEPNoWrapFlags(Flags & Other.Flags);
+  }
+  GEPNoWrapFlags operator|(GEPNoWrapFlags Other) const {
+    return GEPNoWrapFlags(Flags | Other.Flags);
+  }
+  GEPNoWrapFlags &operator&=(GEPNoWrapFlags Other) {
+    Flags &= Other.Flags;
+    return *this;
+  }
+  GEPNoWrapFlags &operator|=(GEPNoWrapFlags Other) {
+    Flags |= Other.Flags;
+    return *this;
+  }
+};
+
+} // end namespace llvm
+
+#endif // LLVM_IR_GEPNOWRAPFLAGS_H

llvm/include/llvm/IR/Instructions.h

@@ -26,6 +26,7 @@
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/GEPNoWrapFlags.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/OperandTraits.h"
@@ -1167,13 +1168,26 @@ class GetElementPtrInst : public Instruction {
   /// a constant offset between them.
   bool hasAllConstantIndices() const;
 
+  /// Set nowrap flags for GEP instruction.
+  void setNoWrapFlags(GEPNoWrapFlags NW);
+
   /// Set or clear the inbounds flag on this GEP instruction.
   /// See LangRef.html for the meaning of inbounds on a getelementptr.
+  /// TODO: Remove this method in favor of setNoWrapFlags().
   void setIsInBounds(bool b = true);
 
+  /// Get the nowrap flags for the GEP instruction.
+  GEPNoWrapFlags getNoWrapFlags() const;
+
   /// 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

llvm/include/llvm/IR/Operator.h

@@ -17,6 +17,7 @@
 #include "llvm/ADT/MapVector.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/FMF.h"
+#include "llvm/IR/GEPNoWrapFlags.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
@@ -399,26 +400,24 @@ class LShrOperator
 };
 
 class GEPOperator
-  : public ConcreteOperator<Operator, Instruction::GetElementPtr> {
-  friend class GetElementPtrInst;
-  friend class ConstantExpr;
-
-  enum {
-    IsInBounds = (1 << 0),
-  };
-
-  void setIsInBounds(bool B) {
-    SubclassOptionalData =
-      (SubclassOptionalData & ~IsInBounds) | (B * IsInBounds);
-  }
-
+    : public ConcreteOperator<Operator, Instruction::GetElementPtr> {
 public:
   /// Transparently provide more efficient getOperand methods.
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(Value);
 
+  GEPNoWrapFlags getNoWrapFlags() const {
+    return GEPNoWrapFlags::fromRaw(SubclassOptionalData);
+  }
+
   /// Test whether this is an inbounds GEP, as defined by LangRef.html.
-  bool isInBounds() const {
-    return SubclassOptionalData & IsInBounds;
+  bool isInBounds() const { return getNoWrapFlags().isInBounds(); }
+
+  bool hasNoUnsignedSignedWrap() const {
+    return getNoWrapFlags().hasNoUnsignedSignedWrap();
+  }
+
+  bool hasNoUnsignedWrap() const {
+    return getNoWrapFlags().hasNoUnsignedWrap();
   }
 
   /// Returns the offset of the index with an inrange attachment, or

llvm/lib/Analysis/ConstantFolding.cpp

@@ -1005,7 +1005,8 @@ Constant *ConstantFoldInstOperandsImpl(const Value *InstOrCE, unsigned Opcode,
       return C;
 
     return ConstantExpr::getGetElementPtr(SrcElemTy, Ops[0], Ops.slice(1),
-                                          GEP->isInBounds(), GEP->getInRange());
+                                          GEP->getNoWrapFlags(),
+                                          GEP->getInRange());
   }
 
   if (auto *CE = dyn_cast<ConstantExpr>(InstOrCE)) {

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);