[InstCombine] Fold overflow bit of [u|s]mul.with.overflow in a poison-safe way

[aqjune]

Summary

As discussed in D101191, this patch adds a poison-safe folding of overflow bit check:

  %Op0 = icmp ne i4 %X, 0
  %Agg = call { i4, i1 } @llvm.[us]mul.with.overflow.i4(i4 %X, i4 %Y)
  %Op1 = extractvalue { i4, i1 } %Agg, 1
  %ret = select i1 %Op0, i1 %Op1, i1 false
=>
  %Y.fr = freeze %Y
  %Agg = call { i4, i1 } @llvm.[us]mul.with.overflow.i4(i4 %X, i4 %Y.fr)
  %Op1 = extractvalue { i4, i1 } %Agg, 1
  %ret = %Op1

https://alive2.llvm.org/ce/z/zgPUGT
https://alive2.llvm.org/ce/z/h2gZ_6

Note that there are cases where inserting freeze is not necessary: e.g. %Y is noundef.
In this case, LLVM is already good because %ret is already successfully folded into and,
triggering the pre-existing optimization in InstSimplify: https://godbolt.org/z/v6qena15K

[aqjune]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2587

The function below is based on InstSimplify's version.

llvm/test/Transforms/PhaseOrdering/unsigned-multiply-overflow-check.ll

8

-instcombine-unsafe-select-transform=0 is added to show the impact of this patch.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2587

Duplicating so much code feels wrong. Could we just call the simplify code from here with something like:

// Simulate a bitwise logic op for a select with a constant operand. 
if (match(FalseVal, m_Zero())) {
  if (Value *Simplified =
          SimplifyAndInst(CondVal, TrueVal, SQ.getWithInstruction(&SI))) {
    // create freeze and return
  }
}
if (match(TrueVal, m_One())) {
  if (Value *Simplified =
          SimplifyOrInst(CondVal, FalseVal, SQ.getWithInstruction(&SI))) {
    // create freeze and return
  }
}

llvm/test/Transforms/PhaseOrdering/unsigned-multiply-overflow-check.ll

8

I would prefer to leave this file unchanged with this patch, so we still have an indicator to represent the overall behavior of the optimizer.

We should add minimal tests for this patch under test/Transforms/InstCombine. These could be based on the test files from D65150 / D65151 - just replace the and/or logic instructions with select?

[aqjune]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2587

I think we still need to look into the expression because we should determine which variable to freeze.

If there is a header file that InstSimplify and InstCombine can share would be a great place for this function?

llvm/test/Transforms/PhaseOrdering/unsigned-multiply-overflow-check.ll

8

No problem, I'll make tests like those in the next iteration.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2587

Don't we only need to freeze the operand that is not in the icmp expression? That is, CondVal of the select must match as an icmp with a non-zero operand, and it's the other value that must be frozen.

I don't know of a shared header in Analysis for general logic. We do have things like VectorUtils and CmpInstAnalysis, so we could invent some new place.

[lebedev.ri]

I have a really 'are you nuts' proposal.
How about we canonicalize all i1 and/or's into select form?
Then that instsimplify code becomes dead, and we can simply port it here.
And we'll weed out all the remaining patterns that still only handle the and/or forms/

[aqjune]

In D101423#2723737, @lebedev.ri wrote:

I have a really 'are you nuts' proposal.
How about we canonicalize all i1 and/or's into select form?
Then that instsimplify code becomes dead, and we can simply port it here.
And we'll weed out all the remaining patterns that still only handle the and/or forms/

This will reduce a lot of complexity between having both select and or/and i1, but one concern is that select isn't commutative. :(
I mean, and i1 %a, %b == and i1 %b, %a, but in case of select select i1 %a, %b, false != select i1 %b, %a, false due to poison.
So we'll have to carefully choose which form to use, which can be a hard problem.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2587

Don't we only need to freeze the operand that is not in the icmp expression? That is, CondVal of the select must match as an icmp with a non-zero operand, and it's the other value that must be frozen.

