[SimplifyCFG][LICM] Preserve nonnull, range and align metadata when speculating

[nikic]

Summary

After D141386, violation of nonnull, range and align metadata results in poison rather than immediate undefined behavior, which means that these are now safe to retain when speculating. We only need to remove UB-implying metadata like noundef.

This is done by adding a dropUBImplyingAttrsAndMetadata() helper, which lists the metadata which is known safe to retain on speculation.

[jdoerfert]

This is valid, I think.
That said, I went to try what could go wrong and it was easy.
At least early-cse is buggy so this will cause issues until those bugs are fixed:
https://godbolt.org/z/Tz98hYa9E

FWIW, instcombine gets it right regardless of the order. Might be luck or logic to properly merge the attributes.

[nikic]

@jdoerfert Yeah, there is a pending patch that fixes this for !range here: https://reviews.llvm.org/D142687 We'll need similar changes for nonnull/align as well.

I'll take this off the review queue until those changes have all landed.

[jdoerfert]

In D146629#4213446, @nikic wrote:

@jdoerfert Yeah, there is a pending patch that fixes this for !range here: https://reviews.llvm.org/D142687 We'll need similar changes for nonnull/align as well.

I'll take this off the review queue until those changes have all landed.

Just a thought:
Can we create a test that we somehow run against all transformations (one by one) to verify they don't combine good metadata?

[StephenFan]

In D146629#4213401, @jdoerfert wrote:

This is valid, I think.
That said, I went to try what could go wrong and it was easy.
At least early-cse is buggy so this will cause issues until those bugs are fixed:
https://godbolt.org/z/Tz98hYa9E

FWIW, instcombine gets it right regardless of the order. Might be luck or logic to properly merge the attributes.

Does it seem that we even haven't called combineMetadata in early-cse?

[nikic]

Rebase. I believe all the known metadata handling issues are now fixed. https://github.com/llvm/llvm-project/commit/9b5ff4436e446e9df97ee37b3bcf9ba029c7d9aa is the fix for EarlyCSE.

[StephenFan]

LGTM

[durin42]

FYI, we're investigating a pretty big regression in Rust when built with LLVM containing this change. I don't have much yet, but

io::buffered::tests::line_buffer_write0_normal
io::tests::chain_zero_length_read_is_not_eof
io::tests::read_exact
net::tcp::tests::timeouts
net::udp::tests::timeouts
os::unix::net::tests::timeouts

in library/std tests are all broken by this change. I'm working on getting something more now.

[durin42]

I got a funny sense of things and decided to try again with

diff --git a/llvm/lib/IR/Instruction.cpp b/llvm/lib/IR/Instruction.cpp
index f14737e73fde..849a22ec9c2f 100644
--- a/llvm/lib/IR/Instruction.cpp
+++ b/llvm/lib/IR/Instruction.cpp
@@ -249,7 +249,7 @@ void Instruction::dropUBImplyingAttrsAndMetadata() {
   // !noundef and various AA metadata must be dropped, as it generally produces
   // immediate undefined behavior.
   unsigned KnownIDs[] = {LLVMContext::MD_annotation, LLVMContext::MD_range,
-                         LLVMContext::MD_nonnull, LLVMContext::MD_align};
+                         LLVMContext::MD_align};
   dropUBImplyingAttrsAndUnknownMetadata(KnownIDs);
 }

and the tests all pass. I hope that's useful insight, I'm going to try and keep working on getting some IR out of one of the broken tests.

[nikic]

Reduced to a miscompile in GVN:

declare void @use(i64)

define void @test(ptr %p) {
  %val = load ptr, ptr %p, align 8, !nonnull !{}
  %val.int = ptrtoint ptr %val to i64
  %val2 = load i64, ptr %p, align 8
  call void @use(i64 %val.int)
  call void @use(i64 %val2)
  ret void
}

; RUN: opt -passes=gvn < %s

define void @test(ptr %p) {
  %val = load ptr, ptr %p, align 8, !nonnull !0
  %val.int = ptrtoint ptr %val to i64
  call void @use(i64 %val.int)
  call void @use(i64 %val.int)
  ret void
}

I believe the problem here is that this goes through VNCoercion and the replacement we do is not with a plain load, but a ptrtoint cast of the load. As such, we end up not performing the combineMetadataForCSE adjustment.

[nikic]

Looks like there's another bug in GVN:

declare void @use(i32)

