[Codegenprepare] Use IV increment instead of IV if we can prove it is not a poisoned value

[mkazantsev]

Summary

We restrained from replacing uses of IV with IV.next if IV.next is potentially poison.
This patch loosens this restricton for one important particular case: memory instruction
is dominated by loop exit by condition icmp iV.next, X. In this case we can safely assume
that, whenever IV.next is poison, we exit the loop at this point. So only non-poisoned
values may reach the memory instruction.

Note that this transform reduces overlap of live ranges of IV and IV.next, so potentially it
can lead to better opt decisions overall.

[reames]

llvm/lib/CodeGen/CodeGenPrepare.cpp

3860

The code that follows looks analogous to programUndefinedIfPoison from ValueTracking.hpp. (Though you handle slightly different cases.) I'd encourage you to split this into a helper function and use analogous naming conventions to make it easier to follow.

Side note - I'd suggest trying the existing implementation from ValueTracking. While the reason is definitely different, if that's enough to catch your interesting cases, you can save a bunch of work. :)

3874

Detail on the proof here -- the fact the branch is a loop exit is definitely irrelevant.

It's not clear to me whether branching on poison is immediate UB. LangRef seems to say "yes", but the code in ValueTracking seems to say "no". We need to draw in the experts here. (Juneyoung)

[spatel]

llvm/lib/CodeGen/CodeGenPrepare.cpp

3874

cc @aqjune @lebedev.ri @nikic for branch-on-poison question.

[nikic]

llvm/lib/CodeGen/CodeGenPrepare.cpp

3874

See D92739. Branch on poison is UB, but we're currently being conservative about it due to some known bugs (imho to our detriment).

[aqjune]

llvm/lib/CodeGen/CodeGenPrepare.cpp

3874

Hi,
Yes, branching on poison is immediate UB, but it isn't listed in getGuaranteedWellDefinedOps (which is called by programUndefinedIfPoison). It's because there are a few known transformations that are poison-unsafe but still exist, and adding the case to getGuaranteedWellDefinedOps globally impacts optimizations, increasing opportunities for miscompilation bug.

The two most frequently happening (mis)transformations are short circuiting in SimplifyCFG (D95026) and select i1 -> and/or i1 in InstCombine (D93840). And then, there are a few optimizations that require either freeze or other workarounds, but they happen less frequently and fixing them is expected to have smaller performance impact.

For the SimplifyCFG patch: I believe it is almost there; combined with the enhancement in loop unswitch (which is mentioned in the patch) its impact is relatively small, as its asmdiff for llvm-testsuite result mentioned in its comments.
It would be great if the patches could be reviewed and get in. :)

For the select -> and/or InstCombine opt: its fix is hidden under a flag, but I and people have already written many patches to resolve possible regressions (D93065 has links).

[reames]

llvm/lib/CodeGen/CodeGenPrepare.cpp

3874

@aqjune Given the proposed use case here is CGP, which is well after the specific mid-level optimizer cases you mention, what is your take on the practicality of using branching on poison as immediate UB here? I'm tempted to say we should go ahead, but you're much more current with the practical status and tradeoffs than I am.

[aqjune]

llvm/lib/CodeGen/CodeGenPrepare.cpp

3874

Since this transformation is correct and does not globally affect others, I think this is fine.
I remember there was no poison-unsafe transformation in CodeGenPrepare (there were a few in the past, but finally fixed). Also, CodeGenPrepare does not use or update SCEV, which is the analysis that was mainly impacted by getGuaranteedWellDefinedOps's branch-poison change.

llvm/lib/CodeGen/CodeGenPrepare.cpp

auto GetIVStep = [this](const Value * V)
     if (match(IVInc,
               m_ExtractValue<0>(m_Intrinsic<Intrinsic::usub_with_overflow>(
                   m_Specific(PN), m_ConstantInt(Step)))))
       return std::make_pair(IVInc, -Step->getValue());
     if (match(IVInc,
               m_ExtractValue<0>(m_Intrinsic<Intrinsic::uadd_with_overflow>(
                   m_Specific(PN), m_ConstantInt(Step)))))
       return std::make_pair(IVInc, Step->getValue());
