[llvm] [LLVM]: Missed Optimization, treat the disjoint Or as an add (PR #86340)

[Shourya Goel via llvm-commits]

https://github.com/Sh0g0-1758 created https://github.com/llvm/llvm-project/pull/86340

Fixes: #84401

>From 39a4a121b80a1e8d5c1cee738607e7fceec33211 Mon Sep 17 00:00:00 2001
From: Sh0g0-1758 <shouryagoel10000 at gmail.com>
Date: Fri, 22 Mar 2024 21:59:21 +0530
Subject: [PATCH] Initial commit

Added tests
---
 .../InstCombine/InstructionCombining.cpp      | 21 ++++++++++++-------
 .../Transforms/InstCombine/gep-disjoint.ll    | 21 +++++++++++++++++++
 2 files changed, 35 insertions(+), 7 deletions(-)
 create mode 100644 llvm/test/Transforms/InstCombine/gep-disjoint.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 7c40fb4fc86082..507d09077a5cca 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2889,9 +2889,12 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
     // Try to replace ADD + GEP with GEP + GEP.
     Value *Idx1, *Idx2;
     if (match(GEP.getOperand(1),
-              m_OneUse(m_Add(m_Value(Idx1), m_Value(Idx2))))) {
+              m_OneUse(m_AddLike(m_Value(Idx1), m_Value(Idx2))))) {
       //   %idx = add i64 %idx1, %idx2
       //   %gep = getelementptr i32, ptr %ptr, i64 %idx
+      // or
+      //   %idx = or disjoint i64 %idx1, %idx2
+      //   %gep = getelementptr i32, ptr %ptr, i64 %idx
       // as:
       //   %newptr = getelementptr i32, ptr %ptr, i64 %idx1
       //   %newgep = getelementptr i32, ptr %newptr, i64 %idx2
@@ -2901,14 +2904,18 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
                                        Idx2);
     }
     ConstantInt *C;
-    if (match(GEP.getOperand(1), m_OneUse(m_SExtLike(m_OneUse(m_NSWAdd(
+    if (match(GEP.getOperand(1), m_OneUse(m_SExtLike(m_OneUse(m_NSWAddLike(
                                      m_Value(Idx1), m_ConstantInt(C))))))) {
-      // %add = add nsw i32 %idx1, idx2
-      // %sidx = sext i32 %add to i64
-      // %gep = getelementptr i32, ptr %ptr, i64 %sidx
+      //   %add = add nsw i32 %idx1, %idx2
+      //   %sidx = sext i32 %add to i64
+      //   %gep = getelementptr i32, ptr %ptr, i64 %sidx
+      // or
+      //   %disjointOr = or disjoint i32 %idx1, %idx2
+      //   %sidx = sext i32 %disjointOr to i64
+      //   %gep = getelementptr i32, ptr %ptr, i64 %sidx
       // as:
-      // %newptr = getelementptr i32, ptr %ptr, i32 %idx1
-      // %newgep = getelementptr i32, ptr %newptr, i32 idx2
+      //   %newptr = getelementptr i32, ptr %ptr, i32 %idx1
+      //   %newgep = getelementptr i32, ptr %newptr, i32 idx2
       auto *NewPtr = Builder.CreateGEP(
           GEP.getResultElementType(), GEP.getPointerOperand(),
           Builder.CreateSExt(Idx1, GEP.getOperand(1)->getType()));
diff --git a/llvm/test/Transforms/InstCombine/gep-disjoint.ll b/llvm/test/Transforms/InstCombine/gep-disjoint.ll
new file mode 100644
index 00000000000000..fe75519822963b
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/gep-disjoint.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+; test that when offset is know while indexing arrays
+; folding optimization takes place on the addition which gets folded into an OR
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+
+define ptr @test(ptr %arr, i32 %conv) {
+; CHECK-LABEL: define ptr @test(
+; CHECK-SAME: ptr [[ARR:%.*]], i32 [[CONV:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[CONV]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i8, ptr [[ARR]], i64 [[TMP1]]
+; CHECK-NEXT:    [[ARRAYIDX9:%.*]] = getelementptr i8, ptr [[TMP2]], i64 1
+; CHECK-NEXT:    ret ptr [[ARRAYIDX9]]
+;
+  %or7 = or disjoint i32 %conv, 1
+  %idxprom8 = zext nneg i32 %or7 to i64
+  %arrayidx9 = getelementptr i8, ptr %arr, i64 %idxprom8
+  ret ptr %arrayidx9
+}