define void @test(i1 %c, ptr %p) {
  %v1 = load i32, ptr %p, !range !{i32 0, i32 10}
  br i1 %c, label %if, label %join

if:
  br label %join

join:
  %v2 = load i32, ptr %p, !range !{i32 20, i32 30}
  call void @use(i32 %v1)
  call void @use(i32 %v2)
  ret void
}

; RUN: opt -S -passes=gvn < %s

declare void @use(i32)

define void @test(i1 %c, ptr %p) {
  %v1 = load i32, ptr %p, align 4, !range !0
  br i1 %c, label %if, label %join

if:                                               ; preds = %0
  br label %join

join:                                             ; preds = %if, %0
  call void @use(i32 %v1)
  call void @use(i32 %v1)
  ret void
}

!0 = !{i32 0, i32 10}

Apparently we don't perform metadata combining when replacing non-local loads :(

llvm/include/llvm/IR/Instruction.h

public:
 
   /// Drops flags and metadata that may generate poison.
   void dropPoisonGeneratingFlagsAndMetadata() {
     dropPoisonGeneratingFlags();
     dropPoisonGeneratingMetadata();
   }
 
   /// This function drops non-debug unknown metadata (through
-  /// dropUnknownNonDebugMetadata). For calls, it also drops parameter and 
+  /// 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 dropUBImplyingAttrsAndUnknownMetadata(ArrayRef<unsigned> KnownIDs = {});
 
+  /// Drop any attributes or metadata that can cause immediate undefined
+  /// behavior. Retain other attributes/metadata on a best-effort basis.
+  /// This should be used when speculating instructions.
+  void dropUBImplyingAttrsAndMetadata();
+
   /// Determine whether the exact flag is set.
   bool isExact() const LLVM_READONLY;
 
   /// Set or clear all fast-math-flags on this instruction, which must be an
   /// operator which supports this flag. See LangRef.html for the meaning of
   /// this flag.
   void setFast(bool B);
 

llvm/lib/IR/Instruction.cpp

if (AL.isEmpty())
     return;
   AttributeMask UBImplyingAttributes =
       AttributeFuncs::getUBImplyingAttributes();
   for (unsigned ArgNo = 0; ArgNo < CB->arg_size(); ArgNo++)
     CB->removeParamAttrs(ArgNo, UBImplyingAttributes);
   CB->removeRetAttrs(UBImplyingAttributes);
 }
 
+void Instruction::dropUBImplyingAttrsAndMetadata() {
+  // !annotation metadata does not impact semantics.
+  // !range, !nonnull and !align produce poison, so they are safe to speculate.
+  // !noundef and various AA metadata must be dropped, as it generally produces
+  // immediate undefined behavior.
+  unsigned KnownIDs[] = {LLVMContext::MD_annotation, LLVMContext::MD_range,
+                         LLVMContext::MD_nonnull, LLVMContext::MD_align};
+  dropUBImplyingAttrsAndUnknownMetadata(KnownIDs);
+}
+
 bool Instruction::isExact() const {
   return cast<PossiblyExactOperator>(this)->isExact();
 }
 
 void Instruction::setFast(bool B) {
   assert(isa<FPMathOperator>(this) && "setting fast-math flag on invalid op");
   cast<FPMathOperator>(this)->setFast(B);
 }

llvm/lib/Transforms/Scalar/LICM.cpp

