[llvm] [SCEV] Use or disjoint flag (PR #74467)

[Nikita Popov via llvm-commits]

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

Use the disjoint flag to convert or to add instead of calling the haveNoCommonBitsSet() ValueTracking query. This ensures that we can reliably undo add -> or canonicalization, even in cases where the necessary information has been lost or is too complex to reinfer in SCEV.

I have updated the bulk of the test coverage to add the necessary disjoint flags in advance.

>From cd90987377ff949d6651ea8e3a11e991356d490b Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Tue, 5 Dec 2023 14:15:25 +0100
Subject: [PATCH] [SCEV] Use or disjoint flag

Use the disjoint flag to convert or to add instead of calling the
haveNoCommonBitsSet() ValueTracking query. This ensures that we
can reliably undo add -> or canonicalization, even in cases where
the necessary information has been lost or is too complex to
reinfer in SCEV.

I have updated the bulk of the test coverage to add the necessary
disjoint flags in advance.
---
 llvm/lib/Analysis/ScalarEvolution.cpp         |  7 +---
 .../Analysis/ScalarEvolution/add-like-or.ll   | 42 +++++++++++++++++--
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 9670e5309b32f..451ae73dbd601 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -5223,11 +5223,8 @@ static std::optional<BinaryOp> MatchBinaryOp(Value *V, const DataLayout &DL,
     return BinaryOp(Op);
 
   case Instruction::Or: {
-    // LLVM loves to convert `add` of operands with no common bits
-    // into an `or`. But SCEV really doesn't deal with `or` that well,
-    // so try extra hard to recognize this `or` as an `add`.
-    if (haveNoCommonBitsSet(Op->getOperand(0), Op->getOperand(1),
-                            SimplifyQuery(DL, &DT, &AC, CxtI)))
+    // Convert or disjoint into add nuw nsw.
+    if (cast<PossiblyDisjointInst>(Op)->isDisjoint())
       return BinaryOp(Instruction::Add, Op->getOperand(0), Op->getOperand(1),
                       /*IsNSW=*/true, /*IsNUW=*/true);
     return BinaryOp(Op);
diff --git a/llvm/test/Analysis/ScalarEvolution/add-like-or.ll b/llvm/test/Analysis/ScalarEvolution/add-like-or.ll
index 38b6c44fdf356..86bb9058832b4 100644
--- a/llvm/test/Analysis/ScalarEvolution/add-like-or.ll
+++ b/llvm/test/Analysis/ScalarEvolution/add-like-or.ll
@@ -6,15 +6,49 @@ define i8 @or-of-constant-with-no-common-bits-set(i8 %x, i8 %y) {
 ; CHECK-NEXT:  Classifying expressions for: @or-of-constant-with-no-common-bits-set
 ; CHECK-NEXT:    %t0 = shl i8 %x, 2
 ; CHECK-NEXT:    --> (4 * %x) U: [0,-3) S: [-128,125)
-; CHECK-NEXT:    %r = or i8 %t0, 3
+; CHECK-NEXT:    %r = or disjoint i8 %t0, 3
 ; CHECK-NEXT:    --> (3 + (4 * %x))<nuw><nsw> U: [3,0) S: [-125,-128)
 ; CHECK-NEXT:  Determining loop execution counts for: @or-of-constant-with-no-common-bits-set
 ;
   %t0 = shl i8 %x, 2
-  %r = or i8 %t0, 3
+  %r = or disjoint i8 %t0, 3
   ret i8 %r
 }
 
+define i8 @or-disjoint(i8 %x, i8 %y) {
+; CHECK-LABEL: 'or-disjoint'
+; CHECK-NEXT:  Classifying expressions for: @or-disjoint
+; CHECK-NEXT:    %or = or disjoint i8 %x, %y
+; CHECK-NEXT:    --> (%x + %y) U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @or-disjoint
+;
+  %or = or disjoint i8 %x, %y
+  ret i8 %or
+}
+
+define i8 @or-no-disjoint(i8 %x, i8 %y) {
+; CHECK-LABEL: 'or-no-disjoint'
+; CHECK-NEXT:  Classifying expressions for: @or-no-disjoint
+; CHECK-NEXT:    %or = or i8 %x, %y
+; CHECK-NEXT:    --> %or U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @or-no-disjoint
+;
+  %or = or i8 %x, %y
+  ret i8 %or
+}
+
+; FIXME: We could add nuw nsw flags here.
+define noundef i8 @or-disjoint-transfer-flags(i8 %x, i8 %y) {
+; CHECK-LABEL: 'or-disjoint-transfer-flags'
+; CHECK-NEXT:  Classifying expressions for: @or-disjoint-transfer-flags
+; CHECK-NEXT:    %or = or disjoint i8 %x, %y
+; CHECK-NEXT:    --> (%x + %y) U: full-set S: full-set
+; CHECK-NEXT:  Determining loop execution counts for: @or-disjoint-transfer-flags
+;
+  %or = or disjoint i8 %x, %y
+  ret i8 %or
+}
+
 define void @mask-high(i64 %arg, ptr dereferenceable(4) %arg1) {
 ; CHECK-LABEL: 'mask-high'
 ; CHECK-NEXT:  Classifying expressions for: @mask-high
@@ -24,7 +58,7 @@ define void @mask-high(i64 %arg, ptr dereferenceable(4) %arg1) {
 ; CHECK-NEXT:    --> (sext i32 %i to i64) U: [-2147483648,2147483648) S: [-2147483648,2147483648)
 ; CHECK-NEXT:    %i3 = and i64 %arg, -16
 ; CHECK-NEXT:    --> (16 * (%arg /u 16))<nuw> U: [0,-15) S: [-9223372036854775808,9223372036854775793)
-; CHECK-NEXT:    %i4 = or i64 1, %i3
+; CHECK-NEXT:    %i4 = or disjoint i64 1, %i3
 ; CHECK-NEXT:    --> (1 + (16 * (%arg /u 16))<nuw>)<nuw><nsw> U: [1,-14) S: [-9223372036854775807,9223372036854775794)
 ; CHECK-NEXT:    %i7 = phi i64 [ %i4, %bb ], [ %i8, %bb6 ]
 ; CHECK-NEXT:    --> {(1 + (16 * (%arg /u 16))<nuw>)<nuw><nsw>,+,1}<%bb6> U: full-set S: full-set Exits: ((sext i32 %i to i64) smax (1 + (16 * (%arg /u 16))<nuw>)<nuw><nsw>) LoopDispositions: { %bb6: Computable }
@@ -42,7 +76,7 @@ bb:
   %i = load i32, ptr %arg1, align 4
   %i2 = sext i32 %i to i64
   %i3 = and i64 %arg, -16
-  %i4 = or i64 1, %i3
+  %i4 = or disjoint i64 1, %i3
   %i5 = icmp sgt i64 %i4, %i2
   br i1 %i5, label %bb10, label %bb6