[AMDGPU] Group multiple single use producers under one single use instruction. (#90713)

Previously each single use producer would be marked with a
"S_SINGLEUSE_VDST 1" instruction. This patch adds support for
larger immediates that encode multiple single use producers into
one S_SINGLEUSE_VDST instruction.

Author: ScottEgerton

llvm/lib/Target/AMDGPU/AMDGPUInsertSingleUseVDST.cpp

@@ -16,10 +16,11 @@
 
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
-#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "SIInstrInfo.h"
+#include "SIRegisterInfo.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -28,10 +29,11 @@
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/Register.h"
-#include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/MC/MCRegister.h"
+#include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/Pass.h"
+#include <array>
 
 using namespace llvm;
 
@@ -41,17 +43,110 @@ namespace {
 class AMDGPUInsertSingleUseVDST : public MachineFunctionPass {
 private:
   const SIInstrInfo *SII;
+  class SingleUseInstruction {
+  private:
+    static const unsigned MaxSkipRange = 0b111;
+    static const unsigned MaxNumberOfSkipRegions = 2;
+
+    unsigned LastEncodedPositionEnd;
+    MachineInstr *ProducerInstr;
+
+    std::array<unsigned, MaxNumberOfSkipRegions + 1> SingleUseRegions;
+    SmallVector<unsigned, MaxNumberOfSkipRegions> SkipRegions;
+
+    // Adds a skip region into the instruction.
+    void skip(const unsigned ProducerPosition) {
+      while (LastEncodedPositionEnd + MaxSkipRange < ProducerPosition) {
+        SkipRegions.push_back(MaxSkipRange);
+        LastEncodedPositionEnd += MaxSkipRange;
+      }
+      SkipRegions.push_back(ProducerPosition - LastEncodedPositionEnd);
+      LastEncodedPositionEnd = ProducerPosition;
+    }
+
+    bool currentRegionHasSpace() {
+      const auto Region = SkipRegions.size();
+      // The first region has an extra bit of encoding space.
+      return SingleUseRegions[Region] <
+             ((Region == MaxNumberOfSkipRegions) ? 0b1111 : 0b111);
+    }
+
+    unsigned encodeImm() {
+      // Handle the first Single Use Region separately as it has an extra bit
+      // of encoding space.
+      unsigned Imm = SingleUseRegions[SkipRegions.size()];
+      unsigned ShiftAmount = 4;
+      for (unsigned i = SkipRegions.size(); i > 0; i--) {
+        Imm |= SkipRegions[i - 1] << ShiftAmount;
+        ShiftAmount += 3;
+        Imm |= SingleUseRegions[i - 1] << ShiftAmount;
+        ShiftAmount += 3;
+      }
+      return Imm;
+    }
+
+  public:
+    SingleUseInstruction(const unsigned ProducerPosition,
+                         MachineInstr *Producer)
+        : LastEncodedPositionEnd(ProducerPosition + 1), ProducerInstr(Producer),
+          SingleUseRegions({1, 0, 0}) {}
+
+    // Returns false if adding a new single use producer failed. This happens
+    // because it could not be encoded, either because there is no room to
+    // encode another single use producer region or that this single use
+    // producer is too far away to encode the amount of instructions to skip.
+    bool tryAddProducer(const unsigned ProducerPosition, MachineInstr *MI) {
+      // Producer is too far away to encode into this instruction or another
+      // skip region is needed and SkipRegions.size() = 2 so there's no room for
+      // another skip region, therefore a new instruction is needed.
+      if (LastEncodedPositionEnd +
+              (MaxSkipRange * (MaxNumberOfSkipRegions - SkipRegions.size())) <
+          ProducerPosition)
+        return false;
+
+      // If a skip region is needed.
+      if (LastEncodedPositionEnd != ProducerPosition ||
+          !currentRegionHasSpace()) {
+        // If the current region is out of space therefore a skip region would
+        // be needed, but there is no room for another skip region.
+        if (SkipRegions.size() == MaxNumberOfSkipRegions)
+          return false;
+        skip(ProducerPosition);
+      }
+
+      SingleUseRegions[SkipRegions.size()]++;
+      LastEncodedPositionEnd = ProducerPosition + 1;
+      ProducerInstr = MI;
+      return true;
+    }
+
+    auto emit(const SIInstrInfo *SII) {
+      return BuildMI(*ProducerInstr->getParent(), ProducerInstr, DebugLoc(),
+                     SII->get(AMDGPU::S_SINGLEUSE_VDST))
+          .addImm(encodeImm());
+    }
+  };
 
 public:
   static char ID;
 
   AMDGPUInsertSingleUseVDST() : MachineFunctionPass(ID) {}
 
-  void emitSingleUseVDST(MachineInstr &MI) const {
-    // Mark the following instruction as a single-use producer:
-    //   s_singleuse_vdst { supr0: 1 }
-    BuildMI(*MI.getParent(), MI, DebugLoc(), SII->get(AMDGPU::S_SINGLEUSE_VDST))
-        .addImm(0x1);
+  void insertSingleUseInstructions(
+      ArrayRef<std::pair<unsigned, MachineInstr *>> SingleUseProducers) const {
+    SmallVector<SingleUseInstruction> Instructions;
+
+    for (auto &[Position, MI] : SingleUseProducers) {
+      // Encode this position into the last single use instruction if possible.
+      if (Instructions.empty() ||
+          !Instructions.back().tryAddProducer(Position, MI)) {
+        // If not, add a new instruction.
+        Instructions.push_back(SingleUseInstruction(Position, MI));
+      }
+    }
+
+    for (auto &Instruction : Instructions)
+      Instruction.emit(SII);
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override {
@@ -78,6 +173,10 @@ class AMDGPUInsertSingleUseVDST : public MachineFunctionPass {
         }
       }
 
+      SmallVector<std::pair<unsigned, MachineInstr *>>
+          SingleUseProducerPositions;
+
+      unsigned VALUInstrCount = 0;
       for (MachineInstr &MI : reverse(MBB.instrs())) {
         // All registers in all operands need to be single use for an
         // instruction to be marked as a single use producer.
@@ -119,13 +218,16 @@ class AMDGPUInsertSingleUseVDST : public MachineFunctionPass {
           for (auto &UsedReg : RegisterUseCount)
             UsedReg.second = 2;
         }
-        if (AllProducerOperandsAreSingleUse && SIInstrInfo::isVALU(MI)) {
-          // TODO: Replace with candidate logging for instruction grouping
-          // later.
-          emitSingleUseVDST(MI);
+
+        if (!SIInstrInfo::isVALU(MI))
+          continue;
+        if (AllProducerOperandsAreSingleUse) {
+          SingleUseProducerPositions.push_back({VALUInstrCount, &MI});
           InstructionEmitted = true;
         }
+        VALUInstrCount++;
       }
+      insertSingleUseInstructions(SingleUseProducerPositions);
     }
     return InstructionEmitted;
   }