[llvm] [JumpThreading] Don't phi translate past loop phi (PR #70664)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/70664

When evaluating comparisons in predecessors, phi operands are translated into the predecessor. If the translation is across a backedge, this means that the two operands of the icmp will be from two different loop iterations, resulting in incorrect simplification.

Fix this by not performing the phi translation for phis in loop headers.

Note: This is not a complete fix. If the
jump-threading-across-loop-headers option is enabled, the LoopHeaders variable does not get populated. Additional changes will be needed to fix that case.

Related to https://github.com/llvm/llvm-project/issues/70651.

>From aaa176843fc435b6a14aa54acc9e4413a2f3d183 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Mon, 30 Oct 2023 15:43:07 +0100
Subject: [PATCH] [JumpThreading] Don't phi translate past loop phi

When evaluating comparisons in predecessors, phi operands are
translated into the predecessor. If the translation is across a
backedge, this means that the two operands of the icmp will be
from two different loop iterations, resulting in incorrect
simplification.

Fix this by not performing the phi translation for phis in loop
headers.

Note: This is not a complete fix. If the
jump-threading-across-loop-headers option is enabled, the
LoopHeaders variable does not get populated. Additional changes
will be needed to fix that case.

Related to https://github.com/llvm/llvm-project/issues/70651.
---
 llvm/lib/Transforms/Scalar/JumpThreading.cpp  |  5 ++++-
 llvm/test/Transforms/JumpThreading/pr70651.ll | 19 +++++++++++++++++--
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index f2b9d784ead8af3..7a8128c5b6c0901 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -761,7 +761,10 @@ bool JumpThreadingPass::computeValueKnownInPredecessorsImpl(
     PHINode *PN = dyn_cast<PHINode>(CmpLHS);
     if (!PN)
       PN = dyn_cast<PHINode>(CmpRHS);
-    if (PN && PN->getParent() == BB) {
+    // Do not perform phi translation across a loop header phi, because this
+    // may result in comparison of values from two different loop iterations.
+    // FIXME: This check is broken if LoopHeaders is not populated.
+    if (PN && PN->getParent() == BB && !LoopHeaders.contains(BB)) {
       const DataLayout &DL = PN->getModule()->getDataLayout();
       // We can do this simplification if any comparisons fold to true or false.
       // See if any do.
diff --git a/llvm/test/Transforms/JumpThreading/pr70651.ll b/llvm/test/Transforms/JumpThreading/pr70651.ll
index a156be541874a6a..2c6059bfdb3d566 100644
--- a/llvm/test/Transforms/JumpThreading/pr70651.ll
+++ b/llvm/test/Transforms/JumpThreading/pr70651.ll
@@ -1,7 +1,8 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
 ; RUN: opt -S -passes=jump-threading < %s | FileCheck %s
+; RUN: opt -S -passes=jump-threading -jump-threading-across-loop-headers < %s | FileCheck %s --check-prefix=THREAD-LOOP
 
-; FIXME: This is a miscompile.
+; FIXME: This is a miscompile if -jump-threading-across-loop-headers is enabled.
 define i64 @test(i64 %v) {
 ; CHECK-LABEL: define i64 @test(
 ; CHECK-SAME: i64 [[V:%.*]]) {
@@ -12,10 +13,24 @@ define i64 @test(i64 %v) {
 ; CHECK-NEXT:    [[SUM:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[SUM_NEXT:%.*]], [[FOR_BODY]] ]
 ; CHECK-NEXT:    [[SUM_NEXT]] = add i64 [[SUM]], [[V]]
 ; CHECK-NEXT:    [[OVERFLOW:%.*]] = icmp ult i64 [[SUM_NEXT]], [[SUM]]
-; CHECK-NEXT:    br i1 [[V_NONNEG]], label [[FOR_BODY]], label [[EXIT:%.*]]
+; CHECK-NEXT:    [[CMP:%.*]] = xor i1 [[V_NONNEG]], [[OVERFLOW]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i64 [[SUM]]
 ;
+; THREAD-LOOP-LABEL: define i64 @test(
+; THREAD-LOOP-SAME: i64 [[V:%.*]]) {
+; THREAD-LOOP-NEXT:  entry:
+; THREAD-LOOP-NEXT:    [[V_NONNEG:%.*]] = icmp sgt i64 [[V]], -1
+; THREAD-LOOP-NEXT:    br label [[FOR_BODY:%.*]]
+; THREAD-LOOP:       for.body:
+; THREAD-LOOP-NEXT:    [[SUM:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[SUM_NEXT:%.*]], [[FOR_BODY]] ]
+; THREAD-LOOP-NEXT:    [[SUM_NEXT]] = add i64 [[SUM]], [[V]]
+; THREAD-LOOP-NEXT:    [[OVERFLOW:%.*]] = icmp ult i64 [[SUM_NEXT]], [[SUM]]
+; THREAD-LOOP-NEXT:    br i1 [[V_NONNEG]], label [[FOR_BODY]], label [[EXIT:%.*]]
+; THREAD-LOOP:       exit:
+; THREAD-LOOP-NEXT:    ret i64 [[SUM]]
+;
 entry:
   %v.nonneg = icmp sgt i64 %v, -1
   br label %for.body