Strip undef implying attributes when moving calls

When hoisting/moving calls to locations, we strip unknown metadata. Such calls are usually marked `speculatable`, i.e. they are guaranteed to not cause undefined behaviour when run anywhere. So, we should strip attributes that can cause immediate undefined behaviour if those attributes are not valid in the context where the call is moved to.

This patch introduces such an API and uses it in relevant passes. See
updated tests.

Fix for PR50744.

Reviewed By: nikic, jdoerfert, lebedev.ri

Differential Revision: https://reviews.llvm.org/D104641

include/llvm/IR/Instruction.h

@@ -332,6 +332,8 @@ public:
   /// @{
   /// Passes are required to drop metadata they don't understand. This is a
   /// convenience method for passes to do so.
+  /// dropUndefImplyingAttrsAndUnknownMetadata should be used instead of
+  /// this API if the Instruction being modified is a call.
   void dropUnknownNonDebugMetadata(ArrayRef<unsigned> KnownIDs);
   void dropUnknownNonDebugMetadata() {
     return dropUnknownNonDebugMetadata(None);
@@ -391,6 +393,13 @@ public:
   /// having non-poison inputs.
   void dropPoisonGeneratingFlags();
 
+  /// This function drops non-debug unknown metadata (through
+  /// dropUnknownNonDebugMetadata). For calls, it also drops parameter and 
+  /// return attributes that can cause undefined behaviour. Both of these should
+  /// be done by passes which move instructions in IR.
+  void
+  dropUndefImplyingAttrsAndUnknownMetadata(ArrayRef<unsigned> KnownIDs = {});
+
   /// Determine whether the exact flag is set.
   bool isExact() const;
 

lib/IR/Instruction.cpp

@@ -165,6 +165,23 @@ void Instruction::dropPoisonGeneratingFlags() {
   // TODO: FastMathFlags!
 }
 
+void Instruction::dropUndefImplyingAttrsAndUnknownMetadata(
+    ArrayRef<unsigned> KnownIDs) {
+  dropUnknownNonDebugMetadata(KnownIDs);
+  auto *CB = dyn_cast<CallBase>(this);
+  if (!CB)
+    return;
+  // For call instructions, we also need to drop parameter and return attributes
+  // that are can cause UB if the call is moved to a location where the
+  // attribute is not valid.
+  AttributeList AL = CB->getAttributes();
+  if (AL.isEmpty())
+    return;
+  AttrBuilder UBImplyingAttributes = AttributeFuncs::getUBImplyingAttributes();
+  for (unsigned ArgNo = 0; ArgNo < CB->getNumArgOperands(); ArgNo++)
+    CB->removeParamAttrs(ArgNo, UBImplyingAttributes);
+  CB->removeAttributes(AttributeList::ReturnIndex, UBImplyingAttributes);
+}
 
 bool Instruction::isExact() const {
   return cast<PossiblyExactOperator>(this)->isExact();

lib/Transforms/Scalar/LICM.cpp

@@ -1810,12 +1810,16 @@ static void hoist(Instruction &I, const DominatorTree *DT, const Loop *CurLoop,
   // Conservatively strip all metadata on the instruction unless we were
   // guaranteed to execute I if we entered the loop, in which case the metadata
   // is valid in the loop preheader.
-  if (I.hasMetadataOtherThanDebugLoc() &&
+  // Similarly, If I is a call and it is not guaranteed to execute in the loop,
+  // then moving to the preheader means we should strip attributes on the call
+  // that can cause UB since we may be hoisting above conditions that allowed
+  // inferring those attributes. They may not be valid at the preheader.
+  if ((I.hasMetadataOtherThanDebugLoc() || isa<CallInst>(I)) &&
       // The check on hasMetadataOtherThanDebugLoc is to prevent us from burning
       // time in isGuaranteedToExecute if we don't actually have anything to
       // drop.  It is a compile time optimization, not required for correctness.
       !SafetyInfo->isGuaranteedToExecute(I, DT, CurLoop))
-    I.dropUnknownNonDebugMetadata();
+    I.dropUndefImplyingAttrsAndUnknownMetadata();
 
   if (isa<PHINode>(I))
     // Move the new node to the end of the phi list in the destination block.

lib/Transforms/Utils/Local.cpp

@@ -2822,7 +2822,7 @@ void llvm::hoistAllInstructionsInto(BasicBlock *DomBlock, Instruction *InsertPt,
 
   for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
     Instruction *I = &*II;
-    I->dropUnknownNonDebugMetadata();
+    I->dropUndefImplyingAttrsAndUnknownMetadata();
     if (I->isUsedByMetadata())
       dropDebugUsers(*I);
     if (I->isDebugOrPseudoInst()) {

lib/Transforms/Utils/SimplifyCFG.cpp

@@ -1083,7 +1083,10 @@ static void CloneInstructionsIntoPredecessorBlockAndUpdateSSAUses(
     // For an analogous reason, we must also drop all the metadata whose
     // semantics we don't understand. We *can* preserve !annotation, because
     // it is tied to the instruction itself, not the value or position.
-    NewBonusInst->dropUnknownNonDebugMetadata(LLVMContext::MD_annotation);
+    // Similarly strip attributes on call parameters that may cause UB in
+    // location the call is moved to.
+    NewBonusInst->dropUndefImplyingAttrsAndUnknownMetadata(
+        LLVMContext::MD_annotation);
 
     PredBlock->getInstList().insert(PTI->getIterator(), NewBonusInst);
     NewBonusInst->takeName(&BonusInst);
@@ -2492,10 +2495,12 @@ bool SimplifyCFGOpt::SpeculativelyExecuteBB(BranchInst *BI, BasicBlock *ThenBB,
   // Conservatively strip all metadata on the instruction. Drop the debug loc
   // to avoid making it appear as if the condition is a constant, which would
   // be misleading while debugging.
+  // Similarly strip attributes that maybe dependent on condition we are
+  // hoisting above.
   for (auto &I : *ThenBB) {
     if (!SpeculatedStoreValue || &I != SpeculatedStore)
       I.setDebugLoc(DebugLoc());
-    I.dropUnknownNonDebugMetadata();
+    I.dropUndefImplyingAttrsAndUnknownMetadata();
   }
 
   // Hoist the instructions.

test/Transforms/LICM/call-hoisting.ll

@@ -23,16 +23,22 @@ exit:
   ret void
 }
 
-declare i32 @spec(i32* %p) readonly argmemonly nounwind speculatable
+declare i32 @spec(i32* %p, i32* %q) readonly argmemonly nounwind speculatable
 
-; FIXME: We should strip the nonnull callsite attribute on spec call's argument since it is
-; may not be valid when hoisted to preheader.
+; We should strip the dereferenceable callsite attribute on spec call's argument since it is
+; can cause UB in the speculatable call when hoisted to preheader.
+; However, we need not strip the nonnull attribute since it just propagates
+; poison if the parameter was indeed null.
 define void @test_strip_attribute(i32* noalias %loc, i32* noalias %sink, i32* %q) {
-; CHECK-LABEL: test_strip_attribute
-; CHECK-LABEL: entry
-; CHECK-NEXT:   %ret = call i32 @load(i32* %loc)
-; CHECK-NEXT:   %nullchk = icmp eq i32* %q, null
-; CHECK-NEXT:   %ret2 = call i32 @spec(i32* nonnull %q)
+; CHECK-LABEL: @test_strip_attribute(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[RET:%.*]] = call i32 @load(i32* [[LOC:%.*]])
+; CHECK-NEXT:    [[NULLCHK:%.*]] = icmp eq i32* [[Q:%.*]], null
+; CHECK-NEXT:    [[RET2:%.*]] = call i32 @spec(i32* nonnull [[Q]], i32* [[LOC]])
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[ISNULL:%.*]] ]
+; CHECK-NEXT:    br i1 [[NULLCHK]], label [[ISNULL]], label [[NONNULLBB:%.*]]
 entry:
   br label %loop
 
@@ -43,7 +49,7 @@ loop:
   br i1 %nullchk, label %isnull, label %nonnullbb
 
 nonnullbb:
-  %ret2 = call i32 @spec(i32* nonnull %q)
+  %ret2 = call i32 @spec(i32* nonnull %q, i32* dereferenceable(12) %loc)
   br label %isnull
 
 isnull:

test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll

@@ -118,8 +118,9 @@ final_right:
   ret void
 }
 
-; FIXME: When we fold the dispatch block into pred, the call is moved to pred
-; and the attribute nonnull is no longer valid on paramater. We should drop it.
+; When we fold the dispatch block into pred, the call is moved to pred
+; and the attribute nonnull propagates poison paramater. However, since the
+; function is speculatable, it can never cause UB. So, we need not technically drop it.
 define void @one_pred_with_spec_call(i8 %v0, i8 %v1, i32* %p) {
 ; CHECK-LABEL: @one_pred_with_spec_call(
 ; CHECK-NEXT:  pred:
@@ -133,7 +134,6 @@ define void @one_pred_with_spec_call(i8 %v0, i8 %v1, i32* %p) {
 ; CHECK:       final_right:
 ; CHECK-NEXT:    call void @sideeffect0()
 ; CHECK-NEXT:    br label [[COMMON_RET]]
-;
 pred:
   %c0 = icmp ne i32* %p, null
   br i1 %c0, label %dispatch, label %final_right
@@ -151,6 +151,29 @@ final_right:
   ret void
 }
 
+; Drop dereferenceable on the parameter
+define void @one_pred_with_spec_call_deref(i8 %v0, i8 %v1, i32* %p) {
+; CHECK-LABEL: one_pred_with_spec_call_deref
+; CHECK-LABEL: pred:
+; CHECK:         %c0 = icmp ne i32* %p, null
+; CHECK:         %x = call i32 @speculate_call(i32* %p)
+pred:
+  %c0 = icmp ne i32* %p, null
+  br i1 %c0, label %dispatch, label %final_right
+
+dispatch:
+  %x = call i32 @speculate_call(i32* dereferenceable(12) %p)
+  %c1 = icmp eq i8 %v1, 0
+  br i1 %c1, label %final_left, label %final_right
+
+final_left:
+  ret void
+
+final_right:
+  call void @sideeffect0()
+  ret void
+}
+
 define void @two_preds_with_extra_op(i8 %v0, i8 %v1, i8 %v2, i8 %v3) {
 ; CHECK-LABEL: @two_preds_with_extra_op(
 ; CHECK-NEXT:  entry:

test/Transforms/SimplifyCFG/speculate-call.ll

@@ -20,20 +20,23 @@ define i32 @func() #0 {
   ret i32 1
 }
 
-; FIXME: We should correctly drop the attribute since it may no longer be valid
+; We should correctly drop the attribute since it may no longer be valid
 ; in the context the call is moved to.
+; Since the function is speculatable, the nonnull attribute need not be dropped
+; since it propagates poison (and call executes fine) if the parameter is indeed
+; null.
 define i32 @strip_attr(i32 * %p) {
 ; CHECK-LABEL: strip_attr  
 ; CHECK-LABEL: entry:
 ; CHECK:         %nullchk = icmp ne i32* %p, null
-; CHECK:         %val = call i32 @func_deref(i32* nonnull %p)
+; CHECK:         %val = call i32 @func_nonnull(i32* nonnull %p)
 ; CHECK:         select 
 entry:
   %nullchk = icmp ne i32* %p, null
   br i1 %nullchk, label %if, label %end
 
 if:
-  %val = call i32 @func_deref(i32* nonnull %p) #1
+  %val = call i32 @func_nonnull(i32* nonnull %p) #1
   br label %end
 
 end:
@@ -41,7 +44,28 @@ end:
   ret i32 %ret
 }
 
-declare i32 @func_deref(i32*) #1
+; We should strip the deref attribute since it can cause UB when the
+; speculatable call is moved.
+define i32 @strip_attr2(i32 * %p) {
+; CHECK-LABEL: strip_attr2  
+; CHECK-LABEL: entry:
+; CHECK:         %nullchk = icmp ne i32* %p, null
+; CHECK:         %val = call i32 @func_nonnull(i32* %p)
+; CHECK:         select 
+entry:
+  %nullchk = icmp ne i32* %p, null
+  br i1 %nullchk, label %if, label %end
+
+if:
+  %val = call i32 @func_nonnull(i32* dereferenceable(12) %p) #1
+  br label %end
+
+end:
+  %ret = phi i32 [%val, %if], [0, %entry]
+  ret i32 %ret
+}
+
+declare i32 @func_nonnull(i32*) #1
 
 attributes #0 = { nounwind readnone speculatable }
 attributes #1 = { nounwind argmemonly speculatable }