1 files changed, 31 insertions, 22 deletions

diff --git a/lib/Target/AMDGPU/SIInstrInfo.cpp b/lib/Target/AMDGPU/SIInstrInfo.cpp
index 61967605432e7..2c127d787260a 100644
--- a/lib/Target/AMDGPU/SIInstrInfo.cpp
+++ b/lib/Target/AMDGPU/SIInstrInfo.cpp
@@ -3756,36 +3756,45 @@ void SIInstrInfo::moveToVALU(MachineInstr &TopInst) const {
       // FIXME: This isn't safe because the addressing mode doesn't work
       // correctly if vaddr is negative.
       //
-      // FIXME: Handle v_add_u32 and VOP3 form. Also don't rely on immediate
-      // being in src0.
-      //
       // FIXME: Should probably be done somewhere else, maybe SIFoldOperands.
       //
       // See if we can extract an immediate offset by recognizing one of these:
       //   V_ADD_I32_e32 dst, imm, src1
       //   V_ADD_I32_e32 dst, (S_MOV_B32 imm), src1
       // V_ADD will be removed by "Remove dead machine instructions".
-      if (Add && Add->getOpcode() == AMDGPU::V_ADD_I32_e32) {
-        const MachineOperand *Src =
-          getNamedOperand(*Add, AMDGPU::OpName::src0);
-
-        if (Src->isReg()) {
-          auto Mov = MRI.getUniqueVRegDef(Src->getReg());
-          if (Mov && Mov->getOpcode() == AMDGPU::S_MOV_B32)
-            Src = &Mov->getOperand(1);
-        }
+      if (Add &&
+          (Add->getOpcode() == AMDGPU::V_ADD_I32_e32 ||
+           Add->getOpcode() == AMDGPU::V_ADD_U32_e64)) {
+        static const unsigned SrcNames[2] = {
+          AMDGPU::OpName::src0,
+          AMDGPU::OpName::src1,
+        };
+
+        // Find a literal offset in one of source operands.
+        for (int i = 0; i < 2; i++) {
+          const MachineOperand *Src =
+            getNamedOperand(*Add, SrcNames[i]);
+
+          if (Src->isReg()) {
+            auto Mov = MRI.getUniqueVRegDef(Src->getReg());
+            if (Mov && Mov->getOpcode() == AMDGPU::S_MOV_B32)
+              Src = &Mov->getOperand(1);
+          }
+
+          if (Src) {
+            if (Src->isImm())
+              Offset = Src->getImm();
+            else if (Src->isCImm())
+              Offset = Src->getCImm()->getZExtValue();
+          }
+
+          if (Offset && isLegalMUBUFImmOffset(Offset)) {
+            VAddr = getNamedOperand(*Add, SrcNames[!i]);
+            break;
+          }
 
-        if (Src) {
-          if (Src->isImm())
-            Offset = Src->getImm();
-          else if (Src->isCImm())
-            Offset = Src->getCImm()->getZExtValue();
-        }
-
-        if (Offset && isLegalMUBUFImmOffset(Offset))
-          VAddr = getNamedOperand(*Add, AMDGPU::OpName::src1);
-        else
           Offset = 0;
+        }
       }
 
       BuildMI(*MBB, Inst, Inst.getDebugLoc(),