[InstCombine] Add optimization to prevent poison from being propagated.

[hyeongyukim]

Summary

In D104569, Freeze was inserted just before br to solve the branching on undef miscompilation problem.
But value analysis was being disturbed by added freeze.

v = load ptr
cond = freeze(icmp (and v, const), const')
br cond, ...

The case in which value analysis disturbed is as above.
By changing freeze to add immediately after load, value analysis will be successful again.

v = load ptr
freeze(icmp (and v, const), const')
=>
v = load ptr
v' = freeze v
icmp (and v', const), const'

In this patch, I propose the above optimization.
With this patch, the poison will not spread as the freeze is performed early.

[nikic]

This optimization is unnecessarily specific to the load + and + icmp combination. I would expect it to go something like this: If the freeze operand is a) one-use, b) does not produce poison, c) has all but one guaranteed-non-poison operands, then push the freeze through to the one operand that is not guaranteed-non-poison. Recursively, this will push the freeze first through the icmp and then the and, but also work for other instruction sequences.

[hyeongyukim]

In D105392#2856695, @nikic wrote:

This optimization is unnecessarily specific to the load + and + icmp combination. I would expect it to go something like this: If the freeze operand is a) one-use, b) does not produce poison, c) has all but one guaranteed-non-poison operands, then push the freeze through to the one operand that is not guaranteed-non-poison. Recursively, this will push the freeze first through the icmp and then the and, but also work for other instruction sequences.

That's a good point!
I will improve my patch to deal with more cases.

[hyeongyukim]

Is there any simple way to check if Operand doesn't produce Poison? @nikic

[nikic]

@hyeongyukim canCreateUndefOrPoison()

[hyeongyukim]

Modified to try to push freeze through instructions that propagate but don't produce poison as far as possible.

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3596

This is the wrong API to use here: You want isGuaranteedNotToBeUndefOrPoison() on the operands of the instruction.

getGuaranteedNonPoisonOps() is about which operands of an instruction would result in immediate UB if they are poison.

[hyeongyukim]

I modified it to use isGuaranteedNotToBeUndefOrPoison() instead of isGuaranteedNotToBeUndefOrPoison().

[hyeongyukim]

isGuaranteedNotToBePoison -> isGuaranteedNotToBeUndefOrPoison

[hyeongyukim]

Reformat Code.

[xbolva00]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3598

Address clang tidy warnings? Plus above.

[nikic]

llvm/lib/Analysis/ValueTracking.cpp

5347 ↗

