[LVI][ValueTracking] Take UB-implying attributes into account in isSafeToSpeculativelyExecute

llvm/include/llvm/Analysis/ValueTracking.h

@@ -539,31 +539,43 @@ bool isNotCrossLaneOperation(const Instruction *I);
 /// move the instruction as long as the correct dominance relationships for
 /// the operands and users hold.
 ///
+/// If \p UseVariableInfo is true, the information from non-constant operands
+/// will be taken into account.
+///
+/// If \p AllowRefinement is true, UB-implying attributes and metadata will be
+/// ignored. The caller is responsible for correctly propagating them after
+/// hoisting.
+///
 /// This method can return true for instructions that read memory;
 /// for such instructions, moving them may change the resulting value.
 bool isSafeToSpeculativelyExecute(const Instruction *I,
                                   const Instruction *CtxI = nullptr,
                                   AssumptionCache *AC = nullptr,
                                   const DominatorTree *DT = nullptr,
                                   const TargetLibraryInfo *TLI = nullptr,
-                                  bool UseVariableInfo = true);
+                                  bool UseVariableInfo = true,
+                                  bool AllowRefinement = true);
 
 inline bool isSafeToSpeculativelyExecute(const Instruction *I,
                                          BasicBlock::iterator CtxI,
                                          AssumptionCache *AC = nullptr,
                                          const DominatorTree *DT = nullptr,
                                          const TargetLibraryInfo *TLI = nullptr,
-                                         bool UseVariableInfo = true) {
+                                         bool UseVariableInfo = true,
+                                         bool AllowRefinement = true) {
   // Take an iterator, and unwrap it into an Instruction *.
-  return isSafeToSpeculativelyExecute(I, &*CtxI, AC, DT, TLI, UseVariableInfo);
+  return isSafeToSpeculativelyExecute(I, &*CtxI, AC, DT, TLI, UseVariableInfo,
+                                      AllowRefinement);
 }
 
 /// Don't use information from its non-constant operands. This helper is used
 /// when its operands are going to be replaced.
 inline bool
-isSafeToSpeculativelyExecuteWithVariableReplaced(const Instruction *I) {
+isSafeToSpeculativelyExecuteWithVariableReplaced(const Instruction *I,
+                                                 bool AllowRefinement = true) {
   return isSafeToSpeculativelyExecute(I, nullptr, nullptr, nullptr, nullptr,
-                                      /*UseVariableInfo=*/false);
+                                      /*UseVariableInfo=*/false,
+                                      AllowRefinement);
 }
 
 /// This returns the same result as isSafeToSpeculativelyExecute if Opcode is
@@ -586,7 +598,8 @@ isSafeToSpeculativelyExecuteWithVariableReplaced(const Instruction *I) {
 bool isSafeToSpeculativelyExecuteWithOpcode(
     unsigned Opcode, const Instruction *Inst, const Instruction *CtxI = nullptr,
     AssumptionCache *AC = nullptr, const DominatorTree *DT = nullptr,
-    const TargetLibraryInfo *TLI = nullptr, bool UseVariableInfo = true);
+    const TargetLibraryInfo *TLI = nullptr, bool UseVariableInfo = true,
+    bool AllowRefinement = true);
 
 /// Returns true if the result or effects of the given instructions \p I
 /// depend values not reachable through the def use graph.

llvm/lib/Analysis/LazyValueInfo.cpp

@@ -1701,7 +1701,8 @@ ValueLatticeElement LazyValueInfoImpl::getValueAtUse(const Use &U) {
     // of a cycle, we might end up reasoning about values from different cycle
     // iterations (PR60629).
     if (!CurrI->hasOneUse() ||
-        !isSafeToSpeculativelyExecuteWithVariableReplaced(CurrI))
+        !isSafeToSpeculativelyExecuteWithVariableReplaced(
+            CurrI, /*AllowRefinement=*/false))
       break;
     CurrU = &*CurrI->use_begin();
   }

llvm/lib/Analysis/ValueTracking.cpp

@@ -7201,20 +7201,19 @@ bool llvm::isNotCrossLaneOperation(const Instruction *I) {
          !isa<CallBase, BitCastInst, ExtractElementInst>(I);
 }
 
