[SCEV] Propagate nuw/nsw flags from or disjoint (#197236)

The Scalar Evolution analysis converts the `or disjoint` instruction
into a synthetic add with the `IsNSW` and `IsNUW` flags set, but the
real instruction reference set to `nullptr`. The flags propagation logic
only reads the real instruction flags, when the instruction reference is
present. This behavior silently ignores the flags for the synthetic add,
which can lead to missed optimization opportunities.

This patch inserts the pointer to the actual instruction into the
`BinaryOp` object created for the synthetic add. It also adjust the
flags handling logic to threat the `disjoint` flag as both `nuw` and
`nsw` as if they were set on the `add` instruction.

Assisted-by: Claude Opus 4.6 <noreply@anthropic.com>

llvm/lib/Analysis/ScalarEvolution.cpp

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 3d17c2a..6b02806 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -5418,9 +5418,14 @@
 
   case Instruction::Or: {
     // 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);
+    if (cast<PossiblyDisjointInst>(Op)->isDisjoint()) {
+      BinaryOp BinOp(Instruction::Add, Op->getOperand(0), Op->getOperand(1),
+                     /*IsNSW=*/true, /*IsNUW=*/true);
+      // Keep the reference to the original instruction so that we can later
+      // check whether it can produce poison value or not.
+      BinOp.Op = Op;
+      return BinOp;
+    }
     return BinaryOp(Op);
   }
 
@@ -7459,10 +7464,15 @@
 
   // Return early if there are no flags to propagate to the SCEV.
   SCEV::NoWrapFlags Flags = SCEV::FlagAnyWrap;
-  if (BinOp->hasNoUnsignedWrap())
-    Flags = ScalarEvolution::setFlags(Flags, SCEV::FlagNUW);
-  if (BinOp->hasNoSignedWrap())
-    Flags = ScalarEvolution::setFlags(Flags, SCEV::FlagNSW);
+  if (auto *PDI = dyn_cast<PossiblyDisjointInst>(BinOp);
+      PDI && PDI->isDisjoint()) {
+    Flags = ScalarEvolution::setFlags(SCEV::FlagNUW, SCEV::FlagNSW);
+  } else {
+    if (BinOp->hasNoUnsignedWrap())
+      Flags = ScalarEvolution::setFlags(Flags, SCEV::FlagNUW);
+    if (BinOp->hasNoSignedWrap())
+      Flags = ScalarEvolution::setFlags(Flags, SCEV::FlagNSW);
+  }
   if (Flags == SCEV::FlagAnyWrap)
     return SCEV::FlagAnyWrap;
 

llvm/test/Analysis/ScalarEvolution/add-like-or.ll

diff --git a/llvm/test/Analysis/ScalarEvolution/add-like-or.ll b/llvm/test/Analysis/ScalarEvolution/add-like-or.ll
index 71feb8a..69059d5 100644
--- a/llvm/test/Analysis/ScalarEvolution/add-like-or.ll
+++ b/llvm/test/Analysis/ScalarEvolution/add-like-or.ll
@@ -37,18 +37,47 @@
   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:    --> (%x + %y)<nuw><nsw> 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
 }
 
+; or disjoint with zext: flags should enable zext simplification.
+define noundef i64 @or-disjoint-zext(i32 noundef %x) {
+; CHECK-LABEL: 'or-disjoint-zext'
+; CHECK-NEXT:  Classifying expressions for: @or-disjoint-zext
+; CHECK-NEXT:    %or = or disjoint i32 %x, 1
+; CHECK-NEXT:    --> (1 + %x)<nuw><nsw> U: [1,0) S: [-2147483647,-2147483648)
+; CHECK-NEXT:    %zext = zext i32 %or to i64
+; CHECK-NEXT:    --> (1 + (zext i32 %x to i64))<nuw><nsw> U: [1,4294967297) S: [1,4294967297)
+; CHECK-NEXT:  Determining loop execution counts for: @or-disjoint-zext
+;
+  %or = or disjoint i32 %x, 1
+  %zext = zext i32 %or to i64
+  ret i64 %zext
+}
+
+; or disjoint with larger constant and zext.
+define noundef i64 @or-disjoint-zext-7(i32 noundef %x) {
+; CHECK-LABEL: 'or-disjoint-zext-7'
+; CHECK-NEXT:  Classifying expressions for: @or-disjoint-zext-7
+; CHECK-NEXT:    %or = or disjoint i32 %x, 7
+; CHECK-NEXT:    --> (7 + %x)<nuw><nsw> U: [7,0) S: [-2147483641,-2147483648)
+; CHECK-NEXT:    %zext = zext i32 %or to i64
+; CHECK-NEXT:    --> (7 + (zext i32 %x to i64))<nuw><nsw> U: [7,4294967303) S: [7,4294967303)
+; CHECK-NEXT:  Determining loop execution counts for: @or-disjoint-zext-7
+;
+  %or = or disjoint i32 %x, 7
+  %zext = zext i32 %or to i64
+  ret i64 %zext
+}
+
 define void @mask-high(i64 %arg, ptr dereferenceable(4) %arg1) {
 ; CHECK-LABEL: 'mask-high'
 ; CHECK-NEXT:  Classifying expressions for: @mask-high

llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll

diff --git a/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll b/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
index 593888f..f7ebfbb 100644
--- a/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
+++ b/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
@@ -2181,3 +2181,17 @@
 exit:
   ret void
 }
+
+define void @add_nuw_zext_no_flags(i32 %x) {
+; CHECK-LABEL: 'add_nuw_zext_no_flags'
+; CHECK-NEXT:  Classifying expressions for: @add_nuw_zext_no_flags
+; CHECK-NEXT:    %add = add nuw i32 %x, 1
+; CHECK-NEXT:    --> (1 + %x) U: full-set S: full-set
+; CHECK-NEXT:    %zext = zext i32 %add to i64
+; CHECK-NEXT:    --> (zext i32 (1 + %x) to i64) U: [0,4294967296) S: [0,4294967296)
+; CHECK-NEXT:  Determining loop execution counts for: @add_nuw_zext_no_flags
+;
+  %add = add nuw i32 %x, 1
+  %zext = zext i32 %add to i64
+  ret void
+}