-    // TODO: The result of the intrinsics above is two-compliment. However when
-    // IV inc is expressed as add or sub, iv.next is potentially a poison value.
-    // If it has nuw or nsw flags, we need to make sure that these flags are
-    // inferrable at the point of memory instruction. Otherwise we are replacing
-    // well-defined two-compliment computation with poison. Currently, to avoid
-    // potentially complex analysis needed to prove this, we reject such cases.
-    auto *OIVInc = dyn_cast<OverflowingBinaryOperator>(IVInc);
-    if (!OIVInc || OIVInc->hasNoSignedWrap() || OIVInc->hasNoUnsignedWrap())
+    if (auto *OIVInc = dyn_cast<OverflowingBinaryOperator>(IVInc))
+      if (OIVInc->hasNoSignedWrap() || OIVInc->hasNoUnsignedWrap()) {
+        bool ProvedNoPoison = false;
+        // If IVInc is represented as wrapping binary operator with no wrap
+        // flags, it may be poison. To legally replace the use of IV with use of
+        // IV.next, we need to prove that memory instuction will never execute
+        // with this poison. In general it's a hard thing to prove, but there is
+        // a widely-used case where the proof is trivial. If MemoryInstr is
+        // dominated by loop-exiting branch by icmp conditon involving IV.next,
+        // we can always assume that poison value will exit this loop at this
+        // point and we never reach MemoryInst.
+        for (auto *Node = DT.getNode(MemoryInst->getParent())->getIDom();
+             L->contains(Node->getBlock()); Node = Node->getIDom()) {
+          BasicBlock *BB = Node->getBlock();
+          BasicBlock *IfTrue, *IfFalse;
+          ICmpInst::Predicate Pred;
+          if (match(BB->getTerminator(),
+                    m_Br(m_ICmp(Pred, m_Specific(IVInc), m_Value()),
+                         m_BasicBlock(IfTrue), m_BasicBlock(IfFalse))) ||
+              match(BB->getTerminator(),
+                    m_Br(m_ICmp(Pred, m_Value(), m_Specific(IVInc)),
+                         m_BasicBlock(IfTrue), m_BasicBlock(IfFalse))))
+            if (!L->contains(IfTrue) || !L->contains(IfFalse)) {
+              ProvedNoPoison = true;
+              break;
+            }
+        }
+        if (!ProvedNoPoison)
           return None;
+      }
     if (match(IVInc, m_Sub(m_Specific(PN), m_ConstantInt(Step))))
       return std::make_pair(IVInc, -Step->getValue());
     if (match(IVInc, m_Add(m_Specific(PN), m_ConstantInt(Step))))
       return std::make_pair(IVInc, Step->getValue());
     return None;
   };
 
   // Try to account for the following special case:

llvm/test/CodeGen/X86/overflowing-iv.ll

 ; CHECK-NEXT:    [[LEN_PLUS_1:%.*]] = add i64 [[LEN:%.*]], 1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[IV_NEXT]] = add nsw i64 [[IV]], 1
 ; CHECK-NEXT:    [[COND_1:%.*]] = icmp eq i64 [[IV_NEXT]], [[LEN_PLUS_1]]
 ; CHECK-NEXT:    br i1 [[COND_1]], label [[EXIT:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
-; CHECK-NEXT:    [[SUNKADDR:%.*]] = mul i64 [[IV]], 4
+; CHECK-NEXT:    [[SUNKADDR:%.*]] = mul i64 [[IV_NEXT]], 4
 ; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[P:%.*]] to i8*
 ; CHECK-NEXT:    [[SUNKADDR1:%.*]] = getelementptr i8, i8* [[TMP0]], i64 [[SUNKADDR]]
-; CHECK-NEXT:    [[SUNKADDR2:%.*]] = getelementptr i8, i8* [[SUNKADDR1]], i64 -4
+; CHECK-NEXT:    [[SUNKADDR2:%.*]] = getelementptr i8, i8* [[SUNKADDR1]], i64 -8
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[SUNKADDR2]] to i32*
 ; CHECK-NEXT:    [[LOADED:%.*]] = load atomic i32, i32* [[TMP1]] unordered, align 4
 ; CHECK-NEXT:    [[COND_2:%.*]] = icmp eq i32 [[LOADED]], [[X:%.*]]
 ; CHECK-NEXT:    br i1 [[COND_2]], label [[FAILURE:%.*]], label [[LOOP]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i32 -1
 ; CHECK:       failure:
 ; CHECK-NEXT:    unreachable

 ; CHECK-NEXT:    [[LEN_PLUS_1:%.*]] = add i64 [[LEN:%.*]], 1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw i64 [[IV]], 1
 ; CHECK-NEXT:    [[COND_1:%.*]] = icmp eq i64 [[IV_NEXT]], [[LEN_PLUS_1]]
 ; CHECK-NEXT:    br i1 [[COND_1]], label [[EXIT:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
-; CHECK-NEXT:    [[SUNKADDR:%.*]] = mul i64 [[IV]], 4
+; CHECK-NEXT:    [[SUNKADDR:%.*]] = mul i64 [[IV_NEXT]], 4
 ; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[P:%.*]] to i8*
 ; CHECK-NEXT:    [[SUNKADDR1:%.*]] = getelementptr i8, i8* [[TMP0]], i64 [[SUNKADDR]]
-; CHECK-NEXT:    [[SUNKADDR2:%.*]] = getelementptr i8, i8* [[SUNKADDR1]], i64 -4
+; CHECK-NEXT:    [[SUNKADDR2:%.*]] = getelementptr i8, i8* [[SUNKADDR1]], i64 -8
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[SUNKADDR2]] to i32*
 ; CHECK-NEXT:    [[LOADED:%.*]] = load atomic i32, i32* [[TMP1]] unordered, align 4
 ; CHECK-NEXT:    [[COND_2:%.*]] = icmp eq i32 [[LOADED]], [[X:%.*]]
 ; CHECK-NEXT:    br i1 [[COND_2]], label [[FAILURE:%.*]], label [[LOOP]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i32 -1
 ; CHECK:       failure:
 ; CHECK-NEXT:    unreachable