Implement compiler option -mamdgpu-expand-waitcnt-profiling to expand waitcnt instruction

Author: PankajDwivedi-25

clang/include/clang/Driver/Options.td

@@ -5497,7 +5497,10 @@ defm wavefrontsize64 : SimpleMFlag<"wavefrontsize64",
   " mode (AMDGPU only)">;
 defm amdgpu_precise_memory_op
     : SimpleMFlag<"amdgpu-precise-memory-op", "Enable", "Disable",
-                  " precise memory mode (AMDGPU only)">;
+                  " precise memory mode (AMDGPU only)", m_amdgpu_Features_Group>;
+defm amdgpu_expand_waitcnt_profiling
+    : SimpleMFlag<"amdgpu-expand-waitcnt-profiling", "Enable", "Disable",
+                  " waitcnt expansion for profiling (AMDGPU only)", m_amdgpu_Features_Group>;
 
 def munsafe_fp_atomics : Flag<["-"], "munsafe-fp-atomics">,
   Visibility<[ClangOption, CC1Option, FlangOption, FC1Option]>, Alias<fatomic_ignore_denormal_mode>;

clang/lib/Driver/ToolChains/AMDGPU.cpp

@@ -700,6 +700,10 @@ void amdgpu::getAMDGPUTargetFeatures(const Driver &D,
                    options::OPT_mno_amdgpu_precise_memory_op, false))
     Features.push_back("+precise-memory");
 
+  if (Args.hasFlag(options::OPT_mamdgpu_expand_waitcnt_profiling,
+                   options::OPT_mno_amdgpu_expand_waitcnt_profiling, false))
+    Features.push_back("+expand-waitcnt-profiling");
+
   handleTargetFeaturesGroup(D, Triple, Args, Features,
                             options::OPT_m_amdgpu_Features_Group);
 }

clang/test/Driver/amdgpu-features.c

@@ -38,3 +38,9 @@
 
 // RUN: %clang -### -target amdgcn -mcpu=gfx1010 -mno-amdgpu-precise-memory-op %s 2>&1 | FileCheck --check-prefix=NO-PREC-MEM %s
 // NO-PREC-MEM-NOT: {{".*precise-memory"}}
+
+// RUN: %clang -### -target amdgcn -mcpu=gfx900 -mamdgpu-expand-waitcnt-profiling %s 2>&1 | FileCheck --check-prefix=EXPAND-WAITCNT %s
+// EXPAND-WAITCNT: "-target-feature" "+expand-waitcnt-profiling"
+
+// RUN: %clang -### -target amdgcn -mcpu=gfx900 -mno-amdgpu-expand-waitcnt-profiling %s 2>&1 | FileCheck --check-prefix=NO-EXPAND-WAITCNT %s
+// NO-EXPAND-WAITCNT-NOT: "{{[+]}}expand-waitcnt-profiling"

llvm/lib/Target/AMDGPU/AMDGPU.td

@@ -223,6 +223,10 @@ def FeaturePreciseMemory
     : SubtargetFeature<"precise-memory", "EnablePreciseMemory",
                        "true", "Enable precise memory mode">;
 
+def FeatureExpandWaitcntProfiling
+    : SubtargetFeature<"expand-waitcnt-profiling", "EnableExpandWaitcntProfiling",
+                       "true", "Expand waitcnt instructions for profiling">;
+
 def FeatureSGPRInitBug : SubtargetFeature<"sgpr-init-bug",
   "SGPRInitBug",
   "true",

llvm/lib/Target/AMDGPU/GCNSubtarget.h

@@ -90,6 +90,7 @@ class GCNSubtarget final : public AMDGPUGenSubtargetInfo,
   bool EnableCuMode = false;
   bool TrapHandler = false;
   bool EnablePreciseMemory = false;
+  bool EnableExpandWaitcntProfiling = false;
 
   // Used as options.
   bool EnableLoadStoreOpt = false;
@@ -674,6 +675,10 @@ class GCNSubtarget final : public AMDGPUGenSubtargetInfo,
 
   bool isPreciseMemoryEnabled() const { return EnablePreciseMemory; }
 