Yep, the variable is the one to freeze.
My concern was about code duplication: the pattern matching code should still be copied here to pick the variable...

I don't know of a shared header in Analysis for general logic. We do have things like VectorUtils and CmpInstAnalysis, so we could invent some new place.

Actually to me having a separate header like CmpInstAnalysis and VectorUtil looks great.
Would creating something like OverflowInstAnalysis make sense?
If it is too much, having InstSimplifyUtils.h in llvm/lib/Analysis & including it from InstCombineSelect.cpp?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2587

An Overflow analysis file sounds good. At this point, I think there's only the mul code that would go in there, but we might find other common code later.

[aqjune]

• Move the function into the new header file
• Update ValueTracking's impliesPoison to convert select to and/or if valid.

This avoids redundant freeze insertion (see div-by-0-guard-before-smul_ov-not.ll 's t1_commutative)

[aqjune]

clang-format

[aqjune]

llvm/lib/Analysis/ValueTracking.cpp

5023 ↗

(On Diff #341761)

This update adds an inferences that if a is poison, extract(mul_with_overflow(a, _), _) is also poison.

This enables transformation

select(extract(mul_with_overflow(a, _), _), (a == 0), false)
=>
and(extract(mul_with_overflow(a, _), _), (a == 0))
```.

[spatel]

LGTM - see inline for some small improvements.

llvm/include/llvm/Analysis/OverflowInstAnalysis.h

33 ↗

(On Diff #341761)

This is confusing if we do not have the context of the callers that we can see in this patch.
The comment should say something like: "Match one of the patterns up to the select/logic op. Callers are expected to align that with the operands of the select/logic op user."

35 ↗

(On Diff #341761)

"is used" -> "are used"

36 ↗

(On Diff #341761)

"is one" -> "match one"

llvm/lib/Analysis/ValueTracking.cpp

5024 ↗

(On Diff #341761)

Use "is_contained" ?
Is it possible to separate this change (with an independent test)?

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

 //===- InstCombineSelect.cpp ----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
static Value *foldSelectWithFrozenICmp(SelectInst &Sel, InstCombiner::BuilderTy &Builder) {
   if (FI->hasOneUse() &&
       match(Cond, m_c_ICmp(Pred, m_Specific(TrueVal), m_Specific(FalseVal))) &&
       (Pred == ICmpInst::ICMP_EQ || Pred == ICmpInst::ICMP_NE)) {
     return Pred == ICmpInst::ICMP_EQ ? FalseVal : TrueVal;
   }
 
   return nullptr;
 }
 
+/// The @llvm.[us]mul.with.overflow intrinsic could have been folded from some
+/// other form of check, e.g. one that was using division; it may have been
+/// guarded against division-by-zero. We can drop that check now.
+///
+/// Look for:
+///   %Op0 = icmp ne i4 %X, 0
+///   %Agg = call { i4, i1 } @llvm.[us]mul.with.overflow.i4(i4 %X, i4 %Y)
+///   %Op1 = extractvalue { i4, i1 } %Agg, 1
+///   %ret = select i1 %Op0, i1 %Op1, i1 false
+///
+///   %Op0 = icmp eq i4 %X, 0
+///   %Agg = call { i4, i1 } @llvm.[us]mul.with.overflow.i4(i4 %X, i4 %Y)
+///   %NotOp1 = extractvalue { i4, i1 } %Agg, 1
+///   %Op1 = xor i1 %NotOp1, true
+///   %ret = select i1 %Op0, i1 true, i1 %Op1
+///
+/// We can freeze %Y and just return %Op1 in both cases.
+
+static Value *
+omitCheckForZeroBeforeMulWithOverflow(Value *Op0, Value *Op1, bool IsAnd,
+                                      InstCombinerImpl &IC) {
+  ICmpInst::Predicate Pred;
+  Value *X, *NotOp1;
+  int XIdx;
+  IntrinsicInst *II;
+
+  if (!match(Op0, m_ICmp(Pred, m_Value(X), m_Zero())))
+    return nullptr;
+
+  ///   %Agg = call { i4, i1 } @llvm.[us]mul.with.overflow.i4(i4 %X, i4 %???)
+  ///   %V = extractvalue { i4, i1 } %Agg, 1
+  auto matchMulOverflowCheck = [X, &II, &XIdx](Value *V) {
+    auto *Extract = dyn_cast<ExtractValueInst>(V);
+    // We should only be extracting the overflow bit.
+    if (!Extract || !Extract->getIndices().equals(1))
+      return false;
+
+    II = dyn_cast<IntrinsicInst>(Extract->getAggregateOperand());
+    if (!match(II, m_CombineOr(m_Intrinsic<Intrinsic::umul_with_overflow>(),
+                               m_Intrinsic<Intrinsic::smul_with_overflow>())))
+      return false;
+
+    if (II->getArgOperand(0) == X)
+      XIdx = 0;
+    else if (II->getArgOperand(1) == X)
+      XIdx = 1;
+    else
+      return false;
+    return true;
+  };
+
+  bool Matched = (IsAnd && Pred == ICmpInst::Predicate::ICMP_NE &&
+        matchMulOverflowCheck(Op1)) ||
+      (!IsAnd && Pred == ICmpInst::Predicate::ICMP_EQ &&
+        match(Op1, m_Not(m_Value(NotOp1))) &&
+        matchMulOverflowCheck(NotOp1));
+
+  if (!Matched)
+    return nullptr;
+
+  int IdxToFreeze = !XIdx;
+  Value *ValueToFr = II->getArgOperand(IdxToFreeze);
+  FreezeInst *FI = new FreezeInst(ValueToFr, ValueToFr->getName() + ".fr");
+  IC.InsertNewInstBefore(FI, *II);
+  IC.replaceUse(II->getArgOperandUse(IdxToFreeze), FI);
+  return Op1;
+}
+
 Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();
 
   // FIXME: Remove this workaround when freeze related patches are done.
   // For select with undef operand which feeds into an equality comparison,
if (match(CondVal, m_Select(m_Value(A), m_Value(B), m_Zero())) &&
       return replaceOperand(SI, 0, A);
 
     if (Value *S = SimplifyWithOpReplaced(TrueVal, CondVal, One, SQ,
                                           /* AllowRefinement */ true))
       return replaceOperand(SI, 1, S);
     if (Value *S = SimplifyWithOpReplaced(FalseVal, CondVal, Zero, SQ,
                                           /* AllowRefinement */ true))
       return replaceOperand(SI, 2, S);
+
+    if (match(FalseVal, m_Zero())) {
+      if (Value *V = omitCheckForZeroBeforeMulWithOverflow(CondVal, TrueVal, true, *this))
+        return replaceInstUsesWith(SI, V);
+    } else if (match(TrueVal, m_One())) {
+      if (Value *V = omitCheckForZeroBeforeMulWithOverflow(CondVal, FalseVal, false, *this))
+        return replaceInstUsesWith(SI, V);
+    }
   }
 
   // Selecting between two integer or vector splat integer constants?
   //
   // Note that we don't handle a scalar select of vectors:
   // select i1 %c, <2 x i8> <1, 1>, <2 x i8> <0, 0>
   // because that may need 3 instructions to splat the condition value:
   // extend, insertelement, shufflevector.

llvm/test/Transforms/PhaseOrdering/unsigned-multiply-overflow-check.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s --check-prefix=SIMPLIFYCFG
 ; RUN: opt -instcombine -S < %s | FileCheck %s --check-prefix=INSTCOMBINEONLY
 ; RUN: opt -instcombine -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s --check-prefix=INSTCOMBINESIMPLIFYCFGONLY
 ; RUN: opt -instcombine -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -instcombine -S < %s | FileCheck %s --check-prefix=INSTCOMBINESIMPLIFYCFGINSTCOMBINE
 ; RUN: opt -instcombine -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -phi-node-folding-threshold=3 -S < %s | FileCheck %s --check-prefix=INSTCOMBINESIMPLIFYCFGONLY
 ; RUN: opt -instcombine -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -instcombine -phi-node-folding-threshold=3 -S < %s | FileCheck %s --check-prefix=INSTCOMBINESIMPLIFYCFGINSTCOMBINE
 ; RUN: opt -instcombine -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -instcombine -instcombine-unsafe-select-transform=0 -phi-node-folding-threshold=3 -S < %s | FileCheck %s --check-prefix=INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-pc-linux-gnu"
 
 ; #include <limits>
 ; #include <cstdint>
 ;
 ; using size_type = std::size_t;

 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-LABEL: @will_not_overflow(
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:  bb:
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:    [[UMUL:%.*]] = call { i64, i1 } @llvm.umul.with.overflow.i64(i64 [[ARG:%.*]], i64 [[ARG1:%.*]])
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:    [[UMUL_OV:%.*]] = extractvalue { i64, i1 } [[UMUL]], 1
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:    ret i1 [[UMUL_OV]]
 ;
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-LABEL: @will_not_overflow(
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:  bb:
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[T0:%.*]] = icmp ne i64 [[ARG:%.*]], 0
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[UMUL:%.*]] = call { i64, i1 } @llvm.umul.with.overflow.i64(i64 [[ARG]], i64 [[ARG1:%.*]])
+; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[ARG1_FR:%.*]] = freeze i64 [[ARG1:%.*]]
+; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[UMUL:%.*]] = call { i64, i1 } @llvm.umul.with.overflow.i64(i64 [[ARG:%.*]], i64 [[ARG1_FR]])
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[UMUL_OV:%.*]] = extractvalue { i64, i1 } [[UMUL]], 1
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[T6:%.*]] = select i1 [[T0]], i1 [[UMUL_OV]], i1 false
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    ret i1 [[T6]]
+; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    ret i1 [[UMUL_OV]]
 ;
 bb:
   %t0 = icmp eq i64 %arg, 0
   br i1 %t0, label %bb5, label %bb2
 
 bb2:                                              ; preds = %bb
   %t3 = udiv i64 -1, %arg
   %t4 = icmp ult i64 %t3, %arg1

 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:  bb:
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:    [[UMUL:%.*]] = call { i64, i1 } @llvm.umul.with.overflow.i64(i64 [[ARG:%.*]], i64 [[ARG1:%.*]])
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:    [[UMUL_OV:%.*]] = extractvalue { i64, i1 } [[UMUL]], 1
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:    [[PHI_BO:%.*]] = xor i1 [[UMUL_OV]], true
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINE-NEXT:    ret i1 [[PHI_BO]]
 ;
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-LABEL: @will_overflow(
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:  bb:
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[T0:%.*]] = icmp eq i64 [[ARG:%.*]], 0
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[UMUL:%.*]] = call { i64, i1 } @llvm.umul.with.overflow.i64(i64 [[ARG]], i64 [[ARG1:%.*]])
+; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[ARG1_FR:%.*]] = freeze i64 [[ARG1:%.*]]
+; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[UMUL:%.*]] = call { i64, i1 } @llvm.umul.with.overflow.i64(i64 [[ARG:%.*]], i64 [[ARG1_FR]])
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[UMUL_OV:%.*]] = extractvalue { i64, i1 } [[UMUL]], 1
 ; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[PHI_BO:%.*]] = xor i1 [[UMUL_OV]], true
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    [[T6:%.*]] = select i1 [[T0]], i1 true, i1 [[PHI_BO]]
-; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    ret i1 [[T6]]
+; INSTCOMBINESIMPLIFYCFGINSTCOMBINEUNSAFE-NEXT:    ret i1 [[PHI_BO]]
 ;
 bb:
   %t0 = icmp eq i64 %arg, 0
   br i1 %t0, label %bb5, label %bb2
 
 bb2:                                              ; preds = %bb
   %t3 = udiv i64 -1, %arg
   %t4 = icmp ult i64 %t3, %arg1
   br label %bb5
 
 bb5:                                              ; preds = %bb2, %bb
   %t6 = phi i1 [ false, %bb ], [ %t4, %bb2 ]
   %t7 = xor i1 %t6, true
   ret i1 %t7
 }