[instcombine] propagate freeze through single use poison producing flag instruction

[reames]

Summary

If we have an instruction which produces poison only when flags are specified on the instruction, then we know that freezing the operands and dropping flags is equivalent to freezing the result. If we know those flags don't result in any undefined behavior being executed, then there's no point in preserving the flags as we gain no knowledge by having them.

This patch extends the existing propagation logic which sinks freeze to single potential non-poison operands to allow dropping of flags when we know the freeze is the sole use of the instruction with poison flags.

The main value is that we tend to sink freezes towards the phi in IV cycles where the incoming value to the phi is the freeze of an IV increment. This will in turn (in a future patch), let us fold the freeze through the phi into the loop preheader. Motivated by eliminating need for CanonicalizeFreezeInLoops for the clearly profitable cases from onephi.ll test case in the test directory.

[nikic]

LGTM apart from the requested test.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

3574

I'm missing a test for this case (i.e. freeze eliminated completely after dropping flags).

[nikic]

Ooops, something I missed and that wasn't covered by tests: dropPoisonGeneratingFlags() currently doesn't drop poison FMF, so there's a mismatch between the analysis and the dropping code.

[reames]

In D111675#3059737, @nikic wrote:

Ooops, something I missed and that wasn't covered by tests: dropPoisonGeneratingFlags() currently doesn't drop poison FMF, so there's a mismatch between the analysis and the dropping code.

Fixed in 4c5702c

llvm/include/llvm/Analysis/ValueTracking.h

constexpr unsigned MaxAnalysisRecursionDepth = 6;
   /// non-undef & non-poison operands.
   /// For vectors, canCreateUndefOrPoison returns true if there is potential
   /// poison or undef in any element of the result when vectors without
   /// undef/poison poison are given as operands.
   /// For example, given `Op = shl <2 x i32> %x, <0, 32>`, this function returns
   /// true. If Op raises immediate UB but never creates poison or undef
   /// (e.g. sdiv I, 0), canCreatePoison returns false.
   ///
+  /// \p ConsiderFlags controls whether poison producing flags on the
+  /// instruction are considered.  This can be used to see if the instruction
+  /// could still introduce undef or poison even without poison generating flags
+  /// which might be on the instruction.  (i.e. could the result of
+  /// Op->dropPoisonGeneratingFlags() still create poison or undef)
+  ///
   /// canCreatePoison returns true if Op can create poison from non-poison
   /// operands.
-  bool canCreateUndefOrPoison(const Operator *Op);
-  bool canCreatePoison(const Operator *Op);
+  bool canCreateUndefOrPoison(const Operator *Op, bool ConsiderFlags = true);
+  bool canCreatePoison(const Operator *Op, bool ConsiderFlags = true);
 
   /// Return true if V is poison given that ValAssumedPoison is already poison.
   /// For example, if ValAssumedPoison is `icmp X, 10` and V is `icmp X, 5`,
   /// impliesPoison returns true.
   bool impliesPoison(const Value *ValAssumedPoison, const Value *V);
 
   /// Return true if this function can prove that V does not have undef bits
   /// and is never poison. If V is an aggregate value or vector, check whether

llvm/lib/Analysis/ValueTracking.cpp

auto AllUsesGuardedByBranch = [&](const BranchInst *BI) {
     }
 
     return true;
   };
 
   return llvm::any_of(GuardingBranches, AllUsesGuardedByBranch);
 }
 
-static bool canCreateUndefOrPoison(const Operator *Op, bool PoisonOnly) {
+static bool canCreateUndefOrPoison(const Operator *Op, bool PoisonOnly,
+                                   bool ConsiderFlags) {
+  if (ConsiderFlags) {
     // See whether I has flags that may create poison
     if (const auto *OvOp = dyn_cast<OverflowingBinaryOperator>(Op)) {
       if (OvOp->hasNoSignedWrap() || OvOp->hasNoUnsignedWrap())
         return true;
     }
     if (const auto *ExactOp = dyn_cast<PossiblyExactOperator>(Op))
       if (ExactOp->isExact())
         return true;
     if (const auto *FP = dyn_cast<FPMathOperator>(Op)) {
       auto FMF = FP->getFastMathFlags();
       if (FMF.noNaNs() || FMF.noInfs())
         return true;
     }
+  }
 
   unsigned Opcode = Op->getOpcode();
 
   // Check whether opcode is a poison/undef-generating operation
   switch (Opcode) {
   case Instruction::Shl:
   case Instruction::AShr:
   case Instruction::LShr: {
default: {
     else if (Instruction::isBinaryOp(Opcode))
       return false;
     // Be conservative and return true.
     return true;
   }
   }
 }
 
-bool llvm::canCreateUndefOrPoison(const Operator *Op) {
-  return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/false);
+bool llvm::canCreateUndefOrPoison(const Operator *Op, bool ConsiderFlags) {
+  return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/false, ConsiderFlags);
 }
 
-bool llvm::canCreatePoison(const Operator *Op) {
-  return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true);
+bool llvm::canCreatePoison(const Operator *Op, bool ConsiderFlags) {
+  return ::canCreateUndefOrPoison(Op, /*PoisonOnly=*/true, ConsiderFlags);
 }
 
 static bool directlyImpliesPoison(const Value *ValAssumedPoison,
                                   const Value *V, unsigned Depth) {
   if (ValAssumedPoison == V)
     return true;
 
   const unsigned MaxDepth = 2;

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

InstCombinerImpl::pushFreezeToPreventPoisonFromPropagating(FreezeInst &OrigFI) {
   //   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)))
+  if (!OrigOpInst || !OrigOpInst->hasOneUse() || isa<PHINode>(OrigOp))
+    return nullptr;
+
+  // We can't push the freeze through an instruction which can itself create
+  // poison.  If the only source of new poison is flags, we can simply
+  // strip them (since we know the only use is the freeze and nothing can
+  // benefit from them.)
+  if (canCreateUndefOrPoison(cast<Operator>(OrigOp), /*ConsiderFlags*/ false))
     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;
   }
 
+  OrigOpInst->dropPoisonGeneratingFlags();
+
   // 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;
 }

llvm/test/Transforms/InstCombine/freeze.ll

;
   %v2 = and i32 %v1, 1
   %cond = icmp eq i32 %v2, 0
   %cond.fr = freeze i1 %cond
   ret i1 %cond.fr
 }
 
 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:    [[V1_FR:%.*]] = freeze i32 [[V1:%.*]]
