[llvm] [SCEV] Allow reuse of provably-disjoint `or` in canReuseInstruction (PR #216355)

[Nick Riasanovsky via llvm-commits]

https://github.com/njriasan created https://github.com/llvm/llvm-project/pull/216355

canReuseInstruction() unconditionally rejects any `or disjoint` found in the operand graph. `disjoint` is a poison-generating flag, so the generic machinery would otherwise just drop it on reuse -- but dropping it changes the instruction's meaning, since SCEV models `or disjoint x, y` as `x + y` and a plain `or` is only an add when the operands share no bits.

Relax the bail-out to the cases where disjointness cannot be re-established independently of the flag. haveNoCommonBitsSet() is queried with UseInstrInfo=false so the proof cannot rest on flags or metadata that DropPoisonGeneratingInsts is about to strip. DC is left null, so no position-dependent branch facts enter the proof either, and the reused instruction is never moved by any of the callers.

This lets SLSR strength-reduce index arithmetic in cases where InstCombine has canonicalized an `add` into an `or disjoint`.

This arose from an internal Triton use case I can't share where the increased register usages resulted in a 33% latency regression (now resolved).

Assisted-by: Codex

>From 1c60b7b4373c0ca98a6fb7a5e4b4481c0266ceec Mon Sep 17 00:00:00 2001
From: Nick Riasanovsky <njriasan at meta.com>
Date: Fri, 14 Aug 2026 09:08:30 -0700
Subject: [PATCH] [SCEV] Allow reuse of provably-disjoint `or` in
 canReuseInstruction

canReuseInstruction() unconditionally rejects any `or disjoint` found in
the operand graph. `disjoint` is a poison-generating flag, so the generic
machinery would otherwise just drop it on reuse -- but dropping it changes
the instruction's meaning, since SCEV models `or disjoint x, y` as `x + y`
and a plain `or` is only an add when the operands share no bits.

Relax the bail-out to the cases where disjointness cannot be re-established
independently of the flag. haveNoCommonBitsSet() is queried with
UseInstrInfo=false so the proof cannot rest on flags or metadata that
DropPoisonGeneratingInsts is about to strip. DC is left null, so no
position-dependent branch facts enter the proof either, and the reused
instruction is never moved by any of the callers.

This lets SLSR strength-reduce index arithmetic in cases where InstCombine
has canonicalized an `add` into an `or disjoint`.

Assisted-by: Codex
---
 llvm/lib/Analysis/ScalarEvolution.cpp         | 13 +++++---
 .../StraightLineStrengthReduce/slsr-mul.ll    | 33 +++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 27a1a20bcdf79..c481bab6fb3ec 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -4315,11 +4315,16 @@ bool ScalarEvolution::canReuseInstruction(
     if (!I)
       return false;
 
-    // Disjoint or instructions are interpreted as adds by SCEV. However, we
-    // can't replace an arbitrary add with disjoint or, even if we drop the
-    // flag. We would need to convert the or into an add.
+    // Disjoint or instructions are interpreted as adds by SCEV. Reusing one is
+    // only safe if the operands are independently known not to overlap; after
+    // dropping the flag, an arbitrary or would not represent the add. Exclude
+    // instruction annotations from the proof because successful reuse may drop
+    // them below.
     if (auto *PDI = dyn_cast<PossiblyDisjointInst>(I))
-      if (PDI->isDisjoint())
+      if (PDI->isDisjoint() &&
+          !haveNoCommonBitsSet(PDI->getOperand(0), PDI->getOperand(1),
+                               SimplifyQuery(getDataLayout(), &DT, &AC, I,
+                                             /*UseInstrInfo=*/false)))
         return false;
 
     // FIXME: Ignore vscale, even though it technically could be poison. Do this
diff --git a/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll b/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll
index 485c666a8779b..355242fbaa377 100644
--- a/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll
+++ b/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll
@@ -68,6 +68,39 @@ define void @or(i32 %a, i32 %s) {
   ret void
 }
 
+define void @disjoint_or_known_disjoint(i32 %a, i32 %s) {
+; CHECK-LABEL: @disjoint_or_known_disjoint(
+  %b = shl i32 %a, 4
+  %b1 = or disjoint i32 %b, 1
+; CHECK: %b1 = or i32 %b, 1
+  %mul1 = mul i32 %b1, %s
+; CHECK: %mul1 = mul i32 %b1, %s
+  call void @foo(i32 %mul1)
+
+  %b2 = add i32 %b, 2
+  %mul2 = mul i32 %b2, %s
+; CHECK-NOT: %mul2 = mul
+; CHECK: %mul2 = add i32 %mul1, %s
+  call void @foo(i32 %mul2)
+
+  ret void
+}
+
+define void @disjoint_or_may_overlap(i32 %b, i32 %s) {
+; CHECK-LABEL: @disjoint_or_may_overlap(
+  %b1 = or disjoint i32 %b, 1
+  %mul1 = mul i32 %b1, %s
+; CHECK: %mul1 = mul i32 %b1, %s
+  call void @foo(i32 %mul1)
+
+  %b2 = add i32 %b, 2
+  %mul2 = mul i32 %b2, %s
+; CHECK: %mul2 = mul i32 %b2, %s
+  call void @foo(i32 %mul2)
+
+  ret void
+}
+
 ; foo(a * b)
 ; foo((a + 1) * b)
 ; foo(a * (b + 1))