AMDGPU/SI: Add s_waitcnt at the end of non-void functions

[mareko]

Summary

v2: Make ReturnsVoid private, so that I can another 8 lines of code and

look more productive.

[mareko]

Ping. This adds s_waitcnt at the end of functions which have loads or stores in flight.

[• tstellarAMD]

LGTM.

lib/Target/AMDGPU/SIISelLowering.cpp

SDValue SITargetLowering::LowerReturn(SDValue Chain,
                                       SDLoc DL, SelectionDAG &DAG) const {
   MachineFunction &MF = DAG.getMachineFunction();
   SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
 
   if (Info->getShaderType() == ShaderType::COMPUTE)
     return AMDGPUTargetLowering::LowerReturn(Chain, CallConv, isVarArg, Outs,
                                              OutVals, DL, DAG);
 
+  Info->setIfReturnsVoid(Outs.size() == 0);
+
   SmallVector<ISD::OutputArg, 48> Splits;
   SmallVector<SDValue, 48> SplitVals;
 
   // Split vectors into their elements.
   for (unsigned i = 0, e = Outs.size(); i != e; ++i) {
     const ISD::OutputArg &Out = Outs[i];
 
     if (Out.VT.isVector()) {

lib/Target/AMDGPU/SIInsertWaits.cpp

private:
   /// \brief Different export instruction types seen since last wait.
   unsigned ExpInstrTypesSeen;
 
   /// \brief Type of the last opcode.
   InstType LastOpcodeType;
 
   bool LastInstWritesM0;
 
+  /// \brief Whether the machine function returns void
+  bool ReturnsVoid;
+
   /// \brief Get increment/decrement amount for this instruction.
   Counters getHwCounts(MachineInstr &MI);
 
   /// \brief Is operand relevant for async execution?
   bool isOpRelevant(MachineOperand &Op);
 
   /// \brief Get register interval an operand affects.
   RegInterval getRegInterval(const TargetRegisterClass *RC,
void SIInsertWaits::pushInstruction(MachineBasicBlock &MBB,
   }
 }
 
 bool SIInsertWaits::insertWait(MachineBasicBlock &MBB,
                                MachineBasicBlock::iterator I,
                                const Counters &Required) {
 
   // End of program? No need to wait on anything
-  if (I != MBB.end() && I->getOpcode() == AMDGPU::S_ENDPGM)
+  // A function not returning void needs to wait, because other bytecode will
+  // be appended after it and we don't know what it will be.
+  if (I != MBB.end() && I->getOpcode() == AMDGPU::S_ENDPGM && ReturnsVoid)
     return false;
 
   // Figure out if the async instructions execute in order
   bool Ordered[3];
 
   // VM_CNT is always ordered
   Ordered[0] = true;
 
TRI =
       static_cast<const SIRegisterInfo *>(MF.getSubtarget().getRegisterInfo());
 
   MRI = &MF.getRegInfo();
 
   WaitedOn = ZeroCounts;
   LastIssued = ZeroCounts;
   LastOpcodeType = OTHER;
   LastInstWritesM0 = false;
+  ReturnsVoid = MF.getInfo<SIMachineFunctionInfo>()->returnsVoid();
 
   memset(&UsedRegs, 0, sizeof(UsedRegs));
   memset(&DefinedRegs, 0, sizeof(DefinedRegs));
 
   for (MachineFunction::iterator BI = MF.begin(), BE = MF.end();
        BI != BE; ++BI) {
 
     MachineBasicBlock &MBB = *BI;

lib/Target/AMDGPU/SIMachineFunctionInfo.h

class SIMachineFunctionInfo : public AMDGPUMachineFunction {
   unsigned WorkGroupIDXSystemSGPR;
   unsigned WorkGroupIDYSystemSGPR;
   unsigned WorkGroupIDZSystemSGPR;
   unsigned WorkGroupInfoSystemSGPR;
   unsigned PrivateSegmentWaveByteOffsetSystemSGPR;
 
   // Graphics info.
   unsigned PSInputAddr;
+  bool ReturnsVoid;
 
 public:
   // FIXME: Make private
   unsigned LDSWaveSpillSize;
   unsigned PSInputEna;
   std::map<unsigned, unsigned> LaneVGPRs;
   unsigned ScratchOffsetReg;
   unsigned NumUserSGPRs;
public:
   bool isPSInputAllocated(unsigned Index) const {
     return PSInputAddr & (1 << Index);
   }
 
   void markPSInputAllocated(unsigned Index) {
     PSInputAddr |= 1 << Index;
   }
 
+  bool returnsVoid() const {
+    return ReturnsVoid;
+  }
+
+  void setIfReturnsVoid(bool Value) {
+    ReturnsVoid = Value;
+  }
+
   unsigned getMaximumWorkGroupSize(const MachineFunction &MF) const;
 };
 
 } // End namespace llvm
 
 
 #endif

lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

: AMDGPUMachineFunction(MF),
     GridWorkGroupCountYUserSGPR(AMDGPU::NoRegister),
     GridWorkGroupCountZUserSGPR(AMDGPU::NoRegister),
     WorkGroupIDXSystemSGPR(AMDGPU::NoRegister),
     WorkGroupIDYSystemSGPR(AMDGPU::NoRegister),
     WorkGroupIDZSystemSGPR(AMDGPU::NoRegister),
     WorkGroupInfoSystemSGPR(AMDGPU::NoRegister),
     PrivateSegmentWaveByteOffsetSystemSGPR(AMDGPU::NoRegister),
     PSInputAddr(0),
+    ReturnsVoid(true),
     LDSWaveSpillSize(0),
     PSInputEna(0),
     NumUserSGPRs(0),
     NumSystemSGPRs(0),
     HasSpilledSGPRs(false),
     HasSpilledVGPRs(false),
     PrivateSegmentBuffer(false),
     DispatchPtr(false),