[DAE] Adjust param/arg attributes when changing parameter to undef

[Carrot]

Summary

In DeadArgumentElimination pass, if a function's argument is never used, corresponding caller's parameter can be changed to undef. If the param/arg has attribute noundef or other related attributes, LLVM LangRef(https://llvm.org/docs/LangRef.html#parameter-attributes) says its behavior is undefined. SimplifyCFG(D97244) takes advantage of this behavior and does bad transformation on valid code.

To avoid this undefined behavior when change caller's parameter to undef, this patch removes noundef attribute and other attributes imply noundef on param/arg.

[spatel]

Adding potential reviewers. I know there are some analysis helpers for this sort of thing, but not sure if those work for this usage.

[nikic]

This looks reasonable to me. Dropping nonnull is not strictly necessary (because it has no immediate UB), but doesn't hurt either.

I was going to suggest using AttrBuilder to drop all attributes at once, but I'm not sure that works with dereferenceable attributes.

[jdoerfert]

llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp

320

There should be a helper that populates a list of attributes we should drop when transforming something to undef.

[Carrot]

In D98899#2640159, @nikic wrote:

This looks reasonable to me. Dropping nonnull is not strictly necessary (because it has no immediate UB), but doesn't hurt either.

I was going to suggest using AttrBuilder to drop all attributes at once, but I'm not sure that works with dereferenceable attributes.

Thanks for the good suggestion.
According to the implementation of AttrBuilder::remove, it can remove dereferenceable attribute from another AttrBuilder with any dereferenceable bytes.

[Carrot]

Reformat the patch.

[jdoerfert]

LGTM with two nits. Others can chime in if they dislike the wording.

llvm/include/llvm/IR/Attributes.h

517–518

I'd make this more explicit, here and below.

520

Again, more explicit what this does, here and below.

[Carrot]

Thanks for the improvement.
Will commit this version.

llvm/include/llvm/IR/Attributes.h

public:
 
   /// Remove the specified attribute at the specified arg index from this
   /// attribute list. Returns a new list because attribute lists are immutable.
   LLVM_NODISCARD AttributeList removeParamAttributes(
       LLVMContext &C, unsigned ArgNo, const AttrBuilder &AttrsToRemove) const {
     return removeAttributes(C, ArgNo + FirstArgIndex, AttrsToRemove);
   }
 
+  /// Remove noundef attribute and other attributes imply undef from this
+  /// attribute list. Returns a new list because attribute lists are immutable.
+  LLVM_NODISCARD AttributeList removeParamNoundefAttributes(
+      LLVMContext &C, unsigned ArgNo) const;
+
   /// Remove all attributes at the specified arg index from this
   /// attribute list. Returns a new list because attribute lists are immutable.
   LLVM_NODISCARD AttributeList removeParamAttributes(LLVMContext &C,
                                                      unsigned ArgNo) const {
     return removeAttributes(C, ArgNo + FirstArgIndex);
   }
 
   /// Replace the type contained by attribute \p AttrKind at index \p ArgNo wih
   /// \p ReplacementTy, preserving all other attributes.
   LLVM_NODISCARD AttributeList replaceAttributeType(LLVMContext &C,
                                                     unsigned ArgNo,
                                                     Attribute::AttrKind Kind,
                                                     Type *ReplacementTy) const {

llvm/include/llvm/IR/Function.h

public:
   void removeParamAttr(unsigned ArgNo, Attribute::AttrKind Kind);
 
   /// removes the attribute from the list of attributes.
   void removeParamAttr(unsigned ArgNo, StringRef Kind);
 
   /// removes the attribute from the list of attributes.
   void removeParamAttrs(unsigned ArgNo, const AttrBuilder &Attrs);
 
+  /// removes noundef and related attributes from the list of attributes.
+  void removeParamNoundefAttrs(unsigned ArgNo);
+
   /// check if an attributes is in the list of attributes.
   bool hasAttribute(unsigned i, Attribute::AttrKind Kind) const {
     return getAttributes().hasAttribute(i, Kind);
   }
 
   /// check if an attributes is in the list of attributes.
   bool hasParamAttribute(unsigned ArgNo, Attribute::AttrKind Kind) const {
     return getAttributes().hasParamAttribute(ArgNo, Kind);

llvm/include/llvm/IR/InstrTypes.h

public:
   /// Removes the attribute from the given argument
   void removeParamAttr(unsigned ArgNo, StringRef Kind) {
     assert(ArgNo < getNumArgOperands() && "Out of bounds");
     AttributeList PAL = getAttributes();
     PAL = PAL.removeParamAttribute(getContext(), ArgNo, Kind);
     setAttributes(PAL);
   }
 
+  /// Removes noundef and related attributes from the given argument
+  void removeParamNoundefAttr(unsigned ArgNo) {
+    assert(ArgNo < getNumArgOperands() && "Out of bounds");
+    AttributeList PAL = getAttributes();
+    PAL = PAL.removeParamNoundefAttributes(getContext(), ArgNo);
+    setAttributes(PAL);
+  }
+
   /// adds the dereferenceable attribute to the list of attributes.
   void addDereferenceableAttr(unsigned i, uint64_t Bytes) {
     AttributeList PAL = getAttributes();
     PAL = PAL.addDereferenceableAttr(getContext(), i, Bytes);
     setAttributes(PAL);
   }
 
   /// adds the dereferenceable_or_null attribute to the list of

llvm/lib/IR/Attributes.cpp

AttributeList AttributeList::removeAttributes(LLVMContext &C,
   WithoutIndex = attrIdxToArrayIdx(WithoutIndex);
   if (WithoutIndex >= getNumAttrSets())
     return *this;
   SmallVector<AttributeSet, 4> AttrSets(this->begin(), this->end());
   AttrSets[WithoutIndex] = AttributeSet();
   return getImpl(C, AttrSets);
 }
 
+AttributeList AttributeList::removeParamNoundefAttributes(
+    LLVMContext &C, unsigned ArgNo) const {
+  AttrBuilder B;
+  B.addAttribute(Attribute::NoUndef);
+  B.addAttribute(Attribute::NonNull);
+  B.addDereferenceableAttr(1);
+  B.addDereferenceableOrNullAttr(1);
+  return removeParamAttributes(C, ArgNo, B);
+}
+
 AttributeList AttributeList::addDereferenceableAttr(LLVMContext &C,
                                                     unsigned Index,
                                                     uint64_t Bytes) const {
   AttrBuilder B;
   B.addDereferenceableAttr(Bytes);
   return addAttributes(C, Index, B);
 }
 
AttributeSet AttributeList::getRetAttributes() const {
   return getAttributes(ReturnIndex);
 }
 
 AttributeSet AttributeList::getFnAttributes() const {
   return getAttributes(FunctionIndex);
 }
 
 bool AttributeList::hasAttribute(unsigned Index,
                                  Attribute::AttrKind Kind) const {
   return getAttributes(Index).hasAttribute(Kind);
 }
 
 bool AttributeList::hasAttribute(unsigned Index, StringRef Kind) const {
   return getAttributes(Index).hasAttribute(Kind);
 }
 
 bool AttributeList::hasAttributes(unsigned Index) const {

llvm/lib/IR/Function.cpp

 }
 
 void Function::removeParamAttrs(unsigned ArgNo, const AttrBuilder &Attrs) {
   AttributeList PAL = getAttributes();
   PAL = PAL.removeParamAttributes(getContext(), ArgNo, Attrs);
   setAttributes(PAL);
 }
 
+void Function::removeParamNoundefAttrs(unsigned ArgNo) {
+  AttributeList PAL = getAttributes();
+  PAL = PAL.removeParamNoundefAttributes(getContext(), ArgNo);
+  setAttributes(PAL);
+}
+
 void Function::addDereferenceableAttr(unsigned i, uint64_t Bytes) {
   AttributeList PAL = getAttributes();
   PAL = PAL.addDereferenceableAttr(getContext(), i, Bytes);
   setAttributes(PAL);
 }
 
 void Function::addDereferenceableParamAttr(unsigned ArgNo, uint64_t Bytes) {
   AttributeList PAL = getAttributes();

llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp

bool DeadArgumentEliminationPass::RemoveDeadArgumentsFromCallers(Function &Fn) {
   for (Argument &Arg : Fn.args()) {
     if (!Arg.hasSwiftErrorAttr() && Arg.use_empty() &&
         !Arg.hasPassPointeeByValueCopyAttr()) {
       if (Arg.isUsedByMetadata()) {
         Arg.replaceAllUsesWith(UndefValue::get(Arg.getType()));
         Changed = true;
       }
       UnusedArgs.push_back(Arg.getArgNo());
+      Fn.removeParamNoundefAttrs(Arg.getArgNo());
     }
   }
 
   if (UnusedArgs.empty())
     return false;
 
   for (Use &U : Fn.uses()) {
     CallBase *CB = dyn_cast<CallBase>(U.getUser());
     if (!CB || !CB->isCallee(&U))
       continue;
 
     // Now go through all unused args and replace them with "undef".
     for (unsigned I = 0, E = UnusedArgs.size(); I != E; ++I) {
       unsigned ArgNo = UnusedArgs[I];
 
       Value *Arg = CB->getArgOperand(ArgNo);
       CB->setArgOperand(ArgNo, UndefValue::get(Arg->getType()));
+      CB->removeParamNoundefAttr(ArgNo);
+
       ++NumArgumentsReplacedWithUndef;
       Changed = true;
     }
   }
 
   return Changed;
 }
 
 /// Convenience function that returns the number of return values. It returns 0
 /// for void functions and 1 for functions not returning a struct. It returns
 /// the number of struct elements for functions returning a struct.

llvm/test/Transforms/DeadArgElim/NoundefAttrs.ll

+; RUN: opt -deadargelim -S < %s | FileCheck %s
+
+; If caller is changed to pass in undef, noundef and related attributes
+; should be deleted.
+
+
+; CHECK:   define i64 @bar(i64* %0, i64 %1)
+define i64 @bar(i64* nonnull dereferenceable(8) %0, i64 %1) {
+entry:
+  %2 = add i64 %1, 8
+  ret i64 %2
+}
+
+define i64 @foo(i64* %p, i64 %v) {
+; CHECK:   %retval = call i64 @bar(i64* undef, i64 %v)
+  %retval = call i64 @bar(i64* nonnull dereferenceable(8) %p, i64 %v)
+  ret i64 %retval
+}

llvm/test/Transforms/InstCombine/unused-nonnull.ll

 
 done:
   %retval = phi i32 [0, %entry], [%1, %do_work], [%1, %null]
   ret i32 %retval
 }
 
 define i32 @compute(i8* noundef nonnull %ptr, i32 %x) #1 {
 ; CHECK-LABEL: define {{[^@]+}}@compute
-; CHECK-SAME: (i8* nocapture noundef nonnull readnone [[PTR:%.*]], i32 returned [[X:%.*]]) local_unnamed_addr #1
+; CHECK-SAME: (i8* nocapture readnone [[PTR:%.*]], i32 returned [[X:%.*]]) local_unnamed_addr #1
 ; CHECK-NEXT:    ret i32 [[X]]
 ;
   ret i32 %x
 }
 
 declare void @call_if_null(i8* %ptr) #0
 
 attributes #0 = { nounwind }
 attributes #1 = { noinline nounwind readonly }