[LVI][ValueTracking] Take UB-implying attributes into account in isSafeToSpeculativelyExecute (#137604)

Closes https://github.com/llvm/llvm-project/issues/137582.

In the original case, LVI uses the edge information in `%entry ->
%if.end` to get a more precise result. However, since the call to `smin`
has an `noundef` return attribute, an immediate UB will be triggered
after optimization.

Currently, `isSafeToSpeculativelyExecuteWithOpcode(%min)` returns true
because
https://github.com/llvm/llvm-project/commit/6a288c1e32351d4be3b7630841af078fa1c3bb8b
only checks whether the function is speculatable. However, it is not
enough in this case.

This patch takes UB-implying attributes into account if
`IgnoreUBImplyingAttrs` is set to false. If it is set to true, the
caller is responsible for correctly propagating UB-implying attributes.

llvm/include/llvm/Analysis/ValueTracking.h

@@ -539,6 +539,12 @@ 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 IgnoreUBImplyingAttrs is true, UB-implying attributes 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,
@@ -546,24 +552,28 @@ bool isSafeToSpeculativelyExecute(const Instruction *I,
                                   AssumptionCache *AC = nullptr,
                                   const DominatorTree *DT = nullptr,
                                   const TargetLibraryInfo *TLI = nullptr,
-                                  bool UseVariableInfo = true);
+                                  bool UseVariableInfo = true,
+                                  bool IgnoreUBImplyingAttrs = 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 IgnoreUBImplyingAttrs = 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,
+                                      IgnoreUBImplyingAttrs);
 }
 
 /// 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) {
+inline bool isSafeToSpeculativelyExecuteWithVariableReplaced(
+    const Instruction *I, bool IgnoreUBImplyingAttrs = true) {
   return isSafeToSpeculativelyExecute(I, nullptr, nullptr, nullptr, nullptr,
-                                      /*UseVariableInfo=*/false);
+                                      /*UseVariableInfo=*/false,
+                                      IgnoreUBImplyingAttrs);
 }
 
 /// This returns the same result as isSafeToSpeculativelyExecute if Opcode is
@@ -586,7 +596,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 IgnoreUBImplyingAttrs = true);
 
 /// Returns true if the result or effects of the given instructions \p I
 /// depend values not reachable through the def use graph.

llvm/include/llvm/IR/Instruction.h

@@ -585,6 +585,10 @@ public:
   /// This should be used when speculating instructions.
   void dropUBImplyingAttrsAndMetadata();
 
+  /// Return true if this instruction has UB-implying attributes
+  /// that can cause immediate undefined behavior.
+  bool hasUBImplyingAttrs() const LLVM_READONLY;
+
   /// Determine whether the exact flag is set.
   bool isExact() const LLVM_READONLY;
 

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, /*IgnoreUBImplyingAttrs=*/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 IgnoreUBImplyingAttrs) {
   return isSafeToSpeculativelyExecuteWithOpcode(Inst->getOpcode(), Inst, CtxI,
-                                                AC, DT, TLI, UseVariableInfo);
+                                                AC, DT, TLI, UseVariableInfo,
+                                                IgnoreUBImplyingAttrs);
 }
 
 bool llvm::isSafeToSpeculativelyExecuteWithOpcode(
     unsigned Opcode, const Instruction *Inst, const Instruction *CtxI,
     AssumptionCache *AC, const DominatorTree *DT, const TargetLibraryInfo *TLI,
-    bool UseVariableInfo) {
+    bool UseVariableInfo, bool IgnoreUBImplyingAttrs) {
 #ifndef NDEBUG
   if (Inst->getOpcode() != Opcode) {
     // Check that the operands are actually compatible with the Opcode override.
@@ -7287,7 +7286,11 @@ 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.
+    return IgnoreUBImplyingAttrs || !CI->hasUBImplyingAttrs();
   }
   case Instruction::VAArg:
   case Instruction::Alloca:

llvm/lib/IR/Instruction.cpp

@@ -532,8 +532,8 @@ void Instruction::dropUBImplyingAttrsAndUnknownMetadata(
   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.
+  // that 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;
@@ -554,6 +554,20 @@ void Instruction::dropUBImplyingAttrsAndMetadata() {
   dropUBImplyingAttrsAndUnknownMetadata(KnownIDs);
 }
 
+bool Instruction::hasUBImplyingAttrs() const {
+  auto *CB = dyn_cast<CallBase>(this);
+  if (!CB)
+    return false;
+  // For call instructions, we also need to check parameter and return
+  // attributes that can cause UB.
+  for (unsigned ArgNo = 0; ArgNo < CB->arg_size(); ArgNo++)
+    if (CB->isPassingUndefUB(ArgNo))
+      return true;
+  return CB->hasRetAttr(Attribute::NoUndef) ||
+         CB->hasRetAttr(Attribute::Dereferenceable) ||
+         CB->hasRetAttr(Attribute::DereferenceableOrNull);
+}
+
 bool Instruction::isExact() const {
   return cast<PossiblyExactOperator>(this)->isExact();
 }

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
+}