[LICM] Preserve Disjoint flag on OR when hoisting.

[fhahn]

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

Fixes #139625.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

Changes

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

Fixes #139625.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/LICM.cpp (+6)
• (modified) llvm/test/Transforms/LICM/hoist-binop.ll (+63)

diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 62ef40c686874..cbdb8bf661a23 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -2877,6 +2877,12 @@ 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 (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 ea7d96c07d5ff..000b1fc1c5ce6 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(

[fhahn]

I am wondering if there's any API we could come up with that would make it easier to avoid missing cases to update when we add new flags

[preames]

LGTM

[preames]

This block of code seems like something we should have a helper for. As a follow up, maybe factor something out of other places we do reassociation?

[preames]

Actually, unless I'm missing something, Reassociate.cpp has the same problem of dropping disjoint flags. Note that there is the existing OverflowTracking helper, which might serve as a starting place for the API.

[efriedma-quic]

This cast is fragile; it only works under certain assumptions about what the IRBuilder will not fold.

Maybe use IRBuilder<NoFolder>, so you can just cast<>?

[fhahn]

Currently this uses the default ConstantFolder, which should be fine, as we either get a constant or the required. OR instruction. This is similar to the code for the other cases above.

Is there any particlar case you are thinking where this would be too fragile? I can imagine that we would have to be more careful if the folder could fold the op to an existing other IR value, in which case we may end up setting the disjoint flags on a different OR than we intended. But I don't think that can be the case with the ConstantFolder.

[efriedma-quic]

It's not actually wrong with ConstantFolder, I guess, but the interaction is pretty subtle. I've seen similar-looking code cause real issues.

Maybe it's enough to spell out that it's using ConstantFolder, and why it's getting used?

[fhahn]

Thanks, added the same comment as we have above for the ADD/MUL case.