[IR] Drop poison-generating return attributes when necessary (#89138)

Rename has/dropPoisonGeneratingFlagsOrMetadata to
has/dropPoisonGeneratingAnnotations and make it also handle
nonnull, align and range return attributes on calls, similar
to the existing handling for !nonnull, !align and !range metadata.

Author: andjo403

llvm/include/llvm/IR/Instruction.h

@@ -496,14 +496,24 @@ class Instruction : public User,
   /// Drops metadata that may generate poison.
   void dropPoisonGeneratingMetadata();
 
-  /// Return true if this instruction has poison-generating flags or metadata.
-  bool hasPoisonGeneratingFlagsOrMetadata() const {
-    return hasPoisonGeneratingFlags() || hasPoisonGeneratingMetadata();
+  /// Return true if this instruction has poison-generating attribute.
+  bool hasPoisonGeneratingReturnAttributes() const LLVM_READONLY;
+
+  /// Drops return attributes that may generate poison.
+  void dropPoisonGeneratingReturnAttributes();
+
+  /// Return true if this instruction has poison-generating flags,
+  /// return attributes or metadata.
+  bool hasPoisonGeneratingAnnotations() const {
+    return hasPoisonGeneratingFlags() ||
+           hasPoisonGeneratingReturnAttributes() ||
+           hasPoisonGeneratingMetadata();
   }
 
-  /// Drops flags and metadata that may generate poison.
-  void dropPoisonGeneratingFlagsAndMetadata() {
+  /// Drops flags, return attributes and metadata that may generate poison.
+  void dropPoisonGeneratingAnnotations() {
     dropPoisonGeneratingFlags();
+    dropPoisonGeneratingReturnAttributes();
     dropPoisonGeneratingMetadata();
   }
 

llvm/include/llvm/IR/Operator.h

@@ -64,9 +64,10 @@ class Operator : public User {
   /// to evaluate to poison despite having non-poison inputs.
   bool hasPoisonGeneratingFlags() const;
 
-  /// Return true if this operator has poison-generating flags or metadata.
-  /// The latter is only possible for instructions.
-  bool hasPoisonGeneratingFlagsOrMetadata() const;
+  /// Return true if this operator has poison-generating flags,
+  /// return attributes or metadata. The latter two is only possible for
+  /// instructions.
+  bool hasPoisonGeneratingAnnotations() const;
 };
 
 /// Utility class for integer operators which may exhibit overflow - Add, Sub,

llvm/lib/Analysis/InstructionSimplify.cpp

@@ -4431,7 +4431,7 @@ static Value *simplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
         return nullptr;
     }
     Constant *Res = ConstantFoldInstOperands(I, ConstOps, Q.DL, Q.TLI);
-    if (DropFlags && Res && I->hasPoisonGeneratingFlagsOrMetadata())
+    if (DropFlags && Res && I->hasPoisonGeneratingAnnotations())
       DropFlags->push_back(I);
     return Res;
   }

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -4239,7 +4239,7 @@ bool ScalarEvolution::canReuseInstruction(
       return false;
 
     // If the instruction can't create poison, we can recurse to its operands.
-    if (I->hasPoisonGeneratingFlagsOrMetadata())
+    if (I->hasPoisonGeneratingAnnotations())
       DropPoisonGeneratingInsts.push_back(I);
 
     for (Value *Op : I->operands())

llvm/lib/Analysis/ValueTracking.cpp

@@ -6951,7 +6951,7 @@ static bool canCreateUndefOrPoison(const Operator *Op, UndefPoisonKind Kind,
                                    bool ConsiderFlagsAndMetadata) {
 
   if (ConsiderFlagsAndMetadata && includesPoison(Kind) &&
-      Op->hasPoisonGeneratingFlagsOrMetadata())
+      Op->hasPoisonGeneratingAnnotations())
     return true;
 
   unsigned Opcode = Op->getOpcode();

llvm/lib/IR/Instruction.cpp

@@ -13,7 +13,9 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/IR/AttributeMask.h"
+#include "llvm/IR/Attributes.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
@@ -472,6 +474,27 @@ void Instruction::dropPoisonGeneratingMetadata() {
   eraseMetadata(LLVMContext::MD_align);
 }
 
+bool Instruction::hasPoisonGeneratingReturnAttributes() const {
+  if (const auto *CB = dyn_cast<CallBase>(this)) {
+    AttributeSet RetAttrs = CB->getAttributes().getRetAttrs();
+    return RetAttrs.hasAttribute(Attribute::Range) ||
+           RetAttrs.hasAttribute(Attribute::Alignment) ||
+           RetAttrs.hasAttribute(Attribute::NonNull);
+  }
+  return false;
+}
+
+void Instruction::dropPoisonGeneratingReturnAttributes() {
+  if (auto *CB = dyn_cast<CallBase>(this)) {
+    AttributeMask AM;
+    AM.addAttribute(Attribute::Range);
+    AM.addAttribute(Attribute::Alignment);
+    AM.addAttribute(Attribute::NonNull);
+    CB->removeRetAttrs(AM);
+  }
+  assert(!hasPoisonGeneratingReturnAttributes() && "must be kept in sync");
+}
+
 void Instruction::dropUBImplyingAttrsAndUnknownMetadata(
     ArrayRef<unsigned> KnownIDs) {
   dropUnknownNonDebugMetadata(KnownIDs);

llvm/lib/IR/Operator.cpp

@@ -56,11 +56,12 @@ bool Operator::hasPoisonGeneratingFlags() const {
   }
 }
 
-bool Operator::hasPoisonGeneratingFlagsOrMetadata() const {
+bool Operator::hasPoisonGeneratingAnnotations() const {
   if (hasPoisonGeneratingFlags())
     return true;
   auto *I = dyn_cast<Instruction>(this);
-  return I && I->hasPoisonGeneratingMetadata();
+  return I && (I->hasPoisonGeneratingReturnAttributes() ||
+               I->hasPoisonGeneratingMetadata());
 }
 
 Type *GEPOperator::getSourceElementType() const {

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -1340,7 +1340,7 @@ Instruction *InstCombinerImpl::foldSelectValueEquivalence(SelectInst &Sel,
                              /* AllowRefinement */ false,
                              &DropFlags) == TrueVal) {
     for (Instruction *I : DropFlags) {
-      I->dropPoisonGeneratingFlagsAndMetadata();
+      I->dropPoisonGeneratingAnnotations();
       Worklist.add(I);
     }
 

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -4330,7 +4330,7 @@ InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
       return nullptr;
   }
 
-  OrigOpInst->dropPoisonGeneratingFlagsAndMetadata();
+  OrigOpInst->dropPoisonGeneratingAnnotations();
 
   // If all operands are guaranteed to be non-poison, we can drop freeze.
   if (!MaybePoisonOperand)
@@ -4401,7 +4401,7 @@ Instruction *InstCombinerImpl::foldFreezeIntoRecurrence(FreezeInst &FI,
   }
 
   for (Instruction *I : DropFlags)
-    I->dropPoisonGeneratingFlagsAndMetadata();
+    I->dropPoisonGeneratingAnnotations();
 
   if (StartNeedsFreeze) {
     Builder.SetInsertPoint(StartBB->getTerminator());

llvm/lib/Transforms/Scalar/BDCE.cpp

@@ -75,7 +75,7 @@ static void clearAssumptionsOfUsers(Instruction *I, DemandedBits &DB) {
     Instruction *J = WorkList.pop_back_val();
 
     // NSW, NUW, and exact are based on operands that might have changed.
-    J->dropPoisonGeneratingFlagsAndMetadata();
+    J->dropPoisonGeneratingAnnotations();
 
     // We do not have to worry about llvm.assume, because it demands its
     // operand, so trivializing can't change it.

llvm/lib/Transforms/Scalar/GuardWidening.cpp

@@ -699,7 +699,7 @@ Value *GuardWideningImpl::freezeAndPush(Value *Orig,
         Worklist.push_back(U.get());
   }
   for (Instruction *I : DropPoisonFlags)
-    I->dropPoisonGeneratingFlagsAndMetadata();
+    I->dropPoisonGeneratingAnnotations();
 
   Value *Result = Orig;
   for (Value *V : NeedFreeze) {

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

@@ -1522,7 +1522,7 @@ Value *SCEVExpander::expand(const SCEV *S) {
   } else {
     for (Instruction *I : DropPoisonGeneratingInsts) {
       rememberFlags(I);
-      I->dropPoisonGeneratingFlagsAndMetadata();
+      I->dropPoisonGeneratingAnnotations();
       // See if we can re-infer from first principles any of the flags we just
       // dropped.
       if (auto *OBO = dyn_cast<OverflowingBinaryOperator>(I))

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

@@ -768,7 +768,7 @@ bool SimplifyIndvar::eliminateIdentitySCEV(Instruction *UseInst,
       return false;
 
     for (Instruction *I : DropPoisonGeneratingInsts)
-      I->dropPoisonGeneratingFlagsAndMetadata();
+      I->dropPoisonGeneratingAnnotations();
   }
 
   LLVM_DEBUG(dbgs() << "INDVARS: Eliminated identity: " << *UseInst << '\n');

llvm/test/Transforms/InstCombine/freeze.ll

@@ -1069,7 +1069,7 @@ define ptr @freeze_load_dereferenceable(ptr %ptr) {
 
 define ptr @freeze_load_dereferenceable_or_null(ptr %ptr) {
 ; CHECK-LABEL: @freeze_load_dereferenceable_or_null(
-; CHECK-NEXT:    [[P:%.*]] = load ptr, ptr [[PTR:%.*]], align 8, !dereferenceable_or_null !1
+; CHECK-NEXT:    [[P:%.*]] = load ptr, ptr [[PTR:%.*]], align 8, !dereferenceable_or_null [[META1]]
 ; CHECK-NEXT:    ret ptr [[P]]
 ;
   %p = load ptr, ptr %ptr, !dereferenceable_or_null !1
@@ -1160,6 +1160,45 @@ define i32 @propagate_drop_flags_trunc(i64 %arg) {
   ret i32 %v1.fr
 }
 
+declare i32 @llvm.umax.i32(i32 %a, i32 %b)
+
+define i32 @freeze_call_with_range_attr(i32 %a) {
+; CHECK-LABEL: @freeze_call_with_range_attr(
+; CHECK-NEXT:    [[Y:%.*]] = lshr i32 2047, [[A:%.*]]
+; CHECK-NEXT:    [[Y_FR:%.*]] = freeze i32 [[Y]]
+; CHECK-NEXT:    [[X:%.*]] = call i32 @llvm.umax.i32(i32 [[Y_FR]], i32 50)
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  %y = lshr i32 2047, %a
+  %x = call range(i32 0, 2048) i32 @llvm.umax.i32(i32 %y, i32 50)
+  %x.fr = freeze i32 %x
+  ret i32 %x.fr
+}
+
+declare ptr @llvm.ptrmask.p0.i64(ptr, i64)
+
+define ptr @freeze_ptrmask_align(ptr %p, i64 noundef %m) {
+; CHECK-LABEL: @freeze_ptrmask_align(
+; CHECK-NEXT:    [[P_FR:%.*]] = freeze ptr [[P:%.*]]
+; CHECK-NEXT:    [[MASK:%.*]] = call ptr @llvm.ptrmask.p0.i64(ptr [[P_FR]], i64 [[M:%.*]])
+; CHECK-NEXT:    ret ptr [[MASK]]
+;
+  %mask = call align(4) ptr @llvm.ptrmask.p0.i64(ptr %p, i64 %m)
+  %fr =  freeze ptr %mask
+  ret ptr %fr
+}
+
+define ptr @freeze_ptrmask_nonnull(ptr %p, i64 noundef %m) {
+; CHECK-LABEL: @freeze_ptrmask_nonnull(
+; CHECK-NEXT:    [[P_FR:%.*]] = freeze ptr [[P:%.*]]
+; CHECK-NEXT:    [[MASK:%.*]] = call ptr @llvm.ptrmask.p0.i64(ptr [[P_FR]], i64 [[M:%.*]])
+; CHECK-NEXT:    ret ptr [[MASK]]
+;
+  %mask = call nonnull ptr @llvm.ptrmask.p0.i64(ptr %p, i64 %m)
+  %fr =  freeze ptr %mask
+  ret ptr %fr
+}
+
 !0 = !{}
 !1 = !{i64 4}
 !2 = !{i32 0, i32 100}