AMDGPU: Improve frame index folding for pre-allocated objects

[arsenm]

Summary

If the LocalStackSlotAlloction pass has already fixed the offset
for the object during operand folding, check if the offset is an
inline immediate.

There are a couple remaining issues separate from this patch
that need to be corrected before this can be committed.
LocalStackSlotAlloctation decides if it didn't need to put any
frame base registers in that it shouldn't create the pre-allocated
area, so no offsets are known.

Additionally there are code regressions when the base frame index
is folded into an instruction at the expensive of the offset.
e.g. an add fi, offset folds the frame index instead of the single
use offset, instead of materializing the frame index which is re-used
in multiple instructions.

[• tstellarAMD]

LGTM.

[arsenm]

I think this isn't necessarily correct. I don't think the local frame object is at offset 0

lib/Target/AMDGPU/SIFoldOperands.cpp

for (MachineRegisterInfo::use_iterator
       foldOperand(OpToFold, RSUseMI, RSUse.getOperandNo(), FoldList,
                   CopiesToReplace, TII, TRI, MRI);
     }
 
     return;
   }
 
 
-  bool FoldingImm = OpToFold.isImm();
+  bool FoldingImm = OpToFold.isImm() || OpToFold.isFI();
 
   // In order to fold immediates into copies, we need to change the
   // copy to a MOV.
   if (FoldingImm && UseMI->isCopy()) {
     unsigned DestReg = UseMI->getOperand(0).getReg();
     const TargetRegisterClass *DestRC
       = TargetRegisterInfo::isVirtualRegister(DestReg) ?
       MRI.getRegClass(DestReg) :
if (FoldingImm && UseMI->isCopy()) {
 
     // Don't fold into target independent nodes.  Target independent opcodes
     // don't have defined register classes.
     if (UseDesc.isVariadic() ||
         UseDesc.OpInfo[UseOpIdx].RegClass == -1)
       return;
   }
 
-  if (!FoldingImm) {
+  if (!OpToFold.isImm()) {
     tryAddToFoldList(FoldList, UseMI, UseOpIdx, &OpToFold, TII);
 
     // FIXME: We could try to change the instruction from 64-bit to 32-bit
     // to enable more folding opportunites.  The shrink operands pass
     // already does this.
     return;
   }
 

 
 bool SIFoldOperands::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(*MF.getFunction()))
     return false;
 
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
 
   MachineRegisterInfo &MRI = MF.getRegInfo();
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+
   const SIInstrInfo *TII = ST.getInstrInfo();
   const SIRegisterInfo &TRI = TII->getRegisterInfo();
 
   for (MachineFunction::iterator BI = MF.begin(), BE = MF.end();
                                                   BI != BE; ++BI) {
 
     MachineBasicBlock &MBB = *BI;
     MachineBasicBlock::iterator I, Next;
for (I = MBB.begin(); I != MBB.end(); I = Next) {
         // Folding immediates with more than one use will increase program size.
         // FIXME: This will also reduce register usage, which may be better
         // in some cases. A better heuristic is needed.
         for (MachineRegisterInfo::use_iterator
                Use = MRI.use_begin(Dst.getReg()), E = MRI.use_end();
              Use != E; ++Use) {
           MachineInstr *UseMI = Use->getParent();
 
-          if (TII->isInlineConstant(OpToFold, OpSize)) {
+          if ((OpToFold.isImm() && TII->isInlineConstant(OpToFold, OpSize)) ||
+              (OpToFold.isFI() &&
+               TII->isFrameIndexPreAllocInlineImm(MFI, OpToFold.getIndex()))) {
             foldOperand(OpToFold, UseMI, Use.getOperandNo(), FoldList,
                         CopiesToReplace, TII, TRI, MRI);
           } else {
             if (++NumLiteralUses == 1) {
               NonInlineUse = &*Use;
               NonInlineUseOpNo = Use.getOperandNo();
             }
           }

lib/Target/AMDGPU/SIInstrInfo.h

public:
   bool isLiteralConstant(const MachineOperand &MO, unsigned OpSize) const;
 
   // Returns true if this operand could potentially require a 32-bit literal
   // operand, but not necessarily. A FrameIndex for example could resolve to an
   // inline immediate value that will not require an additional 4-bytes; this
   // assumes that it will.
   bool isLiteralConstantLike(const MachineOperand &MO, unsigned OpSize) const;
 
+  /// \returns true if the object \p FI has already had its offset determined to
+  /// be a valid inline immediate value.
+  bool isFrameIndexPreAllocInlineImm(const MachineFrameInfo &MFI, int FI) const;
+
   bool isImmOperandLegal(const MachineInstr &MI, unsigned OpNo,
                          const MachineOperand &MO) const;
 
   /// \brief Return true if this 64-bit VALU instruction has a 32-bit encoding.
   /// This function will return false if you pass it a 32-bit instruction.
   bool hasVALU32BitEncoding(unsigned Opcode) const;
 
   /// \brief Returns true if this operand uses the constant bus.

lib/Target/AMDGPU/SIInstrInfo.cpp

case MachineOperand::MO_Register:
     return Op0.getReg() == Op1.getReg();
   case MachineOperand::MO_Immediate:
     return Op0.getImm() == Op1.getImm();
   default:
     llvm_unreachable("Didn't expect to be comparing these operand types");
   }
 }
 
+bool SIInstrInfo::isFrameIndexPreAllocInlineImm(const MachineFrameInfo &MFI,
+                                                int FI) const {
+  // Check if LocalStackSlotAllocation has already determined the offset for
+  // this frame index.
+  if (!MFI.getUseLocalStackAllocationBlock() || !MFI.isObjectPreAllocated(FI))
+    return false;
+
+  for (int I = 0, E = MFI.getLocalFrameObjectCount(); I != E; ++I) {
+    int ObjFI;
+    int64_t ObjOffset;
+    std::tie(ObjFI, ObjOffset) = MFI.getLocalFrameObjectMap(I);
+    if (ObjFI == FI) {
+      assert(isUInt<32>(ObjOffset));
+      dbgs() << "Found FI#" << FI << " at offset " << ObjOffset << '\n';
+      return isInlineConstant(APInt(32, ObjOffset));
+    }
+  }
+
+  return false;
+}
+
 bool SIInstrInfo::isImmOperandLegal(const MachineInstr &MI, unsigned OpNo,
                                     const MachineOperand &MO) const {
   const MCOperandInfo &OpInfo = get(MI.getOpcode()).OpInfo[OpNo];
 
   assert(MO.isImm() || MO.isTargetIndex() || MO.isFI());
 
   if (OpInfo.OperandType == MCOI::OPERAND_IMMEDIATE)
     return true;
 
   if (OpInfo.RegClass < 0)
     return false;
 
-  unsigned OpSize = RI.getRegClass(OpInfo.RegClass)->getSize();
-  if (isLiteralConstant(MO, OpSize))
+  if (MO.isFI()) {
+    assert(AMDGPU::getRegBitWidth(OpInfo.RegClass) == 32);
+
+    if (RI.opCanUseInlineConstant(OpInfo.OperandType) &&
+        isFrameIndexPreAllocInlineImm(MI.getParent()->getParent()->getFrameInfo(),
+                                      MO.getIndex()))
+      return true;
+
     return RI.opCanUseLiteralConstant(OpInfo.OperandType);
+  }
+
+  unsigned OpSize = AMDGPU::getRegBitWidth(OpInfo.RegClass) / 8;
+  if (isLiteralConstant(MO, OpSize)) {
+    return RI.opCanUseLiteralConstant(OpInfo.OperandType);
+  }
 
   return RI.opCanUseInlineConstant(OpInfo.OperandType);
 }
 
 bool SIInstrInfo::hasVALU32BitEncoding(unsigned Opcode) const {
   int Op32 = AMDGPU::getVOPe32(Opcode);
   if (Op32 == -1)
     return false;