R600: Move CONST_ADDRESS folding into AMDGPUDAGToDAGISel::Select()

This increases the number of opportunites we have for folding.  With the
previous implementation we were unable to fold into any instructions
other than the first when multiple instructions were selected from a
single SDNode.

Reviewed-by: Vincent Lejeune <vljn at ovi.com>

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@186919 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/lib/Target/R600/AMDILISelDAGToDAG.cpp b/lib/Target/R600/AMDILISelDAGToDAG.cpp
index 78036a8..e4fb07d 100644
--- a/lib/Target/R600/AMDILISelDAGToDAG.cpp
+++ b/lib/Target/R600/AMDILISelDAGToDAG.cpp
@@ -50,7 +50,7 @@
 private:
   inline SDValue getSmallIPtrImm(unsigned Imm);
   bool FoldOperand(SDValue &Src, SDValue &Sel, SDValue &Neg, SDValue &Abs,
-                   const R600InstrInfo *TII, std::vector<unsigned> Cst);
+                   const R600InstrInfo *TII);
   bool FoldOperands(unsigned, const R600InstrInfo *, std::vector<SDValue> &);
   bool FoldDotOperands(unsigned, const R600InstrInfo *, std::vector<SDValue> &);
 
@@ -158,12 +158,100 @@
 }
 
 SDNode *AMDGPUDAGToDAGISel::Select(SDNode *N) {
+  const R600InstrInfo *TII =
+                      static_cast<const R600InstrInfo*>(TM.getInstrInfo());
   unsigned int Opc = N->getOpcode();
   if (N->isMachineOpcode()) {
     return NULL;   // Already selected.
   }
   switch (Opc) {
   default: break;
+  case AMDGPUISD::CONST_ADDRESS: {
+    for (SDNode::use_iterator I = N->use_begin(), Next = llvm::next(I);
+                              I != SDNode::use_end(); I = Next) {
+      Next = llvm::next(I);
+      if (!I->isMachineOpcode()) {
+        continue;
+      }
+      unsigned Opcode = I->getMachineOpcode();
+      bool HasDst = TII->getOperandIdx(Opcode, AMDGPU::OpName::dst) > -1;
+      int SrcIdx = I.getOperandNo();
+      int SelIdx;
+      // Unlike MachineInstrs, SDNodes do not have results in their operand
+      // list, so we need to increment the SrcIdx, since
+      // R600InstrInfo::getOperandIdx is based on the MachineInstr indices.
+      if (HasDst) {
+        SrcIdx++;
+      }
+
+      SelIdx = TII->getSelIdx(I->getMachineOpcode(), SrcIdx);
+      if (SelIdx < 0) {
+        continue;
+      }
+
+      SDValue CstOffset;
+      if (N->getValueType(0).isVector() ||
+          !SelectGlobalValueConstantOffset(N->getOperand(0), CstOffset))
+        continue;
+
+      // Gather constants values
+      int SrcIndices[] = {
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src0),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src1),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src2),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src0_X),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src0_Y),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src0_Z),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src0_W),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src1_X),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src1_Y),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src1_Z),
+        TII->getOperandIdx(Opcode, AMDGPU::OpName::src1_W)
+      };
+      std::vector<unsigned> Consts;
+      for (unsigned i = 0; i < sizeof(SrcIndices) / sizeof(int); i++) {
+        int OtherSrcIdx = SrcIndices[i];
+        int OtherSelIdx = TII->getSelIdx(Opcode, OtherSrcIdx);
+        if (OtherSrcIdx < 0 || OtherSelIdx < 0) {
+          continue;
+        }
+        if (HasDst) {
+          OtherSrcIdx--;
+          OtherSelIdx--;
+        }
+        if (RegisterSDNode *Reg =
+                         dyn_cast<RegisterSDNode>(I->getOperand(OtherSrcIdx))) {
+          if (Reg->getReg() == AMDGPU::ALU_CONST) {
+            ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(I->getOperand(OtherSelIdx));
+            Consts.push_back(Cst->getZExtValue());
+          }
+        }
+      }
+
+      ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(CstOffset);
+      Consts.push_back(Cst->getZExtValue());
+      if (!TII->fitsConstReadLimitations(Consts))
+        continue;
+
+      // Convert back to SDNode indices
+      if (HasDst) {
+        SrcIdx--;
+        SelIdx--;
+      }
+      std::vector<SDValue> Ops;
+      for (int i = 0, e = I->getNumOperands(); i != e; ++i) {
+        if (i == SrcIdx) {
+          Ops.push_back(CurDAG->getRegister(AMDGPU::ALU_CONST, MVT::f32));
+        } else if (i == SelIdx) {
+          Ops.push_back(CstOffset);
+        } else {
+          Ops.push_back(I->getOperand(i));
+        }
+      }
+      CurDAG->UpdateNodeOperands(*I, Ops.data(), Ops.size());
+    }
+    break;
+  }
   case ISD::BUILD_VECTOR: {
     const AMDGPUSubtarget &ST = TM.getSubtarget<AMDGPUSubtarget>();
     if (ST.getGeneration() > AMDGPUSubtarget::NORTHERN_ISLANDS) {
@@ -224,7 +312,6 @@
     if (ST.getGeneration() > AMDGPUSubtarget::NORTHERN_ISLANDS) {
       break;
     }
-    const R600InstrInfo *TII = static_cast<const R600InstrInfo*>(TM.getInstrInfo());
 
     uint64_t ImmValue = 0;
     unsigned ImmReg = AMDGPU::ALU_LITERAL_X;
@@ -342,7 +429,7 @@
     if (Result && Result->isMachineOpcode() &&
         !(TII->get(Result->getMachineOpcode()).TSFlags & R600_InstFlag::VECTOR)
         && TII->hasInstrModifiers(Result->getMachineOpcode())) {
-      // Fold FNEG/FABS/CONST_ADDRESS
+      // Fold FNEG/FABS
       // TODO: Isel can generate multiple MachineInst, we need to recursively
       // parse Result
       bool IsModified = false;
@@ -382,24 +469,8 @@
 }
 
 bool AMDGPUDAGToDAGISel::FoldOperand(SDValue &Src, SDValue &Sel, SDValue &Neg,
-                                     SDValue &Abs, const R600InstrInfo *TII,
-                                     std::vector<unsigned> Consts) {
+                                     SDValue &Abs, const R600InstrInfo *TII) {
   switch (Src.getOpcode()) {
-  case AMDGPUISD::CONST_ADDRESS: {
-    SDValue CstOffset;
-    if (Src.getValueType().isVector() ||
-        !SelectGlobalValueConstantOffset(Src.getOperand(0), CstOffset))
-      return false;
-
-    ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(CstOffset);
-    Consts.push_back(Cst->getZExtValue());
-    if (!TII->fitsConstReadLimitations(Consts))
-      return false;
-
-    Src = CurDAG->getRegister(AMDGPU::ALU_CONST, MVT::f32);
-    Sel = CstOffset;
-    return true;
-    }
   case ISD::FNEG:
     Src = Src.getOperand(0);
     Neg = CurDAG->getTargetConstant(1, MVT::i32);
@@ -441,19 +512,6 @@
     -1
   };
 
-  // Gather constants values
-  std::vector<unsigned> Consts;
-  for (unsigned j = 0; j < 3; j++) {
-    int SrcIdx = OperandIdx[j];
-    if (SrcIdx < 0)
-      break;
-    if (RegisterSDNode *Reg = dyn_cast<RegisterSDNode>(Ops[SrcIdx - 1])) {
-      if (Reg->getReg() == AMDGPU::ALU_CONST) {
-        ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(Ops[SelIdx[j] - 1]);
-        Consts.push_back(Cst->getZExtValue());
-      }
-    }
-  }
 
   for (unsigned i = 0; i < 3; i++) {
     if (OperandIdx[i] < 0)
@@ -463,7 +521,7 @@
     SDValue &Neg = Ops[NegIdx[i] - 1];
     SDValue FakeAbs;
     SDValue &Abs = (AbsIdx[i] > -1) ? Ops[AbsIdx[i] - 1] : FakeAbs;
-    if (FoldOperand(Src, Sel, Neg, Abs, TII, Consts))
+    if (FoldOperand(Src, Sel, Neg, Abs, TII))
       return true;
   }
   return false;
@@ -512,20 +570,6 @@
     TII->getOperandIdx(Opcode, AMDGPU::OpName::src1_abs_W)
   };
 
-  // Gather constants values
-  std::vector<unsigned> Consts;
-  for (unsigned j = 0; j < 8; j++) {
-    int SrcIdx = OperandIdx[j];
-    if (SrcIdx < 0)
-      break;
-    if (RegisterSDNode *Reg = dyn_cast<RegisterSDNode>(Ops[SrcIdx - 1])) {
-      if (Reg->getReg() == AMDGPU::ALU_CONST) {
-        ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(Ops[SelIdx[j] - 1]);
-        Consts.push_back(Cst->getZExtValue());
-      }
-    }
-  }
-
   for (unsigned i = 0; i < 8; i++) {
     if (OperandIdx[i] < 0)
       return false;
@@ -533,7 +577,7 @@
     SDValue &Sel = Ops[SelIdx[i] - 1];
     SDValue &Neg = Ops[NegIdx[i] - 1];
     SDValue &Abs = Ops[AbsIdx[i] - 1];
-    if (FoldOperand(Src, Sel, Neg, Abs, TII, Consts))
+    if (FoldOperand(Src, Sel, Neg, Abs, TII))
       return true;
   }
   return false;