[InstCombine] prevent miscompiles from select-of-div/rem transform

[spatel]

Summary

This avoids the danger shown in issue #60906. There were no regression tests for these patterns, so these potential failures have been around for a long time.

We freeze the condition and preserve the optimization because getting rid of a div/rem is always a big win.

Here are a couple of examples that can be corrected by freezing the condition:
https://alive2.llvm.org/ce/z/sXHTTC

[aqjune]

We could carve out some exceptions to this bailout. For example, the transform is actually ok with udiv/urem and a common divisor because poison in the dividend would not cause immediate UB - the divisor has to be zero, and that would be UB in the original code too. That doesn't work with signed div/rem because we can choose the poison dividend to be MinSignedValue, and that overflows when divided by -1 causing UB that doesn't exist in the original code.

This makes a lot of sense to me. Should we make this patch support this case as well?

[efriedma]

Alternatively, if you freeze the condition of the select, the transform is legal in all cases.

[spatel]

In D144671#4154013, @efriedma wrote:

Alternatively, if you freeze the condition of the select, the transform is legal in all cases.

Thanks! Yes, that's the likely better fix. This transform does get rid of a div/rem, so that's a potentially big win vs. adding a freeze.

[spatel]

Patch updated:
When necessary, freeze the select condition value to preserve the optimization.

[aqjune]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

440

Should we avoid introducing freeze if MatchIsOpZero is false and the op is udiv/urem?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

440

Sure, we can refine the check.

[spatel]

Patch updated:
Add restriction to freeze creation for unsigned ops because they are easier to determine as safe.

[aqjune]

LGTM, thanks..!

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

Instruction *InstCombinerImpl::foldSelectOpOp(SelectInst &SI, Instruction *TI,
 
   // If the select condition is a vector, the operands of the original select's
   // operands also must be vectors. This may not be the case for getelementptr
   // for example.
   if (CondTy->isVectorTy() && (!OtherOpT->getType()->isVectorTy() ||
                                !OtherOpF->getType()->isVectorTy()))
     return nullptr;
 
+  // If we are sinking div/rem after a select, freeze the condition because
+  // div/rem may induce immediate UB with a poison operand.
+  // For example, the following transform is not safe if Cond can ever be poison
+  // because we can replace poison with zero and then we have div-by-zero that
+  // didn't exist in the original code:
+  // Cond ? x/y : x/z --> x / (Cond ? y : z)
+  auto *BO = dyn_cast<BinaryOperator>(TI);
+  if (BO && BO->isIntDivRem() && !isGuaranteedNotToBePoison(Cond))
+    Cond = Builder.CreateFreeze(Cond);
+
   // If we reach here, they do have operations in common.
   Value *NewSI = Builder.CreateSelect(Cond, OtherOpT, OtherOpF,
                                       SI.getName() + ".v", &SI);
   Value *Op0 = MatchIsOpZero ? MatchOp : NewSI;
   Value *Op1 = MatchIsOpZero ? NewSI : MatchOp;
   if (auto *BO = dyn_cast<BinaryOperator>(TI)) {
     BinaryOperator *NewBO = BinaryOperator::Create(BO->getOpcode(), Op0, Op1);
     NewBO->copyIRFlags(TI);

llvm/test/Transforms/InstCombine/select-divrem.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
+; It is a miscompile in most of these tests if we
+; execute div/rem without freezing the potentially
+; poison condition value.
+
 define i5 @sdiv_common_divisor(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @sdiv_common_divisor(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = sdiv i5 [[SEL_V]], [[X:%.*]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = sdiv i5 %y, %x
   %r2 = sdiv i5 %z, %x
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i5 @srem_common_divisor(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @srem_common_divisor(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = srem i5 [[SEL_V]], [[X:%.*]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = srem i5 %y, %x
   %r2 = srem i5 %z, %x
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i5 @udiv_common_divisor(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @udiv_common_divisor(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = udiv i5 [[SEL_V]], [[X:%.*]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = udiv i5 %y, %x
   %r2 = udiv i5 %z, %x
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i5 @urem_common_divisor(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @urem_common_divisor(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = urem i5 [[SEL_V]], [[X:%.*]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = urem i5 %y, %x
   %r2 = urem i5 %z, %x
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i5 @sdiv_common_dividend(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @sdiv_common_dividend(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = sdiv i5 [[X:%.*]], [[SEL_V]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = sdiv i5 %x, %y
   %r2 = sdiv i5 %x, %z
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i5 @srem_common_dividend(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @srem_common_dividend(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = srem i5 [[X:%.*]], [[SEL_V]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = srem i5 %x, %y
   %r2 = srem i5 %x, %z
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i5 @udiv_common_dividend(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @udiv_common_dividend(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = udiv i5 [[X:%.*]], [[SEL_V]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = udiv i5 %x, %y
   %r2 = udiv i5 %x, %z
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
 define i5 @urem_common_dividend(i1 %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @urem_common_dividend(
-; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i1 [[B:%.*]]
+; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[TMP1]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = urem i5 [[X:%.*]], [[SEL_V]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = urem i5 %x, %y
   %r2 = urem i5 %x, %z
   %sel = select i1 %b, i5 %r2, i5 %r1
   ret i5 %sel
 }
 
+; Repeat the above tests, but guarantee that the select
+; condition is not poison via argument attribute. That
+; makes it safe to execute the select before div/rem
+; without needing to freeze the condition.
+
 define i5 @sdiv_common_divisor_defined_cond(i1 noundef %b, i5 %x, i5 %y, i5 %z) {
 ; CHECK-LABEL: @sdiv_common_divisor_defined_cond(
 ; CHECK-NEXT:    [[SEL_V:%.*]] = select i1 [[B:%.*]], i5 [[Z:%.*]], i5 [[Y:%.*]]
 ; CHECK-NEXT:    [[SEL:%.*]] = sdiv i5 [[SEL_V]], [[X:%.*]]
 ; CHECK-NEXT:    ret i5 [[SEL]]
 ;
   %r1 = sdiv i5 %y, %x
   %r2 = sdiv i5 %z, %x