[AMDGPU] Handle OR as operand of raw load/store

Summary:
Use isBaseWithConstantOffset() which handles OR as an operand
to llvm.amdgcn.raw.buffer.load and llvm.amdgcn.raw.buffer.store.

Change-Id: Ifefb9dc5ded8710d333df07ab1900b230e33539a

Reviewers: nhaehnle, mareko, arsenm

Reviewed By: arsenm

Subscribers: arsenm, kzhuravl, jvesely, wdng, nhaehnle, yaxunl, dstuttard, tpr, t-tye, llvm-commits

Differential Revision: https://reviews.llvm.org/D55999

llvm-svn: 350208

1 files changed, 6 insertions, 4 deletions

diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index 12d0bc528ca..894806b9faa 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -6005,11 +6005,13 @@ std::pair<SDValue, SDValue> SITargetLowering::splitBufferOffsets(
   const unsigned MaxImm = 4095;
   SDValue N0 = Offset;
   ConstantSDNode *C1 = nullptr;
-  if (N0.getOpcode() == ISD::ADD) {
-    if ((C1 = dyn_cast<ConstantSDNode>(N0.getOperand(1))))
-      N0 = N0.getOperand(0);
-  } else if ((C1 = dyn_cast<ConstantSDNode>(N0)))
+
+  if ((C1 = dyn_cast<ConstantSDNode>(N0)))
     N0 = SDValue();
+  else if (DAG.isBaseWithConstantOffset(N0)) {
+    C1 = cast<ConstantSDNode>(N0.getOperand(1));
+    N0 = N0.getOperand(0);
+  }
 
   if (C1) {
     unsigned ImmOffset = C1->getZExtValue();