[PredicateInfo] Reserve adjacent LN_Last defs for the same phi use

[Camsyn]

This patch fixes a missed optimization issue: predicate infos might be lost in phi-use scenarios.

Due to the existence of and-chains, a phi-use might be associated with multiple LN_Last predicate infos.
E.g.,

// TWO LN_Last Predicate Info defs:
// 1. a >= 1
// 2. a < 2
if ( a < 1 || a >= 2) {
  a = 1;
}  
// PHI use of `a`
use(a)

However, previously, popStackUntilDFSScope reserved only ONE LN_Last def for a phi use (i.e., reserve only one of a >= 1 / a < 2), although there might be multiple LN_Last defs for the same phi use.

This patch reserves the adjacent LN_Last defs if they are designated for the same phi use.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Kunqiu Chen (Camsyn)

Changes

Due to the existence of and-chains, a phi-use might be associated with multiple LN_Last predicate infos.
E.g.,

// TWO LN_Last Predicate Info defs:
// 1. a >= 1
// 2. a < 2
if ( a < 1 || a >= 2) {
  a = 1;
}  
// PHI use of `a`
use(a)

However, previously, popStackUntilDFSScope reserved only ONE LN_Last def for a phi use, although there might be multiple LN_Last defs for the same phi use

This patch reserves the adjacent LN_Last defs if they are designated for the same phi use.

---

Full diff: https://github.com/llvm/llvm-project/pull/164577.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/PredicateInfo.cpp (+8-2)
• (modified) llvm/test/Transforms/Util/PredicateInfo/testandor.ll (+27)

diff --git a/llvm/lib/Transforms/Utils/PredicateInfo.cpp b/llvm/lib/Transforms/Utils/PredicateInfo.cpp
index 978d5a25a57c8..8f835e136278c 100644
--- a/llvm/lib/Transforms/Utils/PredicateInfo.cpp
+++ b/llvm/lib/Transforms/Utils/PredicateInfo.cpp
@@ -261,8 +261,14 @@ bool PredicateInfoBuilder::stackIsInScope(const ValueDFSStack &Stack,
   // the stack when we hit the end of the phi uses for a given def.
   const ValueDFS &Top = *Stack.back().V;
   if (Top.LocalNum == LN_Last && Top.PInfo) {
-    if (!VDUse.U)
-      return false;
+    if (!VDUse.U) {
+      assert(VDUse.PInfo && "A non-use VDUse should have a predicate info");
+      // We should reserve adjacent LN_Last defs for the same phi use.
+      return VDUse.LocalNum == LN_Last &&
+             // If the two phi defs have the same edge, they must be designated
+             // for the same succ BB.
+             getBlockEdge(Top.PInfo) == getBlockEdge(VDUse.PInfo);
+    }
     auto *PHI = dyn_cast<PHINode>(VDUse.U->getUser());
     if (!PHI)
       return false;
diff --git a/llvm/test/Transforms/Util/PredicateInfo/testandor.ll b/llvm/test/Transforms/Util/PredicateInfo/testandor.ll
index 2e96a92dd37d4..0020d81dcda3f 100644
--- a/llvm/test/Transforms/Util/PredicateInfo/testandor.ll
+++ b/llvm/test/Transforms/Util/PredicateInfo/testandor.ll
@@ -994,3 +994,30 @@ define void @test_assume_deep_and_tree(i1 %a1) {
   call void @foo(i1 %a15)
   ret void
 }
+
+define i32 @test_and_with_phinode(i32 %x) {
+; CHECK-LABEL: @test_and_with_phinode(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[XGT0:%.*]] = icmp uge i32 [[X:%.*]], 1
+; CHECK-NEXT:    [[XLT1:%.*]] = icmp ult i32 [[X]], 2
+; CHECK-NEXT:    [[AND:%.*]] = and i1 [[XGT0]], [[XLT1]]
+; CHECK:         [[X_0_1:%.*]] = bitcast i32 [[X]] to i32
+; CHECK:         [[X_0_2:%.*]] = bitcast i32 [[X_0_1]] to i32
+; CHECK-NEXT:    br i1 [[AND]], label [[PHI:%.*]], label [[NOPE:%.*]]
+; CHECK:       nope:
+; CHECK-NEXT:    br label [[PHI]]
+; CHECK:       phi:
+; CHECK-NEXT:    [[RES:%.*]] = phi i32 [ [[X_0_2]], [[ENTRY:%.*]] ], [ 1, [[NOPE]] ]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+entry:
+  %xgt0 = icmp uge i32 %x, 1
+  %xlt1 = icmp ult i32 %x, 2
+  %and = and i1 %xgt0, %xlt1
+  br i1 %and, label %phi, label %nope
+nope:
+  br label %phi
+phi:
+  %res = phi i32 [ %x, %entry ], [ 1, %nope ]
+  ret i32 %res
+}

[nikic]

To clarify, as it's not entirely clear from the PR description: Is this fixing a missed optimization issue (i.e. one of the predicates being lost on the phi edge)?

[Camsyn]

To clarify, as it's not entirely clear from the PR description: Is this fixing a missed optimization issue (i.e. one of the predicates being lost on the phi edge)?

Yes, previously, one of the predicates was lost on the phi edge, i.e., "reserved only ONE LN_Last def for a phi use".
You can refer to commit 4835c17 for details.

[nikic]

Thanks. The fix looks right to me. I'd suggest adding a test to SCCP as well, to make sure this works end-to-end.

[nikic]

Just as a side node, this Top.PInfo check looks unnecessary. The stack only contains defs.

[nikic]

LGTM