[InstCombine] Optimize select(freeze(icmp eq/ne x, y), x, y)

[aqjune]

Summary

This patch adds an optimization that folds select(freeze(icmp eq/ne x, y), x, y)
to x or y.
This was needed to resolve slowdown after D84940 is applied.

I tried to bake this logic into foldSelectInstWithICmp, but it wasn't clear.
This patch conservatively writes the pattern in a separate function,
foldSelectWithFrozenICmp.

The output does not need freeze; https://alive2.llvm.org/ce/z/X49hNE (from @nikic)

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2543

Why is it necessary to freeze the result? Alive accepts this without the freeze as well: https://alive2.llvm.org/ce/z/X49hNE

[aqjune]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2543

You are correct - freeze isn't needed.
I'll remove it.

[aqjune]

freeze isn't needed in an output

[aqjune]

Added a missing condition which is that freeze should be used by the select instruction only.

[efriedma]

llvm/test/Transforms/InstCombine/select.ll

2588

If I'm understanding correctly, the following transform would be legal:

define void @select_freeze_icmp_multuses(i32 %x, i32 %y) {
  %x.fr = freeze i32 %x
  %y.fr = freeze i32 %y
  %c = icmp ne i32 %x.fr, %y.fr
  call void @use_i1_i32(i1 %c, i32 %x.fr)
  ret void
}

Is this something we should be looking into?

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2540

I agree the one-use check is necessary, but weirdly alive doesn't seem to require it: https://alive2.llvm.org/ce/z/hzwzM_

[spatel]

Is there a reason to do this in InstCombine rather than InstSimplify? We are not creating any new instructions.

[aqjune]

In D85533#2203926, @spatel wrote:

Is there a reason to do this in InstCombine rather than InstSimplify? We are not creating any new instructions.

Since this transformation takes the number of uses of freeze into consideration, InstCombine is a better place to handle it IMO.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2540

Yes, actually there is a bug regarding this pattern: https://github.com/AliveToolkit/alive2/issues/450
I'll have a look at this.

llvm/test/Transforms/InstCombine/select.ll

2588

Yes it is legal, but it did not appear on my codebase.
If it is beneficial to place comparison just before specific operations, we can add it in SelDag.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (auto *I = dyn_cast<Instruction>(V))
       CandidateBlocks.insert(I->getParent());
 
   for (BasicBlock *BB : CandidateBlocks)
     if (auto *PN = foldSelectToPhiImpl(Sel, BB, DT, Builder))
       return PN;
   return nullptr;
 }
 
+static Value *foldSelectWithFrozenICmp(SelectInst &Sel, InstCombiner::BuilderTy &Builder) {
+  FreezeInst *FI = dyn_cast<FreezeInst>(Sel.getCondition());
+  if (!FI)
+    return nullptr;
+
+  Value *Cond = FI->getOperand(0);
+  Value *TrueVal = Sel.getTrueValue(), *FalseVal = Sel.getFalseValue();
+
+  //   select (freeze(x == y)), x, y --> y
+  //   select (freeze(x != y)), x, y --> x
+  // The freeze should be only used by this select. Otherwise, remaining uses of
+  // the freeze can observe a contradictory value.
+  //   c = freeze(x == y)   ; Let's assume that y = poison & x = 42; c is 0 or 1
+  //   a = select c, x, y   ;
+  //   f(a, c)              ; f(poison, 1) cannot happen, but if a is folded
+  //                        ; to y, this can happen.
+  CmpInst::Predicate Pred;
+  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;
+}
+
 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 (Instruction *Rot = foldSelectRotate(SI, Builder))
     return Rot;
 
   if (Instruction *Copysign = foldSelectToCopysign(SI, Builder))
     return Copysign;
 
   if (Instruction *PN = foldSelectToPhi(SI, DT, Builder))
     return replaceInstUsesWith(SI, PN);
 
+  if (Value *Fr = foldSelectWithFrozenICmp(SI, Builder))
+    return replaceInstUsesWith(SI, Fr);
+
   return nullptr;
 }

llvm/test/Transforms/InstCombine/select.ll

;
   %s = select i1 %cond.fr, i8 %x, i8 %y
   %s2 = select i1 %cond.fr, i8 %y, i8 %x
   call void @foo2(i8 %s, i8 %s2)
   ret void
 }
 
 define i32 @select_freeze_icmp_eq(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_freeze_icmp_eq(
-; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[C_FR:%.*]] = freeze i1 [[C]]
-; CHECK-NEXT:    [[V:%.*]] = select i1 [[C_FR]], i32 [[X]], i32 [[Y]]
-; CHECK-NEXT:    ret i32 [[V]]
+; CHECK-NEXT:    ret i32 [[Y:%.*]]
 ;
   %c = icmp eq i32 %x, %y
   %c.fr = freeze i1 %c
   %v = select i1 %c.fr, i32 %x, i32 %y
   ret i32 %v
 }
 
 define i32 @select_freeze_icmp_ne(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_freeze_icmp_ne(
-; CHECK-NEXT:    [[C:%.*]] = icmp ne i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[C_FR:%.*]] = freeze i1 [[C]]
-; CHECK-NEXT:    [[V:%.*]] = select i1 [[C_FR]], i32 [[X]], i32 [[Y]]
-; CHECK-NEXT:    ret i32 [[V]]
+; CHECK-NEXT:    ret i32 [[X:%.*]]
 ;
   %c = icmp ne i32 %x, %y
   %c.fr = freeze i1 %c
   %v = select i1 %c.fr, i32 %x, i32 %y
   ret i32 %v
 }
 
 define i32 @select_freeze_icmp_else(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_freeze_icmp_else(
 ; CHECK-NEXT:    [[C:%.*]] = icmp ult i32 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[C_FR:%.*]] = freeze i1 [[C]]
 ; CHECK-NEXT:    [[V:%.*]] = select i1 [[C_FR]], i32 [[X]], i32 [[Y]]
 ; CHECK-NEXT:    ret i32 [[V]]
 ;
   %c = icmp ult i32 %x, %y
   %c.fr = freeze i1 %c
   %v = select i1 %c.fr, i32 %x, i32 %y
   ret i32 %v
 }
+
+declare void @use_i1_i32(i1, i32)
+
+define void @select_freeze_icmp_multuses(i32 %x, i32 %y) {
+; CHECK-LABEL: @select_freeze_icmp_multuses(
+; CHECK-NEXT:    [[C:%.*]] = icmp ne i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[C_FR:%.*]] = freeze i1 [[C]]
+; CHECK-NEXT:    [[V:%.*]] = select i1 [[C_FR]], i32 [[X]], i32 [[Y]]
+; CHECK-NEXT:    call void @use_i1_i32(i1 [[C_FR]], i32 [[V]])
+; CHECK-NEXT:    ret void
+;
+  %c = icmp ne i32 %x, %y
+  %c.fr = freeze i1 %c
+  %v = select i1 %c.fr, i32 %x, i32 %y
+  call void @use_i1_i32(i1 %c.fr, i32 %v)
+  ret void
+}