[AMDGPU] Add mark last scratch load pass

mbrkusanin · mbrkusanin · commit ca6845ae0509 · 2024-01-04T12:54:04.000+01:00
diff --git a/llvm/include/llvm/CodeGen/TargetPassConfig.h b/llvm/include/llvm/CodeGen/TargetPassConfig.h
@@ -423,6 +423,9 @@ class TargetPassConfig : public ImmutablePass {
   /// to physical registers.
   virtual void addPostRewrite() { }
 
+  /// Add passes to be run immediately before Stack Slot Coloring pass.
+  virtual void addPreStackSlotColoring() {}
+
   /// This method may be implemented by targets that want to run passes after
   /// register allocation pass pipeline but before prolog-epilog insertion.
   virtual void addPostRegAlloc() { }
diff --git a/llvm/lib/CodeGen/TargetPassConfig.cpp b/llvm/lib/CodeGen/TargetPassConfig.cpp
@@ -1502,6 +1502,8 @@ void TargetPassConfig::addOptimizedRegAlloc() {
   addPass(&MachineSchedulerID);
 
   if (addRegAssignAndRewriteOptimized()) {
+    addPreStackSlotColoring();
+
     // Perform stack slot coloring and post-ra machine LICM.
     addPass(&StackSlotColoringID);
 
diff --git a/llvm/lib/Target/AMDGPU/AMDGPU.h b/llvm/lib/Target/AMDGPU/AMDGPU.h
@@ -166,6 +166,9 @@ extern char &SILowerI1CopiesID;
 void initializeAMDGPUGlobalISelDivergenceLoweringPass(PassRegistry &);
 extern char &AMDGPUGlobalISelDivergenceLoweringID;
 
+void initializeAMDGPUMarkLastScratchLoadPass(PassRegistry &);
+extern char &AMDGPUMarkLastScratchLoadID;
+
 void initializeSILowerSGPRSpillsPass(PassRegistry &);
 extern char &SILowerSGPRSpillsID;
 
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUMarkLastScratchLoad.cpp b/llvm/lib/Target/AMDGPU/AMDGPUMarkLastScratchLoad.cpp
@@ -0,0 +1,143 @@
+//===-- AMDGPUMarkLastScratchLoad.cpp -------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Mark scratch load/spill instructions which are guaranteed to be the last time
+// this scratch slot is used so it can be evicted from caches.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "llvm/CodeGen/LiveStacks.h"
+#include "llvm/CodeGen/MachineOperand.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-mark-last-scratch-load"
+
+namespace {
+
+class AMDGPUMarkLastScratchLoad : public MachineFunctionPass {
+private:
+  LiveStacks *LS = nullptr;
+  SlotIndexes *SI = nullptr;
+  const SIInstrInfo *SII = nullptr;
+
+public:
+  static char ID;
+
+  AMDGPUMarkLastScratchLoad() : MachineFunctionPass(ID) {
+    initializeAMDGPUMarkLastScratchLoadPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<SlotIndexes>();
+    AU.addRequired<LiveStacks>();
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  StringRef getPassName() const override {
+    return "AMDGPU Mark Last Scratch Load";
+  }
+};
+
+} // end anonymous namespace
+
+bool AMDGPUMarkLastScratchLoad::runOnMachineFunction(MachineFunction &MF) {
+  LLVM_DEBUG({
+    dbgs() << "********** Mark Last Scratch Load **********\n"
+           << "********** Function: " << MF.getName() << '\n';
+  });
+
+  if (skipFunction(MF.getFunction()))
+    return false;
+
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  if (ST.getGeneration() < AMDGPUSubtarget::GFX12)
+    return false;
+
+  LS = &getAnalysis<LiveStacks>();
+  SI = &getAnalysis<SlotIndexes>();
+  SII = ST.getInstrInfo();
+
+  const unsigned NumSlots = LS->getNumIntervals();
+  if (NumSlots == 0) {
+    LLVM_DEBUG(dbgs() << "No live slots, skipping\n");
+    return false;
+  }
+
+  LLVM_DEBUG(dbgs() << LS->getNumIntervals() << " intervals\n");
+
+  bool Changed = false;
+
+  for (auto &[SS, LI] : *LS) {
+    LLVM_DEBUG(dbgs() << "Checking interval: " << LI << "\n");
+
+    for (const LiveRange::Segment &Segment : LI.segments) {
+      LLVM_DEBUG(dbgs() << "  Checking segment: " << Segment << "\n");
+
+      // Ignore segments that run to the end of basic block because in this case
+      // slot is still live at the end of it.
+      if (Segment.end.isBlock())
+        continue;
+
+      const int FrameIndex = Register::stackSlot2Index(LI.reg());
+      MachineInstr *LastLoad = nullptr;
+
+      MachineInstr *MISegmentStart = SI->getInstructionFromIndex(Segment.start);
+      MachineInstr *MISegmentEnd = SI->getInstructionFromIndex(Segment.end);
+      if (!MISegmentEnd) {
+        // FIXME: The start and end can refer to deleted instructions. We should
+        // be able to handle this more gracefully by finding the closest real
+        // instructions.
+        continue;
+      }
+      MachineBasicBlock *BB = MISegmentEnd->getParent();
+
+      // Start iteration backwards from segment end until the start of basic
+      // block or start of segment if it is in the same basic block.
+      auto End = BB->instr_rend();
+      if (MISegmentStart && MISegmentStart->getParent() == BB)
+        End = MISegmentStart->getReverseIterator();
+
+      for (auto MI = MISegmentEnd->getReverseIterator(); MI != End; ++MI) {
+        int LoadFI = 0;
+
+        if (SII->isLoadFromStackSlot(*MI, LoadFI) && LoadFI == FrameIndex) {
+          LastLoad = &*MI;
+          break;
+        }
+      }
+
+      if (LastLoad) {
+        MachineOperand *LastUse =
+            SII->getNamedOperand(*LastLoad, AMDGPU::OpName::last_use);
+        assert(LastUse && "This instruction must have a last_use operand");
+        LastUse->setImm(1);
+        Changed = true;
+        LLVM_DEBUG(dbgs() << "  Found last load: " << *LastLoad;);
+      }
+    }
+  }
+
+  return Changed;
+}
+
+char AMDGPUMarkLastScratchLoad::ID = 0;
+
+char &llvm::AMDGPUMarkLastScratchLoadID = AMDGPUMarkLastScratchLoad::ID;
+
+INITIALIZE_PASS_BEGIN(AMDGPUMarkLastScratchLoad, DEBUG_TYPE,
+                      "AMDGPU Mark last scratch load", false, false)
+INITIALIZE_PASS_DEPENDENCY(SlotIndexes)
+INITIALIZE_PASS_DEPENDENCY(LiveStacks)
+INITIALIZE_PASS_END(AMDGPUMarkLastScratchLoad, DEBUG_TYPE,
+                    "AMDGPU Mark last scratch load", false, false)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -382,6 +382,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAMDGPUTarget() {
   initializeSILowerI1CopiesPass(*PR);
   initializeAMDGPUGlobalISelDivergenceLoweringPass(*PR);
   initializeSILowerWWMCopiesPass(*PR);
+  initializeAMDGPUMarkLastScratchLoadPass(*PR);
   initializeSILowerSGPRSpillsPass(*PR);
   initializeSIFixSGPRCopiesPass(*PR);
   initializeSIFixVGPRCopiesPass(*PR);
@@ -962,6 +963,7 @@ class GCNPassConfig final : public AMDGPUPassConfig {
 
   void addPreRegAlloc() override;
   bool addPreRewrite() override;
+  void addPreStackSlotColoring() override;
   void addPostRegAlloc() override;
   void addPreSched2() override;
   void addPreEmitPass() override;
@@ -1346,6 +1348,10 @@ bool GCNPassConfig::addPreRewrite() {
   return true;
 }
 
+void GCNPassConfig::addPreStackSlotColoring() {
+  addPass(&AMDGPUMarkLastScratchLoadID);
+}
+
 FunctionPass *GCNPassConfig::createSGPRAllocPass(bool Optimized) {
   // Initialize the global default.
   llvm::call_once(InitializeDefaultSGPRRegisterAllocatorFlag,
diff --git a/llvm/lib/Target/AMDGPU/CMakeLists.txt b/llvm/lib/Target/AMDGPU/CMakeLists.txt
@@ -79,6 +79,7 @@ add_llvm_target(AMDGPUCodeGen
   AMDGPUMCInstLower.cpp
   AMDGPUIGroupLP.cpp
   AMDGPUInsertSingleUseVDST.cpp
+  AMDGPUMarkLastScratchLoad.cpp
   AMDGPUMIRFormatter.cpp
   AMDGPUOpenCLEnqueuedBlockLowering.cpp
   AMDGPUPerfHintAnalysis.cpp
diff --git a/llvm/test/CodeGen/AMDGPU/llc-pipeline.ll b/llvm/test/CodeGen/AMDGPU/llc-pipeline.ll
@@ -359,6 +359,7 @@
 ; GCN-O1-NEXT:        SI Lower WWM Copies
 ; GCN-O1-NEXT:        GCN NSA Reassign
 ; GCN-O1-NEXT:        Virtual Register Rewriter
+; GCN-O1-NEXT:        AMDGPU Mark Last Scratch Load
 ; GCN-O1-NEXT:        Stack Slot Coloring
 ; GCN-O1-NEXT:        Machine Copy Propagation Pass
 ; GCN-O1-NEXT:        Machine Loop Invariant Code Motion
@@ -655,6 +656,7 @@
 ; GCN-O1-OPTS-NEXT:        SI Lower WWM Copies
 ; GCN-O1-OPTS-NEXT:        GCN NSA Reassign
 ; GCN-O1-OPTS-NEXT:        Virtual Register Rewriter
+; GCN-O1-OPTS-NEXT:        AMDGPU Mark Last Scratch Load
 ; GCN-O1-OPTS-NEXT:        Stack Slot Coloring
 ; GCN-O1-OPTS-NEXT:        Machine Copy Propagation Pass
 ; GCN-O1-OPTS-NEXT:        Machine Loop Invariant Code Motion
@@ -957,6 +959,7 @@
 ; GCN-O2-NEXT:        SI Lower WWM Copies
 ; GCN-O2-NEXT:        GCN NSA Reassign
 ; GCN-O2-NEXT:        Virtual Register Rewriter
+; GCN-O2-NEXT:        AMDGPU Mark Last Scratch Load
 ; GCN-O2-NEXT:        Stack Slot Coloring
 ; GCN-O2-NEXT:        Machine Copy Propagation Pass
 ; GCN-O2-NEXT:        Machine Loop Invariant Code Motion
@@ -1271,6 +1274,7 @@
 ; GCN-O3-NEXT:        SI Lower WWM Copies
 ; GCN-O3-NEXT:        GCN NSA Reassign
 ; GCN-O3-NEXT:        Virtual Register Rewriter
+; GCN-O3-NEXT:        AMDGPU Mark Last Scratch Load
 ; GCN-O3-NEXT:        Stack Slot Coloring
 ; GCN-O3-NEXT:        Machine Copy Propagation Pass
 ; GCN-O3-NEXT:        Machine Loop Invariant Code Motion
diff --git a/llvm/test/CodeGen/AMDGPU/sgpr-regalloc-flags.ll b/llvm/test/CodeGen/AMDGPU/sgpr-regalloc-flags.ll
@@ -25,6 +25,7 @@
 ; DEFAULT-NEXT: SI Lower WWM Copies
 ; DEFAULT-NEXT: GCN NSA Reassign
 ; DEFAULT-NEXT: Virtual Register Rewriter
+; DEFAULT-NEXT: AMDGPU Mark Last Scratch Load
 ; DEFAULT-NEXT: Stack Slot Coloring
 
 ; O0: Fast Register Allocator
@@ -61,6 +62,7 @@
 ; BASIC-DEFAULT-NEXT: SI Lower WWM Copies
 ; BASIC-DEFAULT-NEXT: GCN NSA Reassign
 ; BASIC-DEFAULT-NEXT: Virtual Register Rewriter
+; BASIC-DEFAULT-NEXT: AMDGPU Mark Last Scratch Load
 ; BASIC-DEFAULT-NEXT: Stack Slot Coloring
 
 
@@ -75,6 +77,7 @@
 ; DEFAULT-BASIC-NEXT: SI Lower WWM Copies
 ; DEFAULT-BASIC-NEXT: GCN NSA Reassign
 ; DEFAULT-BASIC-NEXT: Virtual Register Rewriter
+; DEFAULT-BASIC-NEXT: AMDGPU Mark Last Scratch Load
 ; DEFAULT-BASIC-NEXT: Stack Slot Coloring
 
 
@@ -95,6 +98,7 @@
 ; BASIC-BASIC-NEXT: SI Lower WWM Copies
 ; BASIC-BASIC-NEXT: GCN NSA Reassign
 ; BASIC-BASIC-NEXT: Virtual Register Rewriter
+; BASIC-BASIC-NEXT: AMDGPU Mark Last Scratch Load
 ; BASIC-BASIC-NEXT: Stack Slot Coloring
 
 
diff --git a/llvm/test/CodeGen/AMDGPU/vgpr-mark-last-scratch-load.ll b/llvm/test/CodeGen/AMDGPU/vgpr-mark-last-scratch-load.ll
diff --git a/llvm/test/CodeGen/AMDGPU/vgpr-mark-last-scratch-load.mir b/llvm/test/CodeGen/AMDGPU/vgpr-mark-last-scratch-load.mir
