[LICM] Preserve Disjoint flag on OR when hoisting. (#140266)

Update hoistBOAssociation to preserve Disjoint flags on the newly
created instructions if both ORs are disjoint.

Fixes https://github.com/llvm/llvm-project/issues/139625.

PR: https://github.com/llvm/llvm-project/pull/140266

2 files changed, 70 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 62ef40c..7d89a13 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -2877,6 +2877,13 @@ static bool hoistBOAssociation(Instruction &I, Loop &L,
     if (auto *I = dyn_cast<Instruction>(Inv))
       I->setFastMathFlags(Intersect);
     NewBO->setFastMathFlags(Intersect);
+  } else if (Opcode == Instruction::Or) {
+    bool Disjoint = cast<PossiblyDisjointInst>(BO)->isDisjoint() &&
+                    cast<PossiblyDisjointInst>(BO0)->isDisjoint();
+    // If `Inv` was not constant-folded, a new Instruction has been created.
+    if (auto *I = dyn_cast<PossiblyDisjointInst>(Inv))
+      I->setIsDisjoint(Disjoint);
+    cast<PossiblyDisjointInst>(NewBO)->setIsDisjoint(Disjoint);
   }
 
   BO->replaceAllUsesWith(NewBO);
diff --git a/llvm/test/Transforms/LICM/hoist-binop.ll b/llvm/test/Transforms/LICM/hoist-binop.ll
index ea7d96c..000b1fc 100644
--- a/llvm/test/Transforms/LICM/hoist-binop.ll
+++ b/llvm/test/Transforms/LICM/hoist-binop.ll
@@ -723,6 +723,69 @@ loop:
   br label %loop
 }
 
+; Trivially hoist or disjoint.
+define void @or_all_disjoint(i64 %c1, i64 %c2) {
+; CHECK-LABEL: @or_all_disjoint(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = or disjoint i64 [[C1:%.*]], [[C2:%.*]]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDEX_NEXT_REASS:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[INDEX_NEXT_REASS]] = or disjoint i64 [[INDEX]], [[INVARIANT_OP]]
+; CHECK-NEXT:    br label [[LOOP]]
+;
+entry:
+  br label %loop
+
+loop:
+  %index = phi i64 [ 0, %entry ], [ %index.next, %loop ]
+  %step.add = or disjoint i64 %index, %c1
+  %index.next = or disjoint i64 %c2, %step.add
+  br label %loop
+}
+
+; Trivially hoist or, disjoint on first or only .
+define void @or_disjoint_on_first_or_only(i64 %c1, i64 %c2) {
+; CHECK-LABEL: @or_disjoint_on_first_or_only(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = or i64 [[C1:%.*]], [[C2:%.*]]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDEX_NEXT_REASS:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[INDEX_NEXT_REASS]] = or i64 [[INDEX]], [[INVARIANT_OP]]
+; CHECK-NEXT:    br label [[LOOP]]
+;
+entry:
+  br label %loop
+
+loop:
+  %index = phi i64 [ 0, %entry ], [ %index.next, %loop ]
+  %step.add = or i64 %index, %c1
+  %index.next = or disjoint i64 %c2, %step.add
+  br label %loop
+}
+
+; Trivially hoist or, disjoint on second or only .
+define void @or_disjoint_on_second_or_only(i64 %c1, i64 %c2) {
+; CHECK-LABEL: @or_disjoint_on_second_or_only(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = or i64 [[C1:%.*]], [[C2:%.*]]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDEX_NEXT_REASS:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[INDEX_NEXT_REASS]] = or i64 [[INDEX]], [[INVARIANT_OP]]
+; CHECK-NEXT:    br label [[LOOP]]
+;
+entry:
+  br label %loop
+
+loop:
+  %index = phi i64 [ 0, %entry ], [ %index.next, %loop ]
+  %step.add = or disjoint i64 %index, %c1
+  %index.next = or i64 %c2, %step.add
+  br label %loop
+}
+
 ; Trivially hoist xor.
 define void @xor(i64 %c1, i64 %c2) {
 ; CHECK-LABEL: @xor(