[llvm] [SeperateConstOffsetFromGEP] Handle `or disjoint` flags (PR #76997)

[Krzysztof Drewniak via llvm-commits]

https://github.com/krzysz00 created https://github.com/llvm/llvm-project/pull/76997

This commit extends separate-const-offset-from-gep to look at the newly-added `disjoint` flag on `or` instructions so as to preserve addditional opportunities for optimization.

As with other `or disjoint`-handling commits, this does not remove the existing check for the or's operands having no bits in common because `disjoint` is currently not inferred.

The tests were pre-committed in #76972.

>From 5cc46862df42e7d01a2d45ccc18f221744af0b93 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Thu, 4 Jan 2024 20:20:54 +0000
Subject: [PATCH] [SeperateConstOffsetFromGEP] Handle `or disjoint` flags

This commit extends separate-const-offset-from-gep to look at the
newly-added `disjoint` flag on `or` instructions so as to preserve
addditional opportunities for optimization.

As with other `or disjoint`-handling commits, this does not remove the
existing check for the or's operands having no bits in common because
`disjoint` is currently not inferred.

The tests were pre-committed in #76972.
---
 .../Scalar/SeparateConstOffsetFromGEP.cpp           | 13 ++++++++-----
 .../split-gep-or-as-add.ll                          |  8 +++++---
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index 225dd454068c84..9bb42f7be8d70a 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -174,6 +174,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
@@ -519,12 +520,14 @@ bool ConstantOffsetExtractor::CanTraceInto(bool SignExtended,
   }
 
   Value *LHS = BO->getOperand(0), *RHS = BO->getOperand(1);
-  // Do not trace into "or" unless it is equivalent to "add". If LHS and RHS
-  // don't have common bits, (LHS | RHS) is equivalent to (LHS + RHS).
-  // FIXME: this does not appear to be covered by any tests
-  //        (with x86/aarch64 backends at least)
+  // Do not trace into "or" unless it is equivalent to "add".
+  // This is the case if the "or"'s disjoint flag is set, or (because we
+  // currently don't infer the disjoint flags) if its left and right operands
+  // have nothing in commen.
   if (BO->getOpcode() == Instruction::Or &&
-      !haveNoCommonBitsSet(LHS, RHS, SimplifyQuery(DL, DT, /*AC*/ nullptr, BO)))
+      !(cast<PossiblyDisjointInst>(BO)->isDisjoint() ||
+        haveNoCommonBitsSet(LHS, RHS,
+                            SimplifyQuery(DL, DT, /*AC*/ nullptr, BO))))
     return false;
 
   // FIXME: We don't currently support constants from the RHS of subs,
diff --git a/llvm/test/Transforms/SeparateConstOffsetFromGEP/split-gep-or-as-add.ll b/llvm/test/Transforms/SeparateConstOffsetFromGEP/split-gep-or-as-add.ll
index 45154f5a68f92c..5041fed12f5963 100644
--- a/llvm/test/Transforms/SeparateConstOffsetFromGEP/split-gep-or-as-add.ll
+++ b/llvm/test/Transforms/SeparateConstOffsetFromGEP/split-gep-or-as-add.ll
@@ -46,9 +46,11 @@ define void @testDisjointOrSplits(ptr %p) {
 ; CHECK-LABEL: define void @testDisjointOrSplits(
 ; CHECK-SAME: ptr [[P:%.*]]) {
 ; CHECK-NEXT:    [[VAR:%.*]] = tail call i64 @foo()
-; CHECK-NEXT:    [[OFF:%.*]] = or disjoint i64 [[VAR]], 10
-; CHECK-NEXT:    [[Q:%.*]] = getelementptr i8, ptr [[P]], i64 [[OFF]]
-; CHECK-NEXT:    store i8 0, ptr [[Q]], align 1
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[P]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[TMP1]], [[VAR]]
+; CHECK-NEXT:    [[TMP3:%.*]] = add i64 [[TMP2]], 10
+; CHECK-NEXT:    [[TMP4:%.*]] = inttoptr i64 [[TMP3]] to ptr
+; CHECK-NEXT:    store i8 0, ptr [[TMP4]], align 1
 ; CHECK-NEXT:    ret void
 ;
   %var = tail call i64 @foo()