[ValueTracking] Improve impliesPoison

[aqjune]

Summary

This patch improves ValueTracking's impliesPoison(V1, V2) to do this reasoning:

  %res = call { i64, i1 } @llvm.umul.with.overflow.i64(i64 %a, i64 %b)
  %overflow = extractvalue { i64, i1 } %res, 1
  %mul      = extractvalue { i64, i1 } %res, 0

	; If %mul is poison, %overflow is also poison, and vice versa.

This improvement leads to supporting this optimization under -instcombine-unsafe-select-transform=0:

define i1 @test2_logical(i64 %a, i64 %b, i64* %ptr) {
; CHECK-LABEL: @test2_logical(
; CHECK-NEXT:    [[MUL:%.*]] = mul i64 [[A:%.*]], [[B:%.*]]
; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i64 [[A]], 0
; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i64 [[B]], 0
; CHECK-NEXT:    [[OVERFLOW_1:%.*]] = and i1 [[TMP1]], [[TMP2]]
; CHECK-NEXT:    [[NEG:%.*]] = sub i64 0, [[MUL]]
; CHECK-NEXT:    store i64 [[NEG]], i64* [[PTR:%.*]], align 8
; CHECK-NEXT:    ret i1 [[OVERFLOW_1]]
;

  %res = tail call { i64, i1 } @llvm.umul.with.overflow.i64(i64 %a, i64 %b)
  %overflow = extractvalue { i64, i1 } %res, 1
  %mul = extractvalue { i64, i1 } %res, 0
  %cmp = icmp ne i64 %mul, 0
  %overflow.1 = select i1 %overflow, i1 true, i1 %cmp
  %neg = sub i64 0, %mul
  store i64 %neg, i64* %ptr, align 8
  ret i1 %overflow.1
}

Previously, this didn't happen because the flag prevented select i1 %overflow, i1 true, i1 %cmp from being or i1 %overflow, %cmp.
Note that the select -> or conversion happens only when impliesPoison(%cmp, %overflow) returns true.
This improvement allows impliesPoison to do the reasoning.

[aqjune]

llvm/include/llvm/Analysis/ValueTracking.h

587

I chose to not make propagatesPoison return true on extractvalue because otherwise we lose this nice property:
"If propagatesPoison(op x y) is true, and the result is non-poison, x and y are also non-poison"
If extractvalue is added, this doesn't hold any more; imagine

%aggr = { i64 poison, i64 0 }
%I = extractvalue %aggr, 1
; propagatesPoison(extractvalue %aggr, 1) now returns true, and the result is well-defined (constant 0), but %aggr isn't well-defined

[spatel]

llvm/lib/Analysis/ValueTracking.cpp

4883–4885

Are there other patterns that you are planning to add here?
If not, we could make the patch much smaller with a more direct match like we use in InstCombineSelect.cpp:

WithOverflowInst *II;
if (!match(CondVal, m_ExtractValue<1>(m_WithOverflowInst(II))) ||
    !match(FalseVal, m_ExtractValue<0>(m_Specific(II))))
  return nullptr;

[aqjune]

Use m_WithOverflowInst

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

4883–4885

Yep, I think dealing with overflow instructions is enough.
I updated the code to use matches (with updates in PatternMatch.h)

[spatel]

Thanks - the pattern match with special-case -1 is a bit awkward, but I don't have a better suggestion.
I am not sure about subtleties of directlyImpliesPoison vs. impliesPoison. @nikic may want to have a look.

[nikic]

LGTM

llvm/include/llvm/Analysis/ValueTracking.h

constexpr unsigned MaxAnalysisRecursionDepth = 6;
   /// Note that this currently only considers the loop header.
   bool isGuaranteedToExecuteForEveryIteration(const Instruction *I,
                                               const Loop *L);
 
   /// Return true if I yields poison or raises UB if any of its operands is
   /// poison.
   /// Formally, given I = `r = op v1 v2 .. vN`, propagatesPoison returns true
   /// if, for all i, r is evaluated to poison or op raises UB if vi = poison.
+  /// If vi is a vector or an aggregate and r is a single value, any poison
+  /// element in vi should make r poison or raise UB.
   /// To filter out operands that raise UB on poison, you can use
   /// getGuaranteedNonPoisonOp.
   bool propagatesPoison(const Operator *I);
 
   /// Insert operands of I into Ops such that I will trigger undefined behavior
   /// if I is executed and that operand has a poison value.
   void getGuaranteedNonPoisonOps(const Instruction *I,
                                  SmallPtrSetImpl<const Value *> &Ops);
   /// Insert operands of I into Ops such that I will trigger undefined behavior
   /// if I is executed and that operand is not a well-defined value
   /// (i.e. has undef bits or poison).

llvm/include/llvm/IR/PatternMatch.h

 /// Match an instruction, capturing it if we match.
 inline bind_ty<Instruction> m_Instruction(Instruction *&I) { return I; }
 /// Match a unary operator, capturing it if we match.
 inline bind_ty<UnaryOperator> m_UnOp(UnaryOperator *&I) { return I; }
 /// Match a binary operator, capturing it if we match.
 inline bind_ty<BinaryOperator> m_BinOp(BinaryOperator *&I) { return I; }
 /// Match a with overflow intrinsic, capturing it if we match.
 inline bind_ty<WithOverflowInst> m_WithOverflowInst(WithOverflowInst *&I) { return I; }
+inline bind_ty<const WithOverflowInst>
+m_WithOverflowInst(const WithOverflowInst *&I) {
+  return I;
+}
 
 /// Match a Constant, capturing the value if we match.
 inline bind_ty<Constant> m_Constant(Constant *&C) { return C; }
 
 /// Match a ConstantInt, capturing the value if we match.
 inline bind_ty<ConstantInt> m_ConstantInt(ConstantInt *&CI) { return CI; }
 
 /// Match a ConstantFP, capturing the value if we match.
template <typename Val_t> inline Signum_match<Val_t> m_Signum(const Val_t &V) {
   return Signum_match<Val_t>(V);
 }
 
 template <int Ind, typename Opnd_t> struct ExtractValue_match {
   Opnd_t Val;
   ExtractValue_match(const Opnd_t &V) : Val(V) {}
 
   template <typename OpTy> bool match(OpTy *V) {
-    if (auto *I = dyn_cast<ExtractValueInst>(V))
-      return I->getNumIndices() == 1 && I->getIndices()[0] == Ind &&
-             Val.match(I->getAggregateOperand());
+    if (auto *I = dyn_cast<ExtractValueInst>(V)) {
+      // If Ind is -1, don't inspect indices
+      if (Ind != -1 &&
+          !(I->getNumIndices() == 1 && I->getIndices()[0] == (unsigned)Ind))
+        return false;
+      return Val.match(I->getAggregateOperand());
+    }
     return false;
   }
 };
 
 /// Match a single index ExtractValue instruction.
 /// For example m_ExtractValue<1>(...)
 template <int Ind, typename Val_t>
 inline ExtractValue_match<Ind, Val_t> m_ExtractValue(const Val_t &V) {
   return ExtractValue_match<Ind, Val_t>(V);
 }
 
+/// Match an ExtractValue instruction with any index.
+/// For example m_ExtractValue(...)
+template <typename Val_t>
+inline ExtractValue_match<-1, Val_t> m_ExtractValue(const Val_t &V) {
+  return ExtractValue_match<-1, Val_t>(V);
+}
+
 /// Matcher for a single index InsertValue instruction.
 template <int Ind, typename T0, typename T1> struct InsertValue_match {
   T0 Op0;
   T1 Op1;
 
   InsertValue_match(const T0 &Op0, const T1 &Op1) : Op0(Op0), Op1(Op1) {}
 
   template <typename OpTy> bool match(OpTy *V) {

llvm/lib/Analysis/ValueTracking.cpp

static bool directlyImpliesPoison(const Value *ValAssumedPoison,
                                   const Value *V, unsigned Depth) {
   if (ValAssumedPoison == V)
     return true;
 
   const unsigned MaxDepth = 2;
   if (Depth >= MaxDepth)
     return false;
 
-  const auto *I = dyn_cast<Instruction>(V);
-  if (I && propagatesPoison(cast<Operator>(I))) {
+  if (const auto *I = dyn_cast<Instruction>(V)) {
+    if (propagatesPoison(cast<Operator>(I)))
       return any_of(I->operands(), [=](const Value *Op) {
         return directlyImpliesPoison(ValAssumedPoison, Op, Depth + 1);
       });
+
+    // V  = extractvalue V0, idx
+    // V2 = extractvalue V0, idx2
+    // V0's elements are all poison or not. (e.g., add_with_overflow)
+    const WithOverflowInst *II;
+    if (match(I, m_ExtractValue(m_WithOverflowInst(II))) &&
+        match(ValAssumedPoison, m_ExtractValue(m_Specific(II))))
+      return true;
   }
   return false;
 }
 
 static bool impliesPoison(const Value *ValAssumedPoison, const Value *V,
                           unsigned Depth) {
   if (isGuaranteedNotToBeUndefOrPoison(ValAssumedPoison))
     return true;

llvm/test/Transforms/InstCombine/umul-sign-check.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -instcombine -S %s | FileCheck %s
+; RUN: opt -instcombine -instcombine-unsafe-select-transform=0 -S %s | FileCheck %s
 
 ; Check that we simplify llvm.umul.with.overflow, if the overflow check is
 ; weakened by or (icmp ne %res, 0) %overflow. This is generated by code using
 ; __builtin_mul_overflow with negative integer constants, e.g.
 
 ;   bool test(unsigned long long v, unsigned long long *res) {
 ;     return __builtin_mul_overflow(v, -4775807LL, res);
 ;   }