static void hoist(Instruction &I, const DominatorTree *DT, const Loop *CurLoop,
   // 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.dropUBImplyingAttrsAndUnknownMetadata();
+    I.dropUBImplyingAttrsAndMetadata();
 
   if (isa<PHINode>(I))
     // Move the new node to the end of the phi list in the destination block.
     moveInstructionBefore(I, *Dest->getFirstNonPHI(), *SafetyInfo, MSSAU, SE);
   else
     // Move the new node to the destination block, before its terminator.
     moveInstructionBefore(I, *Dest->getTerminator(), *SafetyInfo, MSSAU, SE);
 

llvm/lib/Transforms/Utils/Local.cpp

void llvm::hoistAllInstructionsInto(BasicBlock *DomBlock, Instruction *InsertPt,
   // See PR38762, PR39243 for more details.
   //
   // TODO: Extend llvm.dbg.value to take more than one SSA Value (PR39141) to
   // encode predicated DIExpressions that yield different results on different
   // code paths.
 
   for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
     Instruction *I = &*II;
-    I->dropUBImplyingAttrsAndUnknownMetadata();
+    I->dropUBImplyingAttrsAndMetadata();
     if (I->isUsedByMetadata())
       dropDebugUsers(*I);
     if (I->isDebugOrPseudoInst()) {
       // Remove DbgInfo and pseudo probe Intrinsics.
       II = I->eraseFromParent();
       continue;
     }
     I->setDebugLoc(InsertPt->getDebugLoc());

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

if (PTI->getDebugLoc() != NewBonusInst->getDebugLoc()) {
       // dead code caused by folding dead branches.
       NewBonusInst->setDebugLoc(DebugLoc());
     }
 
     RemapInstruction(NewBonusInst, VMap,
                      RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
     VMap[&BonusInst] = NewBonusInst;
 
-    // If we moved a load, we cannot any longer claim any knowledge about
-    // its potential value. The previous information might have been valid
+    // If we speculated an instruction, we need to drop any metadata that may
+    // result in undefined behavior, as the metadata might have been valid
     // only given the branch precondition.
-    // 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.
     // Similarly strip attributes on call parameters that may cause UB in
     // location the call is moved to.
-    NewBonusInst->dropUBImplyingAttrsAndUnknownMetadata(
-        LLVMContext::MD_annotation);
+    NewBonusInst->dropUBImplyingAttrsAndMetadata();
 
     NewBonusInst->insertInto(PredBlock, PTI->getIterator());
     NewBonusInst->takeName(&BonusInst);
     BonusInst.setName(NewBonusInst->getName() + ".old");
 
     // Update (liveout) uses of bonus instructions,
     // now that the bonus instruction has been cloned into predecessor.
     // Note that we expect to be in a block-closed SSA form for this to work!
if (SpeculatedStoreValue) {
     //   dbg.assign %merge, "x", ..., !2
     for (auto *DAI : at::getAssignmentMarkers(SpeculatedStore)) {
       if (any_of(DAI->location_ops(), [&](Value *V) { return V == OrigV; }))
         DAI->replaceVariableLocationOp(OrigV, S);
     }
   }
 
   // Metadata can be dependent on the condition we are hoisting above.
-  // Conservatively strip all metadata on the instruction. Drop the debug loc
+  // Strip all UB-implying 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 : make_early_inc_range(*ThenBB)) {
     if (!SpeculatedStoreValue || &I != SpeculatedStore) {
       // Don't update the DILocation of dbg.assign intrinsics.
       if (!isa<DbgAssignIntrinsic>(&I))
         I.setDebugLoc(DebugLoc());
     }
-    I.dropUBImplyingAttrsAndUnknownMetadata();
+    I.dropUBImplyingAttrsAndMetadata();
 
     // Drop ephemeral values.
     if (EphTracker.contains(&I)) {
       I.replaceAllUsesWith(PoisonValue::get(I.getType()));
       I.eraseFromParent();
     }
   }
 

llvm/test/Transforms/LICM/hoist-metadata.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --version 2
 ; RUN: opt -S -passes=licm < %s| FileCheck %s
 
 declare void @foo(...) memory(none)
 
+; We can preserve all metadata on instructions that are guaranteed to execute.
 define void @test_unconditional(i1 %c, ptr dereferenceable(8) align 8 %p) {
 ; CHECK-LABEL: define void @test_unconditional
 ; CHECK-SAME: (i1 [[C:%.*]], ptr align 8 dereferenceable(8) [[P:%.*]]) {
 ; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[P]], align 4, !range [[RNG0:![0-9]+]]
 ; CHECK-NEXT:    [[V2:%.*]] = load ptr, ptr [[P]], align 8, !nonnull !1, !noundef !1
 ; CHECK-NEXT:    [[V3:%.*]] = load ptr, ptr [[P]], align 8, !dereferenceable !2, !align !2
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
loop:
   %v3 = load ptr, ptr %p, !align !{i64 4}, !dereferenceable !{i64 4}
   call void @foo(i32 %v1, ptr %v2, ptr %v3)
   br i1 %c, label %loop, label %exit
 
 exit:
   ret void
 }
 
+; We cannot preserve UB-implying metadata on instructions that are speculated.
+; However, we can preserve poison-implying metadata.
 define void @test_conditional(i1 %c, i1 %c2, ptr dereferenceable(8) align 8 %p) {
 ; CHECK-LABEL: define void @test_conditional
 ; CHECK-SAME: (i1 [[C:%.*]], i1 [[C2:%.*]], ptr align 8 dereferenceable(8) [[P:%.*]]) {
-; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[P]], align 4
-; CHECK-NEXT:    [[V2:%.*]] = load ptr, ptr [[P]], align 8
-; CHECK-NEXT:    [[V3:%.*]] = load ptr, ptr [[P]], align 8
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[P]], align 4, !range [[RNG0]]
+; CHECK-NEXT:    [[V2:%.*]] = load ptr, ptr [[P]], align 8, !nonnull !1
+; CHECK-NEXT:    [[V3:%.*]] = load ptr, ptr [[P]], align 8, !align !2
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    br i1 [[C]], label [[IF:%.*]], label [[LATCH:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    call void @foo(i32 [[V1]], ptr [[V2]], ptr [[V3]])
 ; CHECK-NEXT:    br label [[LATCH]]
 ; CHECK:       latch:
 ; CHECK-NEXT:    br i1 [[C2]], label [[LOOP]], label [[EXIT:%.*]]

llvm/test/Transforms/SimplifyCFG/hoist-with-metadata.ll

 else:
   %e = load i8, ptr %p
   br label %out
 
 out:
   ret void
 }
 
+; !range violation only returns poison, and is thus safe to speculate.
 define i32 @speculate_range(i1 %c, ptr dereferenceable(8) align 8 %p) {
 ; CHECK-LABEL: @speculate_range(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[P:%.*]], align 4
+; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[P:%.*]], align 4, !range [[RNG2:![0-9]+]]
 ; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], i32 [[V]], i32 0
 ; CHECK-NEXT:    ret i32 [[SPEC_SELECT]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   %v = load i32, ptr %p, !range !{i32 0, i32 10}
   br label %join
 
 join:
   %phi = phi i32 [ %v, %if ], [ 0, %entry ]
   ret i32 %phi
 }
 
+; !nonnull is safe to speculate, but !noundef is not, as the latter causes
+; immediate undefined behavior.
 define ptr @speculate_nonnull(i1 %c, ptr dereferenceable(8) align 8 %p) {
 ; CHECK-LABEL: @speculate_nonnull(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[V:%.*]] = load ptr, ptr [[P:%.*]], align 8
+; CHECK-NEXT:    [[V:%.*]] = load ptr, ptr [[P:%.*]], align 8, !nonnull !1
 ; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], ptr [[V]], ptr null
 ; CHECK-NEXT:    ret ptr [[SPEC_SELECT]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   %v = load ptr, ptr %p, !nonnull !{}, !noundef !{}
   br label %join
 
 join:
   %phi = phi ptr [ %v, %if ], [ null, %entry ]
   ret ptr %phi
 }
 
-
+; !align is safe to speculate, but !dereferenceable is not, as the latter causes
+; immediate undefined behavior.
 define ptr @speculate_align(i1 %c, ptr dereferenceable(8) align 8 %p) {
 ; CHECK-LABEL: @speculate_align(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[V:%.*]] = load ptr, ptr [[P:%.*]], align 8
+; CHECK-NEXT:    [[V:%.*]] = load ptr, ptr [[P:%.*]], align 8, !align !3
 ; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], ptr [[V]], ptr null
 ; CHECK-NEXT:    ret ptr [[SPEC_SELECT]]
 ;
 entry:
   br i1 %c, label %if, label %join
 
 if:
   %v = load ptr, ptr %p, !align !{i64 4}, !dereferenceable !{i64 4}
   br label %join
 
 join:
   %phi = phi ptr [ %v, %if ], [ null, %entry ]
   ret ptr %phi
 }
 
 define void @hoist_fpmath(i1 %c, double %x) {
 ; CHECK-LABEL: @hoist_fpmath(
 ; CHECK-NEXT:  if:
-; CHECK-NEXT:    [[T:%.*]] = fadd double [[X:%.*]], 1.000000e+00, !fpmath !2
+; CHECK-NEXT:    [[T:%.*]] = fadd double [[X:%.*]], 1.000000e+00, !fpmath !4
 ; CHECK-NEXT:    ret void
 ;
 if:
   br i1 %c, label %then, label %else
 then:
   %t = fadd double %x, 1.0, !fpmath !{ float 2.5 }
   br label %out
 else:
   %e = fadd double %x, 1.0, !fpmath !{ float 5.0 }
   br label %out
 out:
   ret void
 }
 
 !0 = !{ i8 0, i8 1 }
 !1 = !{ i8 3, i8 5 }
 !2 = !{}
 ;.
 ; CHECK: [[RNG0]] = !{i8 0, i8 1, i8 3, i8 5}
 ; CHECK: [[META1:![0-9]+]] = !{}
-; CHECK: [[META2:![0-9]+]] = !{float 2.500000e+00}
+; CHECK: [[RNG2]] = !{i32 0, i32 10}
+; CHECK: [[META3:![0-9]+]] = !{i64 4}
+; CHECK: [[META4:![0-9]+]] = !{float 2.500000e+00}
 ;.