(On Diff #356384)

This change should no longer be needed.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3598

This should be for (Use &U : Inst->operands()) or so.

3611

This is not generally safe. E.g. if Op1 is an invoke, you can't insert after it. I think for now you should insert this before Inst.

3613

What you're doing here is really a separate transform: You're canonicalizing from v to freeze v, if a freeze is present. This is presumably needed for your motivating case, but I think it would be better to handle it as an independent transform, as profitability is less obvious in this case. You also don't have any test coverage for this right now.

For now, I'd just do a replaceUse() for the single use.

[hyeongyukim]

Remove meaningless code, modify the location where freeze inserted and some refactoring.

[hyeongyukim]

Thank you for your detailed review!

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3580

This should go into a new function

3595

3598–3601

3604

These checks should be done first

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3609

You should directly store the Use * rather than the Value * and then pass that to replaceUse(). Right now you're replacing some random use of the value.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3595

And actually, what do we want to do if there are several operands with the same Value?
Unless that is okay for us (we'd still only need a single freeze) and we want to rewrite them all,
i guess we need to record the operand index, not the value.

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3595

This is why it should be Use * rather than Value *.

Just to be perfectly clear, this code should be Use *NotGuaranteedNonPoisonOperand = nullptr and then below NotGuaranteedNonPoisonOperand = &U and replaceUse(*NotGuaranteedNonPoisonOperand, FI).

Also maybe rename NotGuaranteedNonPoisonOperand to MaybePoisonOperand, it's a mouthful :)

[hyeongyukim]

Move code to seperate function.
Fixed the problems that Value* is used instead of Use*.

[hyeongyukim]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3580

I named it PushFreezeToPreventPoisonFromPropagate, is there a more appropriate name?
Please recommend me a better name!

3595

TBH I didn't consider the difference between Use * and Value *, but I realized it thanks to the comments.

[lebedev.ri]

Almost there i think

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3533–3534

Early return please
Also, this needs a comment why we don't want to just adjust all the other uses.

3535

Ok, Optional isn't much better here than just storing a pointer, let's go back to that one.

3541–3542

3545–3546

Can we just drop freeze then?

3549

Not clang-formatted

[hyeongyukim]

Optional was removed and comments were added.

[hyeongyukim]

Fix test code.

[hyeongyukim]

I corrected it wrong while refactoring it. I'll fix it again.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3516–3568

I think it would be good to have a comment explaining why we don't do this iff the OrigOpInst has other uses.
We could just change them all to use OrigFI first. Why do we want to not do this?

3517–3519

Note that it was a question, i don't know what should we be doing here.

3519

You probably want = nullptr;

3523

This looks like the freeze will get the name from it's user.
I would expect that it would get the base name from the non-frozen value.

[nikic]

Can you please add a test for freeze of phi? I expect this will fail because you can't insert an instruction before a phi. It's okay to simply bail on this case for now.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3517–3519

Normally the freeze would get dropped by the basic isGuaranteedNotToBeUndefOrPoison() fold already. I think we can only hit this case due to recursion bailout (we effectively allow one more level of recursion here).

As it doesn't cut us anything to do so, just dropping the freeze here seems reasonable.

[hyeongyukim]

Another fixes with some comments.

[hyeongyukim]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3516–3568

I add an example of a problem that occurs when there is more than one use of OrigOpInst.

[hyeongyukim]

In D105392#2857277, @nikic wrote:

Can you please add a test for freeze of phi? I expect this will fail because you can't insert an instruction before a phi. It's okay to simply bail on this case for now.

define i32 @test(i1 %cond, i32 *%ptrA) {
; CHECK-LABEL: @test(
; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[B:%.*]]
; CHECK:       A:
; CHECK-NEXT:    br label [[C:%.*]]
; CHECK:       B:
; CHECK-NEXT:    [[V:%.*]] = load i32, i32* [[PTR:%.*]], align 4
; CHECK-NEXT:    [[PHI_FR:%.*]] = freeze i32 [[V]]
; CHECK-NEXT:    br label [[C]]
; CHECK:       C:
; CHECK-NEXT:    [[Y:%.*]] = phi i32 [ 0, [[A]] ], [ [[PHI_FR]], [[B]] ]
; CHECK-NEXT:    ret i32 [[Y]]
;
  br i1 %cond, label %A, label %B
A:
  br label %C
B:
  %v = load i32, i32* %ptr
  br label %C
C:
  %y = phi i32 [0, %A], [%v, %B]
  %y.fr = freeze i32 %y
  ret i32 %y.fr
}

I tried it, and it seems to go to the label included in phi and add freeze.
@nikic Can you give me another example that could fail?

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3575

@hyeongyukim Looks like phi nodes are not a problem because they're handled separately here. I'd suggest to add a assert(!isa<PHINode>(...)) check in your transform to make it obvious that it can't occur there.

[hyeongyukim]

Add assertion about PHINode

[nikic]

It looks like you uploaded a partial diff in the last iteration.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3517

I'd move the assert next to the insertBefore call, as that's the only place where this matters.

3517

guranteed -> guaranteed

3525

This doesn't look like a wrong transformation to me. Just a possibly non-profitable one, as we unnecessarily use a frozen value. (I agree that we shouldn't do this, at least not in this form, but the comment does not seem right.)

llvm/test/Transforms/InstCombine/freeze.ll

88–137

It would be good to add a test for not pushing freeze through add nuw %x, 1. (In the future, we could push freeze through by dropping poison flags.)

[hyeongyukim]

Edit comments and add test case.

[hyeongyukim]

change to whole diff

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3517

As the clang-tidy warning indicates, the method should be lowercase. Also Propagate -> Propagating?

[hyeongyukim]

rename function

[hyeongyukim]

@lebedev.ri
Could you check this patch? Thanks

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3519

Hm, why would multiple freeze would be inserted?
The transform i envisioned is:

%a = <whatever> %x
%a.frozen = freeze %a
%b = <whatever> %a.frozen
%c = <whatever> %a
 =>
%a = <whatever> %x
%a.frozen = freeze %a
%b = <whatever> %a.frozen
%c = <whatever> %a.frozen
 =>
%x.frozen = freeze %x
%a.frozen = <whatever> %x.frozen
%b = <whatever> %a.frozen
%c = <whatever> %a.frozen

My question being, why do we not to change other users of %a to use %a.frozen?

3563–3564

I do not expect that this assertion will hold, foldOpIntoPhi() has a number of bailouts.

[hyeongyukim]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3519

In terms of correctness, it is still fine.
However, %aand %a.frozen have different semantics, and the optimization performed in %a may not be performed in %a.frozen.
I thought it might not be good for the frozen values to spread too much, so single-use %a is considered only.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3519

That is what i thought, too, but that's not what the comment says.

[hyeongyukim]

Clarify comments.

[hyeongyukim]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3563–3564

Do you mean OrigOp can be PHINode even if we already tried foldOpIntoPhi()?
Then, is it correct to change the assertion to the below code?

if (is<PHINode>(OrigOp))
  return nullptr;

[aqjune]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3538

Maybe update this sentence as well? Insertion of multiple freezes isn't necessary, as discussed.

3563–3564

I think so - +1 for foldOpIntoPhi handling the freeze(phi) case.

[hyeongyukim]

Modify to return null if OrigOp is PHINode.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3536

While we could change the other users of OrigOp to use freeze(OrigOp), that potentially reduces their optimization potential, so let's only do this iff the OrigOp is only used by the freeze.

3538

The PHI bailout should go here.
Otherwise we will be wasting compile time.

[hyeongyukim]

Update comment.

[lebedev.ri]

LG i guess.
Thanks.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3556–3558

Oh, hm, actually i suspect this should go into instsimplify, without the artificial PHINode restriction..

3580–3581

[hyeongyukim]

Modify pushFreezeToPreventPoisonFromPropagating to return Instruction rather than Value.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3580–3581

I actually forgot to remove this comment, it doesn't make sense.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3539

The one-use check may need refinement given freezeDominatedUses()

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   foldAggregateConstructionIntoAggregateReuse(InsertValueInst &OrigIVI);
   Instruction *visitInsertValueInst(InsertValueInst &IV);
   Instruction *visitInsertElementInst(InsertElementInst &IE);
   Instruction *visitExtractElementInst(ExtractElementInst &EI);
   Instruction *visitShuffleVectorInst(ShuffleVectorInst &SVI);
   Instruction *visitExtractValueInst(ExtractValueInst &EV);
   Instruction *visitLandingPadInst(LandingPadInst &LI);
   Instruction *visitVAEndInst(VAEndInst &I);
+  Value *pushFreezeToPreventPoisonFromPropagating(FreezeInst &FI);
   Instruction *visitFreeze(FreezeInst &I);
 
   /// Specify what to return for unhandled instructions.
   Instruction *visitInstruction(Instruction &I) { return nullptr; }
 
   /// True when DB dominates all uses of DI except UI.
   /// UI must be in the same block as DI.
   /// The routine checks that the DI parent and DB are different.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombinerImpl::visitLandingPadInst(LandingPadInst &LI) {
   if (LI.isCleanup() != CleanupFlag) {
     assert(!CleanupFlag && "Adding a cleanup, not removing one?!");
     LI.setCleanup(CleanupFlag);
     return &LI;
   }
 
   return nullptr;
 }
 