+; CHECK-NEXT:    [[V2:%.*]] = shl i32 [[V1_FR]], 1
+; CHECK-NEXT:    [[V3:%.*]] = add i32 [[V2]], 2
+; CHECK-NEXT:    [[V4:%.*]] = or i32 [[V3]], 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
 }

 
 end:
   ret void
 }
 
 
 define i32 @propagate_drop_flags_add(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_flags_add(
-; CHECK-NEXT:    [[V1:%.*]] = add nuw nsw i32 [[ARG:%.*]], 2
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = add i32 [[ARG_FR]], 2
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = add nsw nuw i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_flags_add_foldaway(i32 noundef %arg) {
 ; CHECK-LABEL: @propagate_drop_flags_add_foldaway(
-; CHECK-NEXT:    [[V1:%.*]] = add nuw nsw i32 [[ARG:%.*]], 2
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[V1:%.*]] = add i32 [[ARG:%.*]], 2
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = add nsw nuw i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_flags_sub(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_flags_sub(
-; CHECK-NEXT:    [[V1:%.*]] = add nsw i32 [[ARG:%.*]], -2
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = add i32 [[ARG_FR]], -2
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = sub nsw nuw i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_flags_mul(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_flags_mul(
-; CHECK-NEXT:    [[V1:%.*]] = shl nuw nsw i32 [[ARG:%.*]], 1
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = shl i32 [[ARG_FR]], 1
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = mul nsw nuw i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_flags_udiv(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_flags_udiv(
-; CHECK-NEXT:    [[V1:%.*]] = lshr exact i32 [[ARG:%.*]], 1
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = lshr i32 [[ARG_FR]], 1
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = udiv exact i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_flags_sdiv(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_flags_sdiv(
-; CHECK-NEXT:    [[V1:%.*]] = ashr exact i32 [[ARG:%.*]], 1
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = ashr i32 [[ARG_FR]], 1
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = sdiv exact i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_shl1(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_shl1(
-; CHECK-NEXT:    [[V1:%.*]] = shl nuw nsw i32 [[ARG:%.*]], 2
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = shl i32 [[ARG_FR]], 2
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = shl nsw nuw i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_shl2(i32 %arg, i32 %unknown) {
 ; CHECK-LABEL: @propagate_drop_shl2(
 ; CHECK-NEXT:    [[V1:%.*]] = shl nuw nsw i32 [[ARG:%.*]], [[UNKNOWN:%.*]]
 ; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
 ; CHECK-NEXT:    ret i32 [[V1_FR]]
 ;
   %v1 = shl nsw nuw i32 %arg, %unknown
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_ashr1(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_ashr1(
-; CHECK-NEXT:    [[V1:%.*]] = ashr exact i32 [[ARG:%.*]], 2
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = ashr i32 [[ARG_FR]], 2
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = ashr exact i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_ashr2(i32 %arg, i32 %unknown) {
 ; CHECK-LABEL: @propagate_drop_ashr2(
 ; CHECK-NEXT:    [[V1:%.*]] = ashr exact i32 [[ARG:%.*]], [[UNKNOWN:%.*]]
 ; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
 ; CHECK-NEXT:    ret i32 [[V1_FR]]
 ;
   %v1 = ashr exact i32 %arg, %unknown
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_lshr1(i32 %arg) {
 ; CHECK-LABEL: @propagate_drop_lshr1(
-; CHECK-NEXT:    [[V1:%.*]] = lshr exact i32 [[ARG:%.*]], 2
-; CHECK-NEXT:    [[V1_FR:%.*]] = freeze i32 [[V1]]
-; CHECK-NEXT:    ret i32 [[V1_FR]]
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = lshr i32 [[ARG_FR]], 2
+; CHECK-NEXT:    ret i32 [[V1]]
 ;
   %v1 = lshr exact i32 %arg, 2
   %v1.fr = freeze i32 %v1
   ret i32 %v1.fr
 }
 
 define i32 @propagate_drop_lshr2(i32 %arg, i32 %unknown) {
 ; CHECK-LABEL: @propagate_drop_lshr2(