[Inliner] Propagate range attributes to params when inlining

goldsteinn · goldsteinn · commit c63b635f0390 · 2024-05-11T00:30:52.000-05:00
diff --git a/llvm/include/llvm/IR/Attributes.h b/llvm/include/llvm/IR/Attributes.h
@@ -752,6 +752,11 @@ class AttributeList {
   [[nodiscard]] AttributeList addRangeRetAttr(LLVMContext &C,
                                               const ConstantRange &CR) const;
 
+  /// Add the range attribute to the attribute set at the given arg index.
+  /// Returns a new list because attribute lists are immutable.
+  [[nodiscard]] AttributeList addRangeParamAttr(LLVMContext &C, unsigned Index,
+                                                const ConstantRange &CR) const;
+
   /// Add the allocsize attribute to the attribute set at the given arg index.
   /// Returns a new list because attribute lists are immutable.
   [[nodiscard]] AttributeList
@@ -906,6 +911,9 @@ class AttributeList {
   /// arg.
   uint64_t getParamDereferenceableOrNullBytes(unsigned ArgNo) const;
 
+  /// Get range (or std::nullopt if unknown) of an arg.
+  std::optional<ConstantRange> getParamRange(unsigned ArgNo) const;
+
   /// Get the disallowed floating-point classes of the return value.
   FPClassTest getRetNoFPClass() const;
 
@@ -1082,6 +1090,10 @@ class AttrBuilder {
   /// invalid if the Kind is not present in the builder.
   Attribute getAttribute(StringRef Kind) const;
 
+  /// Retrieve the range if the attribute exists (std::nullopt is returned
+  /// otherwise).
+  std::optional<ConstantRange> getRange() const;
+
   /// Return raw (possibly packed/encoded) value of integer attribute or
   /// std::nullopt if not set.
   std::optional<uint64_t> getRawIntAttr(Attribute::AttrKind Kind) const;
diff --git a/llvm/include/llvm/IR/InstrTypes.h b/llvm/include/llvm/IR/InstrTypes.h
@@ -2204,6 +2204,10 @@ class CallBase : public Instruction {
   /// parameter.
   FPClassTest getParamNoFPClass(unsigned i) const;
 
+  /// If arg ArgNo has a range attribute, return the value range of the
+  /// argument. Otherwise, std::nullopt is returned.
+  std::optional<ConstantRange> getParamRange(unsigned ArgNo) const;
+
   /// If this return value has a range attribute, return the value range of the
   /// argument. Otherwise, std::nullopt is returned.
   std::optional<ConstantRange> getRange() const;
diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
@@ -1530,6 +1530,13 @@ AttributeList::addDereferenceableOrNullParamAttr(LLVMContext &C, unsigned Index,
   return addParamAttributes(C, Index, B);
 }
 
+AttributeList AttributeList::addRangeParamAttr(LLVMContext &C, unsigned Index,
+                                               const ConstantRange &CR) const {
+  AttrBuilder B(C);
+  B.addRangeAttr(CR);
+  return addParamAttributes(C, Index, B);
+}
+
 AttributeList AttributeList::addRangeRetAttr(LLVMContext &C,
                                              const ConstantRange &CR) const {
   AttrBuilder B(C);
@@ -1658,6 +1665,14 @@ AttributeList::getParamDereferenceableOrNullBytes(unsigned Index) const {
   return getParamAttrs(Index).getDereferenceableOrNullBytes();
 }
 
+std::optional<ConstantRange>
+AttributeList::getParamRange(unsigned Index) const {
+  auto RangeAttr = getParamAttrs(Index).getAttribute(Attribute::Range);
+  if (RangeAttr.isValid())
+    return RangeAttr.getRange();
+  return std::nullopt;
+}
+
 FPClassTest AttributeList::getRetNoFPClass() const {
   return getRetAttrs().getNoFPClass();
 }
@@ -1991,6 +2006,13 @@ Attribute AttrBuilder::getAttribute(StringRef A) const {
   return {};
 }
 
+std::optional<ConstantRange> AttrBuilder::getRange() const {
+  const Attribute RangeAttr = getAttribute(Attribute::Range);
+  if (RangeAttr.isValid())
+    return RangeAttr.getRange();
+  return std::nullopt;
+}
+
 bool AttrBuilder::contains(Attribute::AttrKind A) const {
   return getAttribute(A).isValid();
 }
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
@@ -396,6 +396,13 @@ FPClassTest CallBase::getParamNoFPClass(unsigned i) const {
   return Mask;
 }
 
+std::optional<ConstantRange> CallBase::getParamRange(unsigned ArgNo) const {
+  const Attribute RangeAttr = getParamAttr(ArgNo, llvm::Attribute::Range);
+  if (RangeAttr.isValid())
+    return RangeAttr.getRange();
+  return std::nullopt;
+}
+
 std::optional<ConstantRange> CallBase::getRange() const {
   const Attribute RangeAttr = getRetAttr(llvm::Attribute::Range);
   if (RangeAttr.isValid())
diff --git a/llvm/lib/Transforms/Utils/InlineFunction.cpp b/llvm/lib/Transforms/Utils/InlineFunction.cpp
@@ -1384,6 +1384,8 @@ static void AddParamAndFnBasicAttributes(const CallBase &CB,
       ValidExactParamAttrs.back().addAttribute(Attribute::NonNull);
     if (auto Align = CB.getParamAlign(I))
       ValidExactParamAttrs.back().addAlignmentAttr(Align);
+    if (auto Range = CB.getParamRange(I))
+      ValidExactParamAttrs.back().addRangeAttr(*Range);
 
     HasAttrToPropagate |= ValidObjParamAttrs.back().hasAttributes();
     HasAttrToPropagate |= ValidExactParamAttrs.back().hasAttributes();
@@ -1427,8 +1429,20 @@ static void AddParamAndFnBasicAttributes(const CallBase &CB,
               ValidExactParamAttrs[ArgNo].getAlignment().valueOrOne())
             AL = AL.removeParamAttribute(Context, I, Attribute::Alignment);
 
+          auto ExistingRange = AL.getParamRange(I);
           AL = AL.addParamAttributes(Context, I, ValidExactParamAttrs[ArgNo]);
 
+          // For range we use the exact intersection.
+          if (ExistingRange.has_value()) {
+            if (auto NewRange = ValidExactParamAttrs[ArgNo].getRange()) {
+              auto CombinedRange = ExistingRange->exactIntersectWith(*NewRange);
+              if (!CombinedRange.has_value())
+                CombinedRange =
+                    ConstantRange::getEmpty(NewRange->getBitWidth());
+              AL = AL.removeParamAttribute(Context, I, Attribute::Range);
+              AL = AL.addRangeParamAttr(Context, I, *CombinedRange);
+            }
+          }
         } else {
           // Check if the underlying value for the parameter is an argument.
           const Value *UnderlyingV =
diff --git a/llvm/test/Transforms/Inline/access-attributes-prop.ll b/llvm/test/Transforms/Inline/access-attributes-prop.ll
@@ -585,7 +585,7 @@ define dso_local void @foo4(i32 %v) {
 define void @prop_range_empty_intersect(i32 %v) {
 ; CHECK-LABEL: define {{[^@]+}}@prop_range_empty_intersect
 ; CHECK-SAME: (i32 [[V:%.*]]) {
-; CHECK-NEXT:    call void @bar4(i32 range(i32 0, 10) [[V]])
+; CHECK-NEXT:    call void @bar4(i32 range(i32 0, 0) [[V]])
 ; CHECK-NEXT:    ret void
 ;
   call void @foo4_range_0_10(i32 range(i32 11, 50) %v)
@@ -595,7 +595,7 @@ define void @prop_range_empty_intersect(i32 %v) {
 define void @prop_range_empty(i32 %v) {
 ; CHECK-LABEL: define {{[^@]+}}@prop_range_empty
 ; CHECK-SAME: (i32 [[V:%.*]]) {
-; CHECK-NEXT:    call void @bar4(i32 [[V]])
+; CHECK-NEXT:    call void @bar4(i32 range(i32 1, 0) [[V]])
 ; CHECK-NEXT:    ret void
 ;
   call void @foo4(i32 range(i32 1, 0) %v)
@@ -605,7 +605,7 @@ define void @prop_range_empty(i32 %v) {
 define void @prop_range_empty_with_intersect(i32 %v) {
 ; CHECK-LABEL: define {{[^@]+}}@prop_range_empty_with_intersect
 ; CHECK-SAME: (i32 [[V:%.*]]) {
-; CHECK-NEXT:    call void @bar4(i32 range(i32 0, 10) [[V]])
+; CHECK-NEXT:    call void @bar4(i32 range(i32 1, 10) [[V]])
 ; CHECK-NEXT:    ret void
 ;
   call void @foo4_range_0_10(i32 range(i32 1, 0) %v)
@@ -615,7 +615,7 @@ define void @prop_range_empty_with_intersect(i32 %v) {
 define void @prop_range_intersect1(i32 %v) {
 ; CHECK-LABEL: define {{[^@]+}}@prop_range_intersect1
 ; CHECK-SAME: (i32 [[V:%.*]]) {
-; CHECK-NEXT:    call void @bar4(i32 range(i32 0, 10) [[V]])
+; CHECK-NEXT:    call void @bar4(i32 range(i32 0, 9) [[V]])
 ; CHECK-NEXT:    ret void
 ;
   call void @foo4_range_0_10(i32 range(i32 0, 9) %v)
@@ -625,7 +625,7 @@ define void @prop_range_intersect1(i32 %v) {
 define void @prop_range_intersect2(i32 %v) {
 ; CHECK-LABEL: define {{[^@]+}}@prop_range_intersect2
 ; CHECK-SAME: (i32 [[V:%.*]]) {
-; CHECK-NEXT:    call void @bar4(i32 range(i32 0, 10) [[V]])
+; CHECK-NEXT:    call void @bar4(i32 range(i32 1, 9) [[V]])
 ; CHECK-NEXT:    ret void
 ;
   call void @foo4_range_0_10(i32 range(i32 1, 9) %v)
@@ -635,7 +635,7 @@ define void @prop_range_intersect2(i32 %v) {
 define void @prop_range_intersect3(i32 %v) {
 ; CHECK-LABEL: define {{[^@]+}}@prop_range_intersect3
 ; CHECK-SAME: (i32 [[V:%.*]]) {
-; CHECK-NEXT:    call void @bar4(i32 [[V]])
+; CHECK-NEXT:    call void @bar4(i32 range(i32 0, 11) [[V]])
 ; CHECK-NEXT:    ret void
 ;
   call void @foo4_2_range_0_10(i32 range(i32 0, 11) %v)
@@ -645,7 +645,7 @@ define void @prop_range_intersect3(i32 %v) {
 define void @prop_range_direct(i32 %v) {
 ; CHECK-LABEL: define {{[^@]+}}@prop_range_direct
 ; CHECK-SAME: (i32 [[V:%.*]]) {
-; CHECK-NEXT:    call void @bar4(i32 [[V]])
+; CHECK-NEXT:    call void @bar4(i32 range(i32 1, 11) [[V]])
 ; CHECK-NEXT:    ret void
 ;
   call void @foo4(i32 range(i32 1, 11) %v)
