[Intrinsics] Add support for range attributes (#135642)

nikic · web-flow · commit c3c0b27f2d34 · 2025-04-17T11:11:00.000+02:00
Add support for specifying range attributes in Intrinsics.td. Use this to specify the ucmp/scmp range [-1,2). This case is trickier than existing intrinsic attributes, because we need to create the attribute with the correct bitwidth. As such, the attribute construction now needs to be aware of the function type. We also need to be careful to no longer assign attributes on intrinsics with invalid signatures, as we'd make invalid assumptions about the number of arguments etc otherwise. Fixes #130179.
diff --git a/llvm/include/llvm/IR/Intrinsics.h b/llvm/include/llvm/IR/Intrinsics.h
@@ -86,7 +86,7 @@ namespace Intrinsic {
   ID lookupIntrinsicID(StringRef Name);
 
   /// Return the attributes for an intrinsic.
-  AttributeList getAttributes(LLVMContext &C, ID id);
+  AttributeList getAttributes(LLVMContext &C, ID id, FunctionType *FT);
 
   /// Return the function attributes for an intrinsic.
   AttributeSet getFnAttributes(LLVMContext &C, ID id);
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
@@ -134,6 +134,14 @@ class ReadNone<ArgIndex idx> : IntrinsicProperty {
   int ArgNo = idx.Value;
 }
 
+// The return value or argument is in the range [lower, upper),
+// where lower and upper are interpreted as signed integers.
+class Range<AttrIndex idx, int lower, int upper> : IntrinsicProperty {
+  int ArgNo = idx.Value;
+  int Lower = lower;
+  int Upper = upper;
+}
+
 def IntrNoReturn : IntrinsicProperty;
 
 // Applied by default.
@@ -1620,10 +1628,10 @@ def int_umin : DefaultAttrsIntrinsic<
     [IntrNoMem, IntrSpeculatable, IntrWillReturn]>;
 def int_scmp : DefaultAttrsIntrinsic<
     [llvm_anyint_ty], [llvm_anyint_ty, LLVMMatchType<1>],
-    [IntrNoMem, IntrSpeculatable, IntrWillReturn]>;
+    [IntrNoMem, IntrSpeculatable, IntrWillReturn, Range<RetIndex, -1, 2>]>;
 def int_ucmp : DefaultAttrsIntrinsic<
     [llvm_anyint_ty], [llvm_anyint_ty, LLVMMatchType<1>],
-    [IntrNoMem, IntrSpeculatable, IntrWillReturn]>;
+    [IntrNoMem, IntrSpeculatable, IntrWillReturn, Range<RetIndex, -1, 2>]>;
 
 //===------------------------- Memory Use Markers -------------------------===//
 //
diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp
@@ -1519,8 +1519,13 @@ bool llvm::UpgradeIntrinsicFunction(Function *F, Function *&NewFn,
   // Upgrade intrinsic attributes.  This does not change the function.
   if (NewFn)
     F = NewFn;
-  if (Intrinsic::ID id = F->getIntrinsicID())
-    F->setAttributes(Intrinsic::getAttributes(F->getContext(), id));
+  if (Intrinsic::ID id = F->getIntrinsicID()) {
+    // Only do this if the intrinsic signature is valid.
+    SmallVector<Type *> OverloadTys;
+    if (Intrinsic::getIntrinsicSignature(id, F->getFunctionType(), OverloadTys))
+      F->setAttributes(
+          Intrinsic::getAttributes(F->getContext(), id, F->getFunctionType()));
+  }
   return Upgraded;
 }
 
diff --git a/llvm/lib/IR/Function.cpp b/llvm/lib/IR/Function.cpp
@@ -514,8 +514,15 @@ Function::Function(FunctionType *Ty, LinkageTypes Linkage, unsigned AddrSpace,
   // Ensure intrinsics have the right parameter attributes.
   // Note, the IntID field will have been set in Value::setName if this function
   // name is a valid intrinsic ID.
-  if (IntID)
-    setAttributes(Intrinsic::getAttributes(getContext(), IntID));
+  if (IntID) {
+    // Don't set the attributes if the intrinsic signature is invalid. This
+    // case will either be auto-upgraded or fail verification.
+    SmallVector<Type *> OverloadTys;
+    if (!Intrinsic::getIntrinsicSignature(IntID, Ty, OverloadTys))
+      return;
+
+    setAttributes(Intrinsic::getAttributes(getContext(), IntID, Ty));
+  }
 }
 
 Function::~Function() {
diff --git a/llvm/lib/IR/Intrinsics.cpp b/llvm/lib/IR/Intrinsics.cpp
@@ -13,6 +13,7 @@
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringTable.h"
+#include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IntrinsicsAArch64.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -16506,7 +16506,8 @@ Align SITargetLowering::computeKnownAlignForTargetInstr(
     // site specifies a lower alignment?
     Intrinsic::ID IID = GI->getIntrinsicID();
     LLVMContext &Ctx = VT.getMachineFunction().getFunction().getContext();
-    AttributeList Attrs = Intrinsic::getAttributes(Ctx, IID);
+    AttributeList Attrs =
+        Intrinsic::getAttributes(Ctx, IID, Intrinsic::getType(Ctx, IID));
     if (MaybeAlign RetAlign = Attrs.getRetAlignment())
       return *RetAlign;
   }
diff --git a/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp b/llvm/lib/Transforms/Scalar/RewriteStatepointsForGC.cpp
@@ -2903,7 +2903,7 @@ static void stripNonValidAttributesFromPrototype(Function &F) {
   // assumes that the attributes defined in Intrinsic.td are conservatively
   // correct for both physical and abstract model.
   if (Intrinsic::ID id = F.getIntrinsicID()) {
-    F.setAttributes(Intrinsic::getAttributes(Ctx, id));
+    F.setAttributes(Intrinsic::getAttributes(Ctx, id, F.getFunctionType()));
     return;
   }
 
diff --git a/llvm/test/Assembler/aarch64-intrinsics-attributes.ll b/llvm/test/Assembler/aarch64-intrinsics-attributes.ll
@@ -10,8 +10,8 @@ declare i64 @llvm.aarch64.ldxr.p0(ptr)
 ; CHECK: declare i32 @llvm.aarch64.stxp(i64, i64, ptr) [[NOFREE_NOUNWIND_WILLRETURN]]
 declare i32 @llvm.aarch64.stxp(i64, i64, ptr)
 
-; CHECK: declare i32 @llvm.aarch64.dsb(i32) [[NOFREE_NOUNWIND_WILLRETURN]]
-declare i32 @llvm.aarch64.dsb(i32)
+; CHECK: declare void @llvm.aarch64.dsb(i32) [[NOFREE_NOUNWIND_WILLRETURN]]
+declare void @llvm.aarch64.dsb(i32)
 
 ; CHECK: declare i64 @llvm.aarch64.neon.sqdmulls.scalar(i32, i32) [[NO_CALLBACK_NOFREE_NOSYNC_NOUNWIND_READNONE_WILLRETURN:#[0-9]+]]
 declare i64 @llvm.aarch64.neon.sqdmulls.scalar(i32, i32)
diff --git a/llvm/test/Assembler/autoupgrade-invalid-mem-intrinsics.ll b/llvm/test/Assembler/autoupgrade-invalid-mem-intrinsics.ll
@@ -1,15 +1,13 @@
-; RUN: not llvm-as < %s -o /dev/null 2>&1 | FileCheck %s
+; RUN: llvm-as < %s | llvm-dis | FileCheck %s
 
 ; Check that remangling code doesn't fail on an intrinsic with wrong signature
+; TODO: This should probably produce an error.
 
-; CHECK: Attribute after last parameter!
-; CHECK-NEXT: ptr @llvm.memset.i64
+; CHECK: declare void @llvm.memset.i64
 declare void @llvm.memset.i64(ptr nocapture, i8, i64) nounwind
 
-; CHECK: Attribute after last parameter!
-; CHECK-NEXT: ptr @llvm.memcpy.i64
+; CHECK: declare void @llvm.memcpy.i64
 declare void @llvm.memcpy.i64(ptr nocapture, i8, i64) nounwind
 
-; CHECK: Attribute after last parameter!
-; CHECK-NEXT: ptr @llvm.memmove.i64
+; CHECK: declare void @llvm.memmove.i64
 declare void @llvm.memmove.i64(ptr nocapture, i8, i64) nounwind
diff --git a/llvm/test/TableGen/intrinsic-attrs.td b/llvm/test/TableGen/intrinsic-attrs.td
@@ -6,8 +6,10 @@ def int_random_gen   : Intrinsic<[llvm_i32_ty], [], [IntrNoMem, IntrHasSideEffec
 
 def int_deref_ptr_ret : Intrinsic<[llvm_ptr_ty], [], [Dereferenceable<RetIndex, 16>]>;
 
-// CHECK: static AttributeSet getIntrinsicArgAttributeSet(LLVMContext &C, unsigned ID) {
-// CHECK-NEXT:   switch (ID) {
+// CHECK: static AttributeSet getIntrinsicArgAttributeSet(LLVMContext &C, unsigned ID,
+// CHECK-NEXT: Type *ArgType) {
+// CHECK-NEXT: unsigned BitWidth = ArgType->getScalarSizeInBits();
+// CHECK-NEXT: switch (ID) {
 // CHECK-NEXT: default: llvm_unreachable("Invalid attribute set number");
 // CHECK-NEXT: case 0:
 // CHECK-NEXT:     return AttributeSet::get(C, {
@@ -26,13 +28,14 @@ def int_deref_ptr_ret : Intrinsic<[llvm_ptr_ty], [], [Dereferenceable<RetIndex,
 // CHECK: 0 << 8 | 0, // llvm.deref.ptr.ret
 // CHECK: 1 << 8 | 1, // llvm.random.gen
 
-// CHECK: getAttributes(LLVMContext &C, ID id)
+// CHECK: getAttributes(LLVMContext &C, ID id,
+// CHECK-NEXT: FunctionType *FT) {
 // CHECK: case 1:
 // CHECK-NEXT: return AttributeList::get(C, {
 // CHECK-NEXT:   {AttributeList::FunctionIndex, getIntrinsicFnAttributeSet(C, FnAttrID)}
 // CHECK-NEXT: });
 // CHECK-NEXT: case 0:
 // CHECK-NEXT: return AttributeList::get(C, {
-// CHECK-NEXT:   {0, getIntrinsicArgAttributeSet(C, 0)},
+// CHECK-NEXT:   {0, getIntrinsicArgAttributeSet(C, 0, FT->getContainedType(0))},
 // CHECK-NEXT:   {AttributeList::FunctionIndex, getIntrinsicFnAttributeSet(C, FnAttrID)}
 // CHECK-NEXT: });
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/uscmp.ll b/llvm/test/Transforms/CorrelatedValuePropagation/uscmp.ll
@@ -261,11 +261,9 @@ define i8 @ucmp_switch(i32 %x, i32 %y) {
 ; CHECK-LABEL: @ucmp_switch(
 ; CHECK-NEXT:    [[CMP:%.*]] = call i8 @llvm.ucmp.i8.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
 ; CHECK-NEXT:    switch i8 [[CMP]], label [[DEFAULT_UNREACHABLE:%.*]] [
-; CHECK-NEXT:      i8 -2, label [[BB_NEG2:%.*]]
+; CHECK-NEXT:      i8 1, label [[BB_1:%.*]]
 ; CHECK-NEXT:      i8 -1, label [[BB_NEG1:%.*]]
 ; CHECK-NEXT:      i8 0, label [[BB_0:%.*]]
-; CHECK-NEXT:      i8 1, label [[BB_1:%.*]]
-; CHECK-NEXT:      i8 2, label [[BB_2:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       bb.neg2:
 ; CHECK-NEXT:    ret i8 -2
@@ -277,6 +275,8 @@ define i8 @ucmp_switch(i32 %x, i32 %y) {
 ; CHECK-NEXT:    ret i8 1
 ; CHECK:       bb.2:
 ; CHECK-NEXT:    ret i8 2
+; CHECK:       default.unreachable:
+; CHECK-NEXT:    unreachable
 ; CHECK:       default:
 ; CHECK-NEXT:    ret i8 123
 ;
@@ -312,11 +312,9 @@ define i8 @scmp_switch(i32 %x, i32 %y) {
 ; CHECK-LABEL: @scmp_switch(
 ; CHECK-NEXT:    [[CMP:%.*]] = call i8 @llvm.scmp.i8.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
 ; CHECK-NEXT:    switch i8 [[CMP]], label [[DEFAULT_UNREACHABLE:%.*]] [
-; CHECK-NEXT:      i8 -2, label [[BB_NEG2:%.*]]
+; CHECK-NEXT:      i8 1, label [[BB_1:%.*]]
 ; CHECK-NEXT:      i8 -1, label [[BB_NEG1:%.*]]
 ; CHECK-NEXT:      i8 0, label [[BB_0:%.*]]
-; CHECK-NEXT:      i8 1, label [[BB_1:%.*]]
-; CHECK-NEXT:      i8 2, label [[BB_2:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       bb.neg2:
 ; CHECK-NEXT:    ret i8 -2
@@ -328,6 +326,8 @@ define i8 @scmp_switch(i32 %x, i32 %y) {
 ; CHECK-NEXT:    ret i8 1
 ; CHECK:       bb.2:
 ; CHECK-NEXT:    ret i8 2
+; CHECK:       default.unreachable:
+; CHECK-NEXT:    unreachable
 ; CHECK:       default:
 ; CHECK-NEXT:    ret i8 123
 ;
diff --git a/llvm/utils/TableGen/Basic/CodeGenIntrinsics.cpp b/llvm/utils/TableGen/Basic/CodeGenIntrinsics.cpp
@@ -439,6 +439,11 @@ void CodeGenIntrinsic::setProperty(const Record *R) {
     unsigned ArgNo = R->getValueAsInt("ArgNo");
     uint64_t Bytes = R->getValueAsInt("Bytes");
     addArgAttribute(ArgNo, Dereferenceable, Bytes);
+  } else if (R->isSubClassOf("Range")) {
+    unsigned ArgNo = R->getValueAsInt("ArgNo");
+    int64_t Lower = R->getValueAsInt("Lower");
+    int64_t Upper = R->getValueAsInt("Upper");
+    addArgAttribute(ArgNo, Range, Lower, Upper);
   } else
     llvm_unreachable("Unknown property!");
 }
@@ -455,14 +460,14 @@ bool CodeGenIntrinsic::isParamImmArg(unsigned ParamIdx) const {
   ++ParamIdx;
   if (ParamIdx >= ArgumentAttributes.size())
     return false;
-  ArgAttribute Val{ImmArg, 0};
+  ArgAttribute Val{ImmArg, 0, 0};
   return std::binary_search(ArgumentAttributes[ParamIdx].begin(),
                             ArgumentAttributes[ParamIdx].end(), Val);
 }
 
-void CodeGenIntrinsic::addArgAttribute(unsigned Idx, ArgAttrKind AK,
-                                       uint64_t V) {
+void CodeGenIntrinsic::addArgAttribute(unsigned Idx, ArgAttrKind AK, uint64_t V,
+                                       uint64_t V2) {
   if (Idx >= ArgumentAttributes.size())
     ArgumentAttributes.resize(Idx + 1);
-  ArgumentAttributes[Idx].emplace_back(AK, V);
+  ArgumentAttributes[Idx].emplace_back(AK, V, V2);
 }
diff --git a/llvm/utils/TableGen/Basic/CodeGenIntrinsics.h b/llvm/utils/TableGen/Basic/CodeGenIntrinsics.h
@@ -128,24 +128,29 @@ struct CodeGenIntrinsic {
     ReadNone,
     ImmArg,
     Alignment,
-    Dereferenceable
+    Dereferenceable,
+    Range,
   };
 
   struct ArgAttribute {
     ArgAttrKind Kind;
     uint64_t Value;
+    uint64_t Value2;
 
-    ArgAttribute(ArgAttrKind K, uint64_t V) : Kind(K), Value(V) {}
+    ArgAttribute(ArgAttrKind K, uint64_t V, uint64_t V2)
+        : Kind(K), Value(V), Value2(V2) {}
 
     bool operator<(const ArgAttribute &Other) const {
-      return std::tie(Kind, Value) < std::tie(Other.Kind, Other.Value);
+      return std::tie(Kind, Value, Value2) <
+             std::tie(Other.Kind, Other.Value, Other.Value2);
     }
   };
 
   /// Vector of attributes for each argument.
   SmallVector<SmallVector<ArgAttribute, 0>> ArgumentAttributes;
 
-  void addArgAttribute(unsigned Idx, ArgAttrKind AK, uint64_t V = 0);
+  void addArgAttribute(unsigned Idx, ArgAttrKind AK, uint64_t V = 0,
+                       uint64_t V2 = 0);
 
   bool hasProperty(enum SDNP Prop) const { return Properties & (1 << Prop); }
 
diff --git a/llvm/utils/TableGen/Basic/IntrinsicEmitter.cpp b/llvm/utils/TableGen/Basic/IntrinsicEmitter.cpp
@@ -493,6 +493,8 @@ static StringRef getArgAttrEnumName(CodeGenIntrinsic::ArgAttrKind Kind) {
     return "Alignment";
   case CodeGenIntrinsic::Dereferenceable:
     return "Dereferenceable";
+  case CodeGenIntrinsic::Range:
+    return "Range";
   }
   llvm_unreachable("Unknown CodeGenIntrinsic::ArgAttrKind enum");
 }
@@ -502,7 +504,9 @@ void IntrinsicEmitter::EmitAttributes(const CodeGenIntrinsicTable &Ints,
                                       raw_ostream &OS) {
   OS << R"(// Add parameter attributes that are not common to all intrinsics.
 #ifdef GET_INTRINSIC_ATTRIBUTES
-static AttributeSet getIntrinsicArgAttributeSet(LLVMContext &C, unsigned ID) {
+static AttributeSet getIntrinsicArgAttributeSet(LLVMContext &C, unsigned ID,
+                                                Type *ArgType) {
+  unsigned BitWidth = ArgType->getScalarSizeInBits();
   switch (ID) {
   default: llvm_unreachable("Invalid attribute set number");)";
   // Compute unique argument attribute sets.
@@ -535,6 +539,17 @@ static AttributeSet getIntrinsicArgAttributeSet(LLVMContext &C, unsigned ID) {
             Attr.Kind == CodeGenIntrinsic::Dereferenceable)
           OS << formatv("      Attribute::get(C, Attribute::{}, {}),\n",
                         AttrName, Attr.Value);
+        else if (Attr.Kind == CodeGenIntrinsic::Range)
+          // This allows implicitTrunc because the range may only fit the
+          // type based on rules implemented in the IR verifier. E.g. the
+          // [-1, 1] range for ucmp/scmp intrinsics requires a minimum i2 type.
+          // Give the verifier a chance to diagnose this instead of asserting
+          // here.
+          OS << formatv("      Attribute::get(C, Attribute::{}, "
+                        "ConstantRange(APInt(BitWidth, {}, /*isSigned=*/true, "
+                        "/*implicitTrunc=*/true), APInt(BitWidth, {}, "
+                        "/*isSigned=*/true, /*implicitTrunc=*/true))),\n",
+                        AttrName, (int64_t)Attr.Value, (int64_t)Attr.Value2);
         else
           OS << formatv("      Attribute::get(C, Attribute::{}),\n", AttrName);
       }
@@ -635,7 +650,8 @@ static constexpr uint16_t IntrinsicsToAttributesMap[] = {)";
   OS << R"(
 };
 
-AttributeList Intrinsic::getAttributes(LLVMContext &C, ID id) {)";
+AttributeList Intrinsic::getAttributes(LLVMContext &C, ID id,
+                                       FunctionType *FT) {)";
 
   OS << formatv(R"(
   if (id == 0)
@@ -669,8 +685,9 @@ AttributeList Intrinsic::getAttributes(LLVMContext &C, ID id) {)";
 
       unsigned ArgAttrID = UniqArgAttributes.find(Attrs)->second;
       OS << LS
-         << formatv("      {{{}, getIntrinsicArgAttributeSet(C, {})}", AttrIdx,
-                    ArgAttrID);
+         << formatv("      {{{}, getIntrinsicArgAttributeSet(C, {}, "
+                    "FT->getContainedType({}))}",
+                    AttrIdx, ArgAttrID, AttrIdx);
     }
 
     if (hasFnAttributes(Int)) {

Original file line number	Diff line number	Diff line change
`@@ -2903,7 +2903,7 @@ static void stripNonValidAttributesFromPrototype(Function &F) {`
`2903`	`2903`	`// assumes that the attributes defined in Intrinsic.td are conservatively`
`2904`	`2904`	`// correct for both physical and abstract model.`
`2905`	`2905`	`if (Intrinsic::ID id = F.getIntrinsicID()) {`
`2906`		`- F.setAttributes(Intrinsic::getAttributes(Ctx, id));`
	`2906`	`+ F.setAttributes(Intrinsic::getAttributes(Ctx, id, F.getFunctionType()));`
`2907`	`2907`	`return;`
`2908`	`2908`	`}`
`2909`	`2909`