-bool llvm::isSafeToSpeculativelyExecute(const Instruction *Inst,
-                                        const Instruction *CtxI,
-                                        AssumptionCache *AC,
-                                        const DominatorTree *DT,
-                                        const TargetLibraryInfo *TLI,
-                                        bool UseVariableInfo) {
+bool llvm::isSafeToSpeculativelyExecute(
+    const Instruction *Inst, const Instruction *CtxI, AssumptionCache *AC,
+    const DominatorTree *DT, const TargetLibraryInfo *TLI, bool UseVariableInfo,
+    bool AllowRefinement) {
   return isSafeToSpeculativelyExecuteWithOpcode(Inst->getOpcode(), Inst, CtxI,
-                                                AC, DT, TLI, UseVariableInfo);
+                                                AC, DT, TLI, UseVariableInfo,
+                                                AllowRefinement);
 }
 
 bool llvm::isSafeToSpeculativelyExecuteWithOpcode(
     unsigned Opcode, const Instruction *Inst, const Instruction *CtxI,
     AssumptionCache *AC, const DominatorTree *DT, const TargetLibraryInfo *TLI,
-    bool UseVariableInfo) {
+    bool UseVariableInfo, bool AllowRefinement) {
 #ifndef NDEBUG
   if (Inst->getOpcode() != Opcode) {
     // Check that the operands are actually compatible with the Opcode override.
@@ -7266,7 +7265,7 @@ bool llvm::isSafeToSpeculativelyExecuteWithOpcode(
     return false;
   }
   case Instruction::Load: {
-    if (!UseVariableInfo)
+    if (!UseVariableInfo || !AllowRefinement)
       return false;
 
     const LoadInst *LI = dyn_cast<LoadInst>(Inst);
@@ -7287,7 +7286,20 @@ bool llvm::isSafeToSpeculativelyExecuteWithOpcode(
 
     // The called function could have undefined behavior or side-effects, even
     // if marked readnone nounwind.
-    return Callee && Callee->isSpeculatable();
+    if (!Callee || !Callee->isSpeculatable())
+      return false;
+    // Since the operands may be changed after hoisting, undefined behavior may
+    // be triggered by some UB-implying attributes.
+    if (!AllowRefinement) {
+      if (CI->hasRetAttr(Attribute::NoUndef) ||
+          CI->getRetDereferenceableBytes() > 0 ||
+          CI->getRetDereferenceableOrNullBytes() > 0 ||
+          any_of(CI->args(), [&](const Use &U) {
+            return CI->isPassingUndefUB(CI->getArgOperandNo(&U));
+          }))
+        return false;
+    }
+    return true;
   }
   case Instruction::VAArg:
   case Instruction::Alloca:

llvm/test/Transforms/CorrelatedValuePropagation/pr137582.ll

@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=correlated-propagation -S | FileCheck %s
+
+; Make sure that the optimization does not introduce immediate UB.
+
+define i8 @test(i16 %x) {
+; CHECK-LABEL: define range(i8 -128, 1) i8 @test(
+; CHECK-SAME: i16 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[OR:%.*]] = or i16 [[X]], 1
+; CHECK-NEXT:    [[CONV:%.*]] = trunc i16 [[OR]] to i8
+; CHECK-NEXT:    [[MIN:%.*]] = call noundef i8 @llvm.smin.i8(i8 [[CONV]], i8 0)
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i16 [[X]], 0
+; CHECK-NEXT:    br i1 [[COND]], label %[[IF_END:.*]], label %[[IF_THEN:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    br label %[[IF_END]]
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[RES:%.*]] = phi i8 [ [[MIN]], %[[ENTRY]] ], [ 0, %[[IF_THEN]] ]
+; CHECK-NEXT:    ret i8 [[RES]]
+;
+entry:
+  %or = or i16 %x, 1
+  %conv = trunc i16 %or to i8
+  %min = call noundef i8 @llvm.smin.i8(i8 %conv, i8 0)
+  %cond = icmp eq i16 %x, 0
+  br i1 %cond, label %if.end, label %if.then
+
+if.then:
+  br label %if.end
+
+if.end:
+  %res = phi i8 [ %min, %entry ], [ 0, %if.then ]
+  ret i8 %res
+}