[AMDGPU] Add mark last scratch load pass

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;
 

llvm/lib/Target/AMDGPU/AMDGPUMarkLastScratchLoad.cpp

@@ -0,0 +1,142 @@
+//===-- 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.
+//
+// TODO: Handle general stack accesses not just spilling.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "llvm/CodeGen/LiveIntervals.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;
+  LiveIntervals *LIS = 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<LiveIntervals>();
+    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) {
+  if (skipFunction(MF.getFunction()))
+    return false;
+
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  if (ST.getGeneration() < AMDGPUSubtarget::GFX12)
+    return false;
+
+  LS = &getAnalysis<LiveStacks>();
+  LIS = &getAnalysis<LiveIntervals>();
+  SI = &getAnalysis<SlotIndexes>();
+  SII = ST.getInstrInfo();
+  SlotIndexes &Slots = *LIS->getSlotIndexes();
+
+  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) {
+    for (const LiveRange::Segment &Segment : LI.segments) {
+
+      // 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 *MISegmentEnd = SI->getInstructionFromIndex(Segment.end);
+
+      // If there is no instruction at this slot because it was deleted take the
+      // instruction from the next slot.
+      if (!MISegmentEnd) {
+        SlotIndex NextSlot = Slots.getNextNonNullIndex(Segment.end);
+        MISegmentEnd = SI->getInstructionFromIndex(NextSlot);
+      }
+
+      MachineInstr *MISegmentStart = SI->getInstructionFromIndex(Segment.start);
+      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->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 && !LastLoad->memoperands_empty()) {
+        MachineMemOperand *MMO = *LastLoad->memoperands_begin();
+        MMO->setFlags(MOLastUse);
+        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)

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);
@@ -1424,6 +1425,8 @@ bool GCNPassConfig::addRegAssignAndRewriteOptimized() {
   addPreRewrite();
   addPass(&VirtRegRewriterID);
 
+  addPass(&AMDGPUMarkLastScratchLoadID);
+
   return true;
 }
 

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

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

@@ -8756,6 +8756,7 @@ SIInstrInfo::getSerializableMachineMemOperandTargetFlags() const {
   static const std::pair<MachineMemOperand::Flags, const char *> TargetFlags[] =
       {
           {MONoClobber, "amdgpu-noclobber"},
+          {MOLastUse, "amdgpu-last-use"},
       };
 
   return ArrayRef(TargetFlags);

llvm/lib/Target/AMDGPU/SIInstrInfo.h

@@ -41,6 +41,10 @@ class ScheduleHazardRecognizer;
 static const MachineMemOperand::Flags MONoClobber =
     MachineMemOperand::MOTargetFlag1;
 
+/// Mark the MMO of a load as the last use.
+static const MachineMemOperand::Flags MOLastUse =
+    MachineMemOperand::MOTargetFlag2;
+
 /// Utility to store machine instructions worklist.
 struct SIInstrWorklist {
   SIInstrWorklist() = default;

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

@@ -1657,8 +1657,12 @@ void SIRegisterInfo::buildSpillLoadStore(
     } else {
       MIB.addReg(SOffset, SOffsetRegState);
     }
-    MIB.addImm(Offset + RegOffset)
-       .addImm(0); // cpol
+
+    MIB.addImm(Offset + RegOffset);
+
+    bool LastUse = MMO->getFlags() & MOLastUse;
+    MIB.addImm(LastUse ? AMDGPU::CPol::TH_LU : 0); // cpol
+
     if (!IsFlat)
       MIB.addImm(0); // swz
     MIB.addMemOperand(NewMMO);
@@ -2241,6 +2245,7 @@ bool SIRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator MI,
         TII->insertScratchExecCopy(*MF, *MBB, MI, DL, MFI->getSGPRForEXECCopy(),
                                   RS->isRegUsed(AMDGPU::SCC));
       }
+
       buildSpillLoadStore(
           *MBB, MI, DL, Opc, Index, VData->getReg(), VData->isKill(), FrameReg,
           TII->getNamedOperand(*MI, AMDGPU::OpName::offset)->getImm(),

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

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