+  bool isExpandWaitcntProfilingEnabled() const {
+    return EnableExpandWaitcntProfiling;
+  }
+
   bool hasFlatAddressSpace() const {
     return FlatAddressSpace;
   }

llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp

@@ -494,6 +494,16 @@ class SIInsertWaitcnts {
   bool isVMEMOrFlatVMEM(const MachineInstr &MI) const;
   bool run(MachineFunction &MF);
 
+  // Methods for expanding waitcnt instructions for profiling
+  bool expandWaitcntsForProfiling(MachineFunction &MF);
+  bool expandSingleWaitcnt(MachineInstr &MI, MachineBasicBlock &MBB);
+  bool expandSingleCounterWait(MachineInstr &MI, MachineBasicBlock &MBB,
+                               InstCounterType CT);
+  bool expandCounterSequence(MachineBasicBlock &MBB,
+                             MachineBasicBlock::iterator InsertPos,
+                             InstCounterType CT, unsigned CountValue,
+                             DebugLoc DL);
+
   void setForceEmitWaitcnt() {
 // For non-debug builds, ForceEmitWaitcnt has been initialized to false;
 // For debug builds, get the debug counter info and adjust if need be
@@ -2725,6 +2735,156 @@ SIInsertWaitcntsPass::run(MachineFunction &MF,
       .preserve<AAManager>();
 }
 
+/// Expand waitcnt instructions for profiling by inserting a sequence of
+/// decreasing counter values. This helps identify which specific memory
+/// operation is a bottleneck during PC sampling.
+bool SIInsertWaitcnts::expandWaitcntsForProfiling(MachineFunction &MF) {
+  if (!ST->isExpandWaitcntProfilingEnabled())
+    return false;
+
+  bool Modified = false;
+
+  // Iterate through all basic blocks
+  for (MachineBasicBlock &MBB : MF) {
+    for (auto I = MBB.begin(), E = MBB.end(); I != E;) {
+      MachineInstr &MI = *I;
+      ++I; // Advance iterator before potential expansion
+
+      if (ST->hasExtendedWaitCounts()) {
+        // GFX12+: Handle separate wait instructions
+        if (auto CT = counterTypeForInstr(MI.getOpcode())) {
+          Modified |= expandSingleCounterWait(MI, MBB, *CT);
+        }
+      } else {
+        // Pre-GFX12: Handle combined S_WAITCNT
+        if (MI.getOpcode() == AMDGPU::S_WAITCNT) {
+          Modified |= expandSingleWaitcnt(MI, MBB);
+        }
+      }
+    }
+  }
+
+  return Modified;
+}
+
+/// Expand a single S_WAITCNT instruction (pre-GFX12)
+bool SIInsertWaitcnts::expandSingleWaitcnt(MachineInstr &MI,
+                                           MachineBasicBlock &MBB) {
+  assert(MI.getOpcode() == AMDGPU::S_WAITCNT);
+
+  // Decode the waitcnt immediate
+  unsigned Imm = MI.getOperand(0).getImm();
+  AMDGPU::IsaVersion IV = AMDGPU::getIsaVersion(ST->getCPU());
+  AMDGPU::Waitcnt Wait = AMDGPU::decodeWaitcnt(IV, Imm);
+
+  // Insert expanded waitcnts BEFORE the original instruction
+  auto InsertPos = MI.getIterator();
+  DebugLoc DL = MI.getDebugLoc();
+
+  bool Modified = false;
+
+  // Expand each counter independently
+  // For independent counters (Case 2 from requirements):
+  // vmcnt and lgkmcnt can be separated
+  Modified |= expandCounterSequence(MBB, InsertPos, LOAD_CNT, Wait.LoadCnt, DL);
+  Modified |= expandCounterSequence(MBB, InsertPos, DS_CNT, Wait.DsCnt, DL);
+  Modified |= expandCounterSequence(MBB, InsertPos, EXP_CNT, Wait.ExpCnt, DL);
+  Modified |=
+      expandCounterSequence(MBB, InsertPos, STORE_CNT, Wait.StoreCnt, DL);
+
+  // If we expanded anything, remove the original waitcnt
+  if (Modified) {
+    MI.eraseFromParent();
+  }
+
+  return Modified;
+}
+
+/// Expand a single counter wait instruction (GFX12+)
+bool SIInsertWaitcnts::expandSingleCounterWait(MachineInstr &MI,
+                                               MachineBasicBlock &MBB,
+                                               InstCounterType CT) {
+  // Get the counter value from the instruction
+  unsigned CountValue = MI.getOperand(0).getImm();
+
+  // Insert expanded waitcnts BEFORE the original instruction
+  auto InsertPos = MI.getIterator();
+  DebugLoc DL = MI.getDebugLoc();
+
+  bool Modified = expandCounterSequence(MBB, InsertPos, CT, CountValue, DL);
+
+  // If we expanded, remove the original instruction
+  if (Modified) {
+    MI.eraseFromParent();
+  }
+
+  return Modified;
+}
+
+/// Insert a sequence of wait instructions with decreasing counter values
+bool SIInsertWaitcnts::expandCounterSequence(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertPos,
+    InstCounterType CT, unsigned CountValue, DebugLoc DL) {
+  // Skip if counter is already at zero, not active, or at max (wait not needed)
+  if (CountValue == 0 || CountValue == ~0u)
+    return false;
+
+  unsigned MaxCount = getWaitCountMax(CT);
+  if (CountValue >= MaxCount)
+    return false;
+
+  bool Modified = false;
+
+  // Generate decreasing sequence: CountValue-1, CountValue-2, ..., 1, 0
+  // We start from CountValue-1 because the original waitcnt already handles
+  // CountValue
+  for (int i = CountValue - 1; i >= 0; --i) {
+    if (ST->hasExtendedWaitCounts()) {
+      // GFX12+: Use separate wait instructions
+      unsigned Opcode = instrsForExtendedCounterTypes[CT];
+      BuildMI(MBB, InsertPos, DL, TII->get(Opcode)).addImm(i);
+    } else {
+      // Pre-GFX12: Use combined S_WAITCNT with only this counter set
+      AMDGPU::Waitcnt Wait;
+      switch (CT) {
+      case LOAD_CNT:
+        Wait.LoadCnt = i;
+        break;
+      case DS_CNT:
+        Wait.DsCnt = i;
+        break;
+      case EXP_CNT:
+        Wait.ExpCnt = i;
+        break;
+      case STORE_CNT:
+        Wait.StoreCnt = i;
+        break;
+      case SAMPLE_CNT:
+        Wait.SampleCnt = i;
+        break;
+      case BVH_CNT:
+        Wait.BvhCnt = i;
+        break;
+      case KM_CNT:
+        Wait.KmCnt = i;
+        break;
+      case X_CNT:
+        Wait.XCnt = i;
+        break;
+      default:
+        break;
+      }
+
+      AMDGPU::IsaVersion IV = AMDGPU::getIsaVersion(ST->getCPU());
+      unsigned Enc = AMDGPU::encodeWaitcnt(IV, Wait);
+      BuildMI(MBB, InsertPos, DL, TII->get(AMDGPU::S_WAITCNT)).addImm(Enc);
+    }
+    Modified = true;
+  }
+
+  return Modified;
+}
+
 bool SIInsertWaitcnts::run(MachineFunction &MF) {
   ST = &MF.getSubtarget<GCNSubtarget>();
   TII = ST->getInstrInfo();
@@ -2963,5 +3123,10 @@ bool SIInsertWaitcnts::run(MachineFunction &MF) {
   PreheadersToFlush.clear();
   SLoadAddresses.clear();
 
+  // Expand waitcnts for profiling if requested
+  if (ST->isExpandWaitcntProfilingEnabled()) {
+    Modified |= expandWaitcntsForProfiling(MF);
+  }
+
   return Modified;
 }