+Value *
+InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
+  // Try to push freeze through instructions that propagate but don't produce
+  // poison as far as possible.  If an operand of freeze follows three
+  // conditions 1) one-use, 2) does not produce poison, and 3) has all but one
+  // guaranteed-non-poison operands then push the freeze through to the one
+  // operand that is not guaranteed non-poison.  The actual transform is as
+  // follows.
+  //   Op1 = ...                        ; Op1 can be posion
+  //   Op0 = Inst(Op1, NonPoisonOps...) ; Op0 has only one use and only have
+  //                                    ; single guaranteed-non-poison operands
+  //   ... = Freeze(Op0)
+  // =>
+  //   Op1 = ...
+  //   Op1.fr = Freeze(Op1)
+  //   ... = Inst(Op1.fr, NonPoisonOps...)
+  auto *OrigOp = OrigFI.getOperand(0);
+  auto *OrigOpInst = dyn_cast<Instruction>(OrigOp);
+
+  // While we could change the other users of OrigOp to use freeze(OrigOp), that
+  // potentially reduces their optimization potential, so let's only do this iff
+  // the OrigOp is only used by the freeze.
+  if (!OrigOpInst || !OrigOpInst->hasOneUse() || isa<PHINode>(OrigOp) ||
+      canCreateUndefOrPoison(dyn_cast<Operator>(OrigOp)))
+    return nullptr;
+
+  // If operand is guaranteed not to be poison, there is no need to add freeze
+  // to the operand. So we first find the operand that is not guaranteed to be
+  // poison.
+  Use *MaybePoisonOperand = nullptr;
+  for (Use &U : OrigOpInst->operands()) {
+    if (isGuaranteedNotToBeUndefOrPoison(U.get()))
+      continue;
+    if (!MaybePoisonOperand)
+      MaybePoisonOperand = &U;
+    else
+      return nullptr;
+  }
+
+  // If all operands are guaranteed to be non-poison, we can drop freeze.
+  if (!MaybePoisonOperand)
+    return OrigOp;
+
+  auto *FrozenMaybePoisonOperand = new FreezeInst(
+      MaybePoisonOperand->get(), MaybePoisonOperand->get()->getName() + ".fr");
+
+  replaceUse(*MaybePoisonOperand, FrozenMaybePoisonOperand);
+  FrozenMaybePoisonOperand->insertBefore(OrigOpInst);
+  return OrigOp;
+}
+
 Instruction *InstCombinerImpl::visitFreeze(FreezeInst &I) {
   Value *Op0 = I.getOperand(0);
 
   if (Value *V = SimplifyFreezeInst(Op0, SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   // freeze (phi const, x) --> phi const, (freeze x)
   if (auto *PN = dyn_cast<PHINode>(Op0)) {
     if (Instruction *NV = foldOpIntoPhi(I, PN))
       return NV;
   }
 
+  if (Value *NI = pushFreezeToPreventPoisonFromPropagating(I))
+    return replaceInstUsesWith(I, NI);
+
   if (match(Op0, m_Undef())) {
     // If I is freeze(undef), see its uses and fold it to the best constant.
     // - or: pick -1
     // - select's condition: pick the value that leads to choosing a constant
     // - other ops: pick 0
     Constant *BestValue = nullptr;
     Constant *NullValue = Constant::getNullValue(I.getType());
     for (const auto *U : I.users()) {
       Constant *C = NullValue;
 
       if (match(U, m_Or(m_Value(), m_Value())))
         C = Constant::getAllOnesValue(I.getType());
       else if (const auto *SI = dyn_cast<SelectInst>(U)) {
         if (SI->getCondition() == &I) {
           APInt CondVal(1, isa<Constant>(SI->getFalseValue()) ? 0 : 1);
           C = Constant::getIntegerValue(I.getType(), CondVal);
         }
       }
 
       if (!BestValue)
         BestValue = C;
       else if (BestValue != C)
         BestValue = NullValue;
     }
 
     return replaceInstUsesWith(I, BestValue);
   }
 
   return nullptr;
 }
 
 /// Try to move the specified instruction from its current block into the
 /// beginning of DestBlock, which can only happen if it's safe to move the
 /// instruction past all of the instructions between it and the end of its
 /// block.
 static bool TryToSinkInstruction(Instruction *I, BasicBlock *DestBlock) {
   assert(I->getSingleUndroppableUse() && "Invariants didn't hold!");
   BasicBlock *SrcBlock = I->getParent();
 

llvm/test/Transforms/InstCombine/freeze.ll

 ; CHECK-NEXT:    call void @use_i32_i1(i32 32, i1 [[Y:%.*]])
 ; CHECK-NEXT:    ret void
 ;
   %f = freeze i1 undef
   %a = select i1 %f, i32 %x, i32 32 ; prefers %f to be false
   %b = select i1 %f, i1 true, i1 %y ; prefers %f to be true
   call void @use_i32_i1(i32 %a, i1 %b)
   ret void
 }
+
+; Move the freeze forward to prevent poison from spreading.
+
+define i32 @early_freeze_test1(i32 %x, i32 %y) {
+; CHECK-LABEL: @early_freeze_test1(
+; CHECK-NEXT:    [[V1:%.*]] = add i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
+; CHECK-NEXT:    [[V2:%.*]] = shl i32 [[V1_FR]], 1
+; CHECK-NEXT:    [[V3:%.*]] = and i32 [[V2]], 2
+; CHECK-NEXT:    ret i32 [[V3]]
+;
+  %v1 = add i32 %x, %y
+  %v2 = shl i32 %v1, 1
+  %v3 = and i32 %v2, 2
+  %v3.fr = freeze i32 %v3
+  ret i32 %v3.fr
+}
+
+define i1 @early_freeze_test2(i32* %ptr) {
+; CHECK-LABEL: @early_freeze_test2(
+; CHECK-NEXT:    [[V1:%.*]] = load i32, i32* [[PTR:%.*]], align 4
+; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
+; CHECK-NEXT:    [[V2:%.*]] = and i32 [[V1_FR]], 1
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[V2]], 0
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %v1 = load i32, i32* %ptr
+  %v2 = and i32 %v1, 1
+  %cond = icmp eq i32 %v2, 0
+  %cond.fr = freeze i1 %cond
+  ret i1 %cond.fr
+}
+
+; add can overflows, so we cannot move freeze beyond add
+
+define i32 @early_freeze_test3(i32 %v1) {
+; CHECK-LABEL: @early_freeze_test3(
+; CHECK-NEXT:    [[V2:%.*]] = shl i32 [[V1:%.*]], 1
+; CHECK-NEXT:    [[V3:%.*]] = add nuw i32 [[V2]], 2
+; CHECK-NEXT:    [[V3_FR:%.*]] = freeze i32 [[V3]]
+; CHECK-NEXT:    [[V4:%.*]] = or i32 [[V3_FR]], 1
+; CHECK-NEXT:    ret i32 [[V4]]
+;
+  %v2 = shl i32 %v1, 1
+  %v3 = add nuw i32 %v2, 2
+  %v4 = or i32 %v3, 1
+  %v4.fr = freeze i32 %v4
+  ret i32 %v4.fr
+}