[llvm] [AMDGPU] Explicit instruction prefetch for large entry functions (PR #227995)

[Lukas Sommer via llvm-commits]

================
@@ -0,0 +1,324 @@
+//===- AMDGPUInsertICachePrefetch.cpp - Insert ICache prefetches ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// Insert instruction-cache prefetches for large AMDHSA entry functions.
+//
+// The prefetch instruction is fire-and-forget: it updates no wave wait counter
+// and returns neither a result nor an error. It is therefore safe to insert
+// after the wait-counter and hazard passes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "SIMachineFunctionInfo.h"
+#include "SIProgramInfo.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/TargetParser/Triple.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-insert-icache-prefetch"
+
+static cl::opt<bool>
+    EnableICachePrefetch("amdgpu-icache-prefetch",
+                         cl::desc("Insert ICache prefetch instructions"),
+                         cl::init(true), cl::Hidden);
+
+static cl::opt<unsigned> ICachePrefetchInitialSize(
+    "amdgpu-icache-prefetch-initial-size",
+    cl::desc("Override the initial instruction prefetch size in bytes"),
+    cl::init(0), cl::Hidden);
+
+static cl::opt<unsigned> ICachePrefetchThreshold(
+    "amdgpu-icache-prefetch-threshold",
+    cl::desc("Override the explicit instruction prefetch threshold in bytes"),
+    cl::init(0), cl::Hidden);
+
+namespace {
+
+class AMDGPUInsertICachePrefetch {
+  MachineLoopInfo &MLI;
+  MachinePostDominatorTree &PDT;
+
+  bool isLoopFreeEntryPostDominator(const MachineBasicBlock &MBB,
+                                    const MachineBasicBlock &EntryBB) const {
+    return !MLI.getLoopFor(&MBB) && PDT.dominates(&MBB, &EntryBB);
+  }
+
+public:
+  // These analyses describe the final machine CFG at this late insertion
+  // point. CFG-based placement will use them to select loop-free blocks on
+  // the entry block's post-dominator chain.
+  AMDGPUInsertICachePrefetch(MachineLoopInfo &MLI,
+                             MachinePostDominatorTree &PDT)
+      : MLI(MLI), PDT(PDT) {}
+
+  bool run(MachineFunction &MF);
+};
+
+class AMDGPUInsertICachePrefetchLegacy : public MachineFunctionPass {
+public:
+  static char ID;
+
+  AMDGPUInsertICachePrefetchLegacy() : MachineFunctionPass(ID) {}
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesCFG();
+    AU.addRequired<MachineLoopInfoWrapperPass>();
+    AU.addRequired<MachinePostDominatorTreeWrapperPass>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override {
+    auto &MLI = getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+    auto &PDT =
+        getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();
+    return AMDGPUInsertICachePrefetch(MLI, PDT).run(MF);
+  }
+};
+
+} // end anonymous namespace
+
+struct ICachePrefetchConfig {
+  uint64_t InitialSize;
+  uint64_t Threshold;
+};
+
+static ICachePrefetchConfig getICachePrefetchConfig(const GCNSubtarget &ST) {
+  assert(ST.hasInstPrefSize());
+
+  uint32_t Mask, Shift, Width, CacheLineSize;
+  ST.getInstPrefSizeArgs(Mask, Shift, Width, CacheLineSize);
+  uint64_t DescriptorPrefetchCapacity =
+      ((uint64_t{1} << Width) - 1) * CacheLineSize;
+  uint64_t InitialSize = ICachePrefetchInitialSize.getNumOccurrences()
+                             ? ICachePrefetchInitialSize
+                             : ST.getInitialInstPrefSize();
+  uint64_t Threshold = ICachePrefetchThreshold.getNumOccurrences()
+                           ? ICachePrefetchThreshold
+                           : DescriptorPrefetchCapacity;
+
+  if (InitialSize == 0 || InitialSize % CacheLineSize != 0 ||
+      InitialSize > DescriptorPrefetchCapacity)
+    report_fatal_error(
+        Twine("-amdgpu-icache-prefetch-initial-size must be a non-zero "
+              "multiple of ") +
+        Twine(CacheLineSize) + " bytes not exceeding " +
+        Twine(DescriptorPrefetchCapacity) + " bytes for " + ST.getCPU());
+
+  uint64_t ICacheSize = ST.getInstCacheSize();
+  if (Threshold == 0 || Threshold % CacheLineSize != 0 ||
+      Threshold > ICacheSize)
+    report_fatal_error(
+        Twine("-amdgpu-icache-prefetch-threshold must be a non-zero multiple "
+              "of ") +
+        Twine(CacheLineSize) + " bytes not exceeding " + Twine(ICacheSize) +
+        " bytes for " + ST.getCPU());
+
+  if (InitialSize > Threshold)
+    report_fatal_error(
+        Twine("-amdgpu-icache-prefetch-initial-size must not exceed "
+              "-amdgpu-icache-prefetch-threshold for ") +
+        ST.getCPU());
+
+  return {InitialSize, Threshold};
+}
+
+static MachineBasicBlock::iterator findMBBInsertionPoint(MachineBasicBlock &MBB,
+                                                         const GCNSubtarget &ST,
+                                                         bool IsEntryBlock) {
+  MachineBasicBlock::iterator InsertPt = MBB.begin();
+
+  // Skip past any instructions that must remain at the very beginning:
+  // - Debug values and CFI instructions
+  // - In the entry block only, the gfx1250 initial unclaused-VMEM workaround
+  //   and S_SETREG_IMM32_B32 instructions that set up MODE register bits
+  //   (e.g., REPLAY_MODE bit 25 from SIFrameLowering)
+  // In the entry block, we want the prefetches to come after all initial MODE
+  // setup.
+  bool SkippedInitialUnclausedVmemPrologue =
+      !IsEntryBlock || !ST.hasRequiresInitialUnclausedVmem();
+  while (InsertPt != MBB.end()) {
+    if (InsertPt->isDebugValue() || InsertPt->isCFIInstruction() ||
+        (IsEntryBlock && InsertPt->getOpcode() == AMDGPU::S_SETREG_IMM32_B32)) {
+      ++InsertPt;
+      continue;
+    }
+    if (!SkippedInitialUnclausedVmemPrologue &&
+        InsertPt->getOpcode() == AMDGPU::S_MOV_B64) {
+      auto Vnop = std::next(InsertPt);
+      auto GlobalPrefetch = Vnop == MBB.end() ? MBB.end() : std::next(Vnop);
+      if (Vnop != MBB.end() && Vnop->getOpcode() == AMDGPU::V_NOP_e32 &&
+          GlobalPrefetch != MBB.end() &&
+          GlobalPrefetch->getOpcode() == AMDGPU::GLOBAL_PREFETCH_B8_SADDR) {
+        InsertPt = std::next(GlobalPrefetch);
+        SkippedInitialUnclausedVmemPrologue = true;
+        continue;
+      }
+    }
+    break;
+  }
+  return InsertPt;
+}
+
+bool AMDGPUInsertICachePrefetch::run(MachineFunction &MF) {
+  if (!EnableICachePrefetch)
----------------
sommerlukas wrote:

I've removed the option.

https://github.com/llvm/llvm-project/pull/227995