From 7ce3af0c41ebffba83298f73003d337e4943521a Mon Sep 17 00:00:00 2001
From: Basem Mohammed <95645899+Basemism@users.noreply.github.com>
Date: Wed, 2 Sep 2026 11:44:53 +0200
Subject: [PATCH 1/4] arch-amdgpu: add a neighbour PWC and configurable walk
 fetch width

Extend the Vega page-table walker with a second page-walk cache for valid
neighbouring final-level PTEs fetched in the same memory block. Keep the
existing PWC for non-final walk levels and defer insertions until the walker
knows whether an entry terminates the walk.

Add pwc_fetch_bytes to control aligned page-table memory requests, propagate
it to GPU, command-processor, and SDMA walkers, and validate that the width is
a supported power-of-two multiple of a PTE. Track accesses, hits, misses,
insertions, invalidations, and rejected invalid neighbours separately for the
two caches.
---
 configs/common/GPUTLBConfig.py           |   3 +
 configs/example/gpufs/system/system.py   |  11 +-
 src/arch/amdgpu/vega/VegaGPUTLB.py       |  18 ++
 src/arch/amdgpu/vega/pagetable_walker.cc | 237 +++++++++++++++++++++--
 src/arch/amdgpu/vega/pagetable_walker.hh | 130 ++++++++++---
 5 files changed, 354 insertions(+), 45 deletions(-)

diff --git a/configs/common/GPUTLBConfig.py b/configs/common/GPUTLBConfig.py
index 82ee5c182c9..646abc747b3 100644
--- a/configs/common/GPUTLBConfig.py
+++ b/configs/common/GPUTLBConfig.py
@@ -39,6 +39,9 @@ def TLB_constructor(options, level, gpu_ctrl=None, full_system=False):
     if full_system:
         constructor_call = "VegaGPUTLB(\
                 gpu_device = gpu_ctrl, \
+                walker = VegaPagetableWalker(\
+                    pwc_fetch_bytes = getattr(\
+                        options, 'pwc_fetch_bytes', 64)), \
                 size = options.L%(level)dTLBentries, \
                 assoc = options.L%(level)dTLBassoc, \
                 hitLatency = options.L%(level)dAccessLatency,\
diff --git a/configs/example/gpufs/system/system.py b/configs/example/gpufs/system/system.py
index 31e1e5eb914..4957999e005 100644
--- a/configs/example/gpufs/system/system.py
+++ b/configs/example/gpufs/system/system.py
@@ -115,7 +115,10 @@ def makeGpuFSSystem(args):
 
     # This arbitrary address is something in the X86 I/O hole
     hsapp_gpu_map_paddr = 0xE0000000
-    hsapp_pt_walker = VegaPagetableWalker()
+    pwc_fetch_bytes = getattr(args, "pwc_fetch_bytes", 64)
+    hsapp_pt_walker = VegaPagetableWalker(
+        pwc_fetch_bytes=pwc_fetch_bytes
+    )
     gpu_hsapp = HSAPacketProcessor(
         pioAddr=hsapp_gpu_map_paddr,
         numHWQueues=args.num_hw_queues,
@@ -127,7 +130,7 @@ def makeGpuFSSystem(args):
     if args.exit_after_gpu_kernel > -1:
         dispatcher_exit_events = True
     dispatcher = GPUDispatcher(kernel_exit_events=dispatcher_exit_events)
-    cp_pt_walker = VegaPagetableWalker()
+    cp_pt_walker = VegaPagetableWalker(pwc_fetch_bytes=pwc_fetch_bytes)
     target_kernel = args.skip_until_gpu_kernel
     gpu_cmd_proc = GPUCommandProcessor(
         hsapp=gpu_hsapp,
@@ -210,7 +213,9 @@ def makeGpuFSSystem(args):
     sdma_pt_walkers = []
     sdma_engines = []
     for sdma_idx in range(num_sdmas):
-        sdma_pt_walker = VegaPagetableWalker()
+        sdma_pt_walker = VegaPagetableWalker(
+            pwc_fetch_bytes=pwc_fetch_bytes
+        )
         sdma_engine = SDMAEngine(
             walker=sdma_pt_walker,
             mmio_base=sdma_bases[sdma_idx],
diff --git a/src/arch/amdgpu/vega/VegaGPUTLB.py b/src/arch/amdgpu/vega/VegaGPUTLB.py
index de204ef600e..d7b0ba2fe05 100644
--- a/src/arch/amdgpu/vega/VegaGPUTLB.py
+++ b/src/arch/amdgpu/vega/VegaGPUTLB.py
@@ -72,6 +72,24 @@ class VegaPagetableWalker(ClockedObject):
     )
     enable_pwc = Param.Bool(True, "Enable page walk cache")
 
+    # Second PWC for caching neighbouring final-level entries fetched from
+    # the same page-table memory fetch. The fetch width is controlled by
+    # pwc_fetch_bytes.
+    neighbour_pwc_entries = Param.Int(256, "Neighbour PWC entries")
+    neighbour_pwc_replacement_policy = Param.BaseReplacementPolicy(
+        LRURP(), "Replacement policy of the neighbour PWC"
+    )
+    neighbour_pwc_indexing_policy = Param.VegaPWCIndexingPolicy(
+        VegaPWCIndexingPolicy(
+            entries=Parent.neighbour_pwc_entries,
+            assoc=Parent.neighbour_pwc_entries,
+        ),
+        "Indexing policy of the neighbour PWC",
+    )
+
+    pwc_fetch_bytes = Param.Unsigned(
+        128, "Page-walk fetch granularity in bytes"
+    )
 
 class VegaGPUTLB(ClockedObject):
     type = "VegaGPUTLB"
diff --git a/src/arch/amdgpu/vega/pagetable_walker.cc b/src/arch/amdgpu/vega/pagetable_walker.cc
index d56f8eb214f..0cbd3656917 100644
--- a/src/arch/amdgpu/vega/pagetable_walker.cc
+++ b/src/arch/amdgpu/vega/pagetable_walker.cc
@@ -91,6 +91,31 @@ Walker::WalkerState::startFunctional(Addr base, Addr vaddr,
         assert(devmem);
         devmem->access(read);
 
+        // Extract the requested 8B entry from the fetched block.
+        // Copy the full line before rewriting word 0 so
+        // pendingLineEntries preserves the original data.
+        assert(lineIndex < walker->pwcFetchEntries);
+        assert((read->getAddr() & (walker->pwcFetchBytes - 1)) == 0);
+        const uint64_t *lineData = read->getConstPtr<uint64_t>();
+        pendingLineAddr = read->getAddr();
+        pendingLineIndex = lineIndex;
+        for (unsigned i = 0; i < walker->pwcFetchEntries; i++)
+            pendingLineEntries[i] = letoh(lineData[i]);
+        pendingLineValid = true;
+
+        // Rewrite word 0 so stepWalk()'s getLE<uint64_t>() sees the
+        // correct entry.
+        uint64_t selectedEntry = letoh(lineData[lineIndex]);
+        read->setLE<uint64_t>(selectedEntry);
+
+        // Populate pending PWC metadata for deferred insertion in stepWalk()
+        Addr originalEntryAddr =
+            read->getAddr() + lineIndex * sizeof(uint64_t);
+        assert((originalEntryAddr & 0x7) == 0);
+        pendingEntryAddr = originalEntryAddr;
+        pendingEntry = selectedEntry;
+        pendingFromPwc = false;
+
         fault = stepWalk();
         assert(fault == NoFault || read == NULL);
 
@@ -136,7 +161,7 @@ Walker::WalkerState::initState(BaseMMU::Mode _mode, Addr baseAddr, Addr vaddr,
     mode = _mode;
     timing = !is_functional;
     enableNX = true;
-    dataSize = 8; // 64-bit PDEs / PTEs
+    dataSize = walker->pwcFetchBytes;
     nextState = PDE2;
 
     DPRINTF(GPUPTWalker, "Setup walk with base %#lx\n", baseAddr);
@@ -148,14 +173,25 @@ Walker::WalkerState::initState(BaseMMU::Mode _mode, Addr baseAddr, Addr vaddr,
     Addr pde2Addr = (((baseAddr >> 6) << 3) + (logical_addr >> 3 * 9)) << 3;
     DPRINTF(GPUPTWalker, "Walk PDE2 address is %#lx\n", pde2Addr);
 
+    // Align to the configured page-walk fetch block boundary.
+    const Addr blkSize = walker->pwcFetchBytes;
+    assert((pde2Addr & 0x7) == 0);           // 8-byte aligned
+    Addr alignedAddr = pde2Addr & ~(blkSize - 1);
+    assert((alignedAddr & (blkSize - 1)) == 0);
+    unsigned idx = (pde2Addr - alignedAddr) / sizeof(uint64_t);
+    assert(idx < walker->pwcFetchEntries);
+    lineIndex = idx;
+    DPRINTF(GPUPTWalker, "Aligned PDE2 %#lx -> %#lx lineIndex %u\n",
+            pde2Addr, alignedAddr, lineIndex);
+
     // Start populating the VegaTlbEntry response
     entry.vaddr = logical_addr;
 
     // Prepare the read packet that will be used at each level
     Request::Flags flags = Request::PHYSICAL;
 
-    RequestPtr request = std::make_shared<Request>(pde2Addr, dataSize, flags,
-                                                   walker->deviceRequestorId);
+    RequestPtr request = std::make_shared<Request>(
+        alignedAddr, dataSize, flags, walker->deviceRequestorId);
 
     read = new Packet(request, MemCmd::ReadReq);
     read->allocate();
@@ -211,17 +247,64 @@ Walker::WalkerState::stepWalk()
 
     walkStateMachine(pte, nextRead, doEndWalk, fault);
 
+    // Deferred PWC insertion: insert the previous response into the original
+    // PWC only if it was not final (doEndWalk means it IS the final entry)
+    // and not already from the PWC.
+    if (!doEndWalk && !pendingFromPwc && walker->enable_pwc
+        && walker->pwc.findEntry(pendingEntryAddr) == nullptr) {
+        walker->pwc.insert(pendingEntryAddr, pendingEntry);
+        walker->stats.pwcInsertions++;
+    }
+
+    // Deferred PWC2 neighbour insertion: on final entries from real memory
+    // responses, insert the other entries from the fetched line into PWC2.
+    if (doEndWalk && !pendingFromPwc && pendingLineValid
+        && walker->enable_pwc) {
+        DPRINTF(GPUPTWalker, "Skipping PWC insert for final entry %#lx, "
+                "inserting neighbours into PWC2\n", pendingEntryAddr);
+        for (unsigned i = 0; i < walker->pwcFetchEntries; i++) {
+            if (i == pendingLineIndex)
+                continue;
+
+            Addr neighbourAddr =
+                pendingLineAddr + i * sizeof(uint64_t);
+            PageTableEntry neighbourPte = pendingLineEntries[i];
+
+            // Skip invalid/non-present page-table entries
+            if (!neighbourPte.v) {
+                walker->stats.pwc2InvalidNeighboursSkipped++;
+                continue;
+            }
+
+            if (walker->neighbourPwc.findEntry(neighbourAddr) == nullptr) {
+                walker->neighbourPwc.insert(neighbourAddr, neighbourPte);
+                walker->stats.pwc2Insertions++;
+            }
+        }
+    }
+
     if (doEndWalk) {
         DPRINTF(GPUPTWalker, "ending walk\n");
         endWalk();
     } else {
         PacketPtr oldRead = read;
 
+        // Align the next read to the configured page-walk fetch block.
+        const Addr blkSize = walker->pwcFetchBytes;
+        assert((nextRead & 0x7) == 0);           // 8-byte aligned
+        Addr alignedAddr = nextRead & ~(blkSize - 1);
+        assert((alignedAddr & (blkSize - 1)) == 0);
+        unsigned idx = (nextRead - alignedAddr) / sizeof(uint64_t);
+        assert(idx < walker->pwcFetchEntries);
+        lineIndex = idx;
+        DPRINTF(GPUPTWalker, "Aligned next read %#lx -> %#lx lineIndex %u\n",
+                nextRead, alignedAddr, lineIndex);
         // If we didn't return, we're setting up another read.
         Request::Flags flags = oldRead->req->getFlags();
         flags.set(Request::UNCACHEABLE, uncacheable);
         RequestPtr request = std::make_shared<Request>(
-            nextRead, oldRead->getSize(), flags, walker->deviceRequestorId);
+            alignedAddr, oldRead->getSize(), flags,
+            walker->deviceRequestorId);
 
         read = new Packet(request, MemCmd::ReadReq);
         read->allocate();
@@ -385,18 +468,48 @@ Walker::WalkerState::sendPackets()
 bool
 Walker::sendTiming(WalkerState *sending_walker, PacketPtr pkt)
 {
-    auto walker_state = new WalkerSenderState(sending_walker);
+    auto walker_state = new WalkerSenderState(sending_walker,
+                                              sending_walker->lineIndex);
     pkt->pushSenderState(walker_state);
 
-    // If hit, send the response pkt immediately.
-    PWCEntry *entry = pwc.findEntry(pkt->getAddr());
-    if (entry != nullptr) {
-        DPRINTF(GPUPTWalker, "PTE found in buffer, skipping timing request.");
-        pkt->setLE<uint64_t>(entry->pteEntry);
+    // Reconstruct original 8B entry address from aligned fetch block address
+    Addr originalEntryAddr =
+        pkt->getAddr() + sending_walker->lineIndex * sizeof(uint64_t);
+    assert(sending_walker->lineIndex < pwcFetchEntries);
+    assert((originalEntryAddr & 0x7) == 0);
+
+    if (enable_pwc) {
+        // Check original PWC first (non-final walk entries)
+        stats.pwcAccesses++;
+        PWCEntry *entry = pwc.findEntry(originalEntryAddr);
+        if (entry != nullptr) {
+            stats.pwcHits++;
+            DPRINTF(GPUPTWalker,
+                    "PTE found in PWC, skipping timing request.");
+            pkt->setLE<uint64_t>(entry->pteEntry);
+            walker_state->fromPwc = true;
 
-        recvTimingResp(pkt);
+            recvTimingResp(pkt);
 
-        return true;
+            return true;
+        }
+        stats.pwcMisses++;
+
+        // Check PWC2 second (neighbouring final-level entries)
+        stats.pwc2Accesses++;
+        PWCEntry *entry2 = neighbourPwc.findEntry(originalEntryAddr);
+        if (entry2 != nullptr) {
+            stats.pwc2Hits++;
+            DPRINTF(GPUPTWalker,
+                    "PTE found in PWC2, skipping timing request.");
+            pkt->setLE<uint64_t>(entry2->pteEntry);
+            walker_state->fromPwc = true;
+
+            recvTimingResp(pkt);
+
+            return true;
+        }
+        stats.pwc2Misses++;
     }
 
     if (port.sendTimingReq(pkt)) {
@@ -426,14 +539,70 @@ Walker::recvTimingResp(PacketPtr pkt)
     WalkerSenderState *senderState =
         safe_cast<WalkerSenderState *>(pkt->popSenderState());
 
-    DPRINTF(GPUPTWalker, "Got response for %#lx from walker %p -- %#lx\n",
-            pkt->getAddr(), senderState->senderWalk, pkt->getLE<uint64_t>());
-    // on PWC miss, add the entry to PWC
-    if (enable_pwc && pwc.findEntry(pkt->getAddr()) == nullptr) {
-        pwc.insert(pkt->getAddr(), pkt->getLE<uint64_t>());
+    assert(senderState->lineIndex < pwcFetchEntries);
+    // Reconstruct original 8B entry address from aligned fetch block address
+    Addr originalEntryAddr =
+        pkt->getAddr() + senderState->lineIndex * sizeof(uint64_t);
+    assert((originalEntryAddr & 0x7) == 0);
+
+    WalkerState *ws = senderState->senderWalk;
+    uint64_t selectedEntry;
+
+    if (senderState->fromPwc) {
+        // PWC hit: word 0 already contains the correct entry
+        selectedEntry = pkt->getLE<uint64_t>();
+        ws->pendingLineValid = false;
+    } else {
+        // Real memory response. The timing read port can return stale data:
+        // page tables are written functionally to device memory (by the
+        // driver/CP) and are NOT coherent with the GPU timing cache hierarchy,
+        // so the timing read may observe an old/invalid PDE (e.g. ppn=0, v=0)
+        // even though device memory holds the correct entry. Re-read the
+        // configured page-table fetch block directly from device memory (as
+        // the functional walk does) so the walk observes the correct PTEs.
+        // The timing request is still used to model walk latency.
+        std::vector<uint64_t> devline(pwcFetchEntries, 0);
+        RequestPtr freq = std::make_shared<Request>(
+            pkt->getAddr(), pwcFetchBytes, Request::PHYSICAL,
+            deviceRequestorId);
+        PacketPtr fpkt = new Packet(freq, MemCmd::ReadReq);
+        fpkt->dataStatic(reinterpret_cast<uint8_t *>(devline.data()));
+        auto *devmem = system->getDeviceMemory(fpkt);
+        const uint64_t *lineData;
+        if (devmem) {
+            devmem->access(fpkt);
+            lineData = devline.data();
+        } else {
+            lineData = pkt->getConstPtr<uint64_t>();
+        }
+
+        ws->pendingLineAddr = pkt->getAddr();
+        ws->pendingLineIndex = senderState->lineIndex;
+        assert((ws->pendingLineAddr & (pwcFetchBytes - 1)) == 0);
+        for (unsigned i = 0; i < pwcFetchEntries; i++)
+            ws->pendingLineEntries[i] = letoh(lineData[i]);
+        ws->pendingLineValid = true;
+
+        // Extract from the fetched block and rewrite word 0
+        // so stepWalk()'s getLE<uint64_t>() sees the correct entry.
+        selectedEntry = letoh(lineData[senderState->lineIndex]);
+        pkt->setLE<uint64_t>(selectedEntry);
+
+        delete fpkt;
     }
 
-    senderState->senderWalk->startWalk();
+    DPRINTF(GPUPTWalker, "Got response for %#lx (entry %#lx) from walker %p "
+            "lineIndex %u fromPwc %d -- %#lx\n",
+            pkt->getAddr(), originalEntryAddr, ws,
+            senderState->lineIndex, senderState->fromPwc, selectedEntry);
+
+    // Store response metadata in WalkerState for deferred PWC insertion.
+    // stepWalk() will insert into the original PWC only for non-final entries.
+    ws->pendingEntryAddr = originalEntryAddr;
+    ws->pendingEntry = selectedEntry;
+    ws->pendingFromPwc = senderState->fromPwc;
+
+    ws->startWalk();
 
     delete senderState;
 }
@@ -444,6 +613,13 @@ Walker::invalidatePWC()
     for (auto &i : pwc) {
         if (i.valid) {
             pwc.invalidate(&i);
+            stats.pwcInvalidations++;
+        }
+    }
+    for (auto &i : neighbourPwc) {
+        if (i.valid) {
+            neighbourPwc.invalidate(&i);
+            stats.pwc2Invalidations++;
         }
     }
 }
@@ -519,6 +695,31 @@ Walker::WalkerState::offsetFunc(Addr logicalAddr, int top, int lsb)
     return ((logicalAddr & ((1 << top) - 1)) >> lsb);
 }
 
+/**
+ * Stats
+ */
+Walker::WalkerStats::WalkerStats(statistics::Group *parent)
+    : statistics::Group(parent),
+      ADD_STAT(pwcAccesses, "Number of accesses to the original PWC"),
+      ADD_STAT(pwcHits, "Number of hits in the original PWC"),
+      ADD_STAT(pwcMisses, "Number of misses in the original PWC"),
+      ADD_STAT(pwcInsertions, "Number of insertions into the original PWC"),
+      ADD_STAT(pwcInvalidations,
+               "Number of invalidations in the original PWC"),
+      ADD_STAT(pwc2Accesses, "Number of accesses to the neighbour PWC"),
+      ADD_STAT(pwc2Hits, "Number of hits in the neighbour PWC"),
+      ADD_STAT(pwc2Misses, "Number of misses in the neighbour PWC"),
+      ADD_STAT(pwc2Insertions,
+               "Number of insertions into the neighbour PWC"),
+      ADD_STAT(pwc2Invalidations,
+               "Number of invalidations in the neighbour PWC"),
+      ADD_STAT(pwc2InvalidNeighboursSkipped,
+               "Number of invalid neighbour entries skipped during "
+               "PWC2 insertion")
+{
+}
+
+
 /**
  * gem5 methods
  */
diff --git a/src/arch/amdgpu/vega/pagetable_walker.hh b/src/arch/amdgpu/vega/pagetable_walker.hh
index eb8369c82bf..5d34569504c 100644
--- a/src/arch/amdgpu/vega/pagetable_walker.hh
+++ b/src/arch/amdgpu/vega/pagetable_walker.hh
@@ -37,6 +37,8 @@
 #include "arch/amdgpu/vega/page_walk_cache.hh"
 #include "arch/amdgpu/vega/pagetable.hh"
 #include "arch/amdgpu/vega/tlb.hh"
+#include "base/logging.hh"
+#include "base/statistics.hh"
 #include "base/types.hh"
 #include "debug/GPUPTWalker.hh"
 #include "mem/packet.hh"
@@ -55,8 +57,13 @@ namespace VegaISA
 class Walker : public ClockedObject
 {
   protected:
+    // PWC for non-final page-table walk levels (Global, Upper, Middle).
     PageWalkCache pwc;
 
+    // PWC for neighbouring final-level entries from the same page-table
+    // memory fetch. The fetch width is controlled by pwcFetchBytes.
+    PageWalkCache neighbourPwc;
+
     // Port for accessing memory
     class WalkerPort : public RequestPort
     {
@@ -75,6 +82,14 @@ class Walker : public ClockedObject
     friend class WalkerPort;
     WalkerPort port;
 
+    static constexpr Addr MaxPwcFetchBytes = 1024;
+    static constexpr unsigned MaxPwcFetchEntries =
+        MaxPwcFetchBytes / sizeof(uint64_t);
+
+    // Width of page-table memory fetches used by the walker/PWC2 model.
+    Addr pwcFetchBytes;
+    unsigned pwcFetchEntries;
+
     // State to track each walk of the page table
     class WalkerState
     {
@@ -106,18 +121,30 @@ class Walker : public ClockedObject
         bool timing;
         PacketPtr tlbPkt;
         int blockFragmentSize;
+        // Index of the requested 8B entry within the current fetch block.
+        unsigned lineIndex;
+
+        // Response metadata for deferred PWC insertion (set by
+        // recvTimingResp, consumed by stepWalk after walkStateMachine).
+        Addr pendingEntryAddr;
+        uint64_t pendingEntry;
+        bool pendingFromPwc;
+
+        // Full fetched line data for deferred PWC2 neighbour insertion.
+        // Populated from real memory responses; invalid for PWC hits.
+        Addr pendingLineAddr;
+        uint64_t pendingLineEntries[MaxPwcFetchEntries];
+        unsigned pendingLineIndex;
+        bool pendingLineValid;
 
       public:
         WalkerState(Walker *_walker, PacketPtr pkt, bool is_functional = false)
-            : walker(_walker),
-              state(Ready),
-              nextState(Ready),
-              dataSize(8),
-              enableNX(true),
-              retrying(false),
-              started(false),
-              tlbPkt(pkt),
-              blockFragmentSize(0)
+            : walker(_walker), state(Ready), nextState(Ready), dataSize(0),
+              enableNX(true), retrying(false), started(false), tlbPkt(pkt),
+              blockFragmentSize(0), lineIndex(0),
+              pendingEntryAddr(0), pendingEntry(0), pendingFromPwc(false),
+              pendingLineAddr(0), pendingLineIndex(0),
+              pendingLineValid(false)
         {
             DPRINTF(GPUPTWalker, "Walker::WalkerState %p %p %d\n", this,
                     walker, state);
@@ -154,6 +181,8 @@ class Walker : public ClockedObject
     };
 
     friend class WalkerState;
+
+
     // State for timing and atomic accesses (need multiple per walker in
     // the case of multiple outstanding requests in timing mode)
     std::list<WalkerState *> currStates;
@@ -162,9 +191,15 @@ class Walker : public ClockedObject
 
     struct WalkerSenderState : public Packet::SenderState
     {
-        WalkerState *senderWalk;
-        WalkerSenderState(WalkerState *_senderWalk) : senderWalk(_senderWalk)
-        {}
+        WalkerState * senderWalk;
+        // Index of the requested 8B entry within the fetch block.
+        unsigned lineIndex;
+        // True when the response comes from a PWC hit, not real memory
+        bool fromPwc;
+        WalkerSenderState(WalkerState * _senderWalk, unsigned _lineIndex = 0,
+                          bool _fromPwc = false)
+            : senderWalk(_senderWalk), lineIndex(_lineIndex),
+              fromPwc(_fromPwc) {}
     };
 
   public:
@@ -223,6 +258,26 @@ class Walker : public ClockedObject
     // System pointer for functional accesses
     System *system;
 
+    struct WalkerStats : public statistics::Group
+    {
+        WalkerStats(statistics::Group *parent);
+
+        // Original PWC stats (non-final walk levels)
+        statistics::Scalar pwcAccesses;
+        statistics::Scalar pwcHits;
+        statistics::Scalar pwcMisses;
+        statistics::Scalar pwcInsertions;
+        statistics::Scalar pwcInvalidations;
+
+        // Neighbour PWC stats (final-level neighbour entries)
+        statistics::Scalar pwc2Accesses;
+        statistics::Scalar pwc2Hits;
+        statistics::Scalar pwc2Misses;
+        statistics::Scalar pwc2Insertions;
+        statistics::Scalar pwc2Invalidations;
+        statistics::Scalar pwc2InvalidNeighboursSkipped;
+    } stats;
+
   public:
     void
     setTLB(GpuTLB *_tlb)
@@ -232,19 +287,46 @@ class Walker : public ClockedObject
     }
 
     Walker(const VegaPagetableWalkerParams &p)
-        : ClockedObject(p),
-          pwc(name() + ".pwc", p.page_walk_cache_entries,
-              p.page_walk_cache_entries, p.pwc_replacement_policy,
-              p.pwc_indexing_policy),
-          port(name() + ".port", this),
-          funcState(this, nullptr, true),
-          enable_pwc(p.enable_pwc),
-          tlb(nullptr),
-          requestorId(p.system->getRequestorId(this)),
-          deviceRequestorId(999),
-          system(p.system)
+      : ClockedObject(p),
+        pwc(name()+".pwc", p.page_walk_cache_entries,
+            p.page_walk_cache_entries, p.pwc_replacement_policy,
+            p.pwc_indexing_policy),
+        neighbourPwc(name()+".neighbourPwc", p.neighbour_pwc_entries,
+            p.neighbour_pwc_entries, p.neighbour_pwc_replacement_policy,
+            p.neighbour_pwc_indexing_policy),
+        port(name() + ".port", this),
+        pwcFetchBytes(p.pwc_fetch_bytes),
+        pwcFetchEntries(p.pwc_fetch_bytes / sizeof(uint64_t)),
+        funcState(this, nullptr, true),
+        enable_pwc(p.enable_pwc),
+	tlb(nullptr),
+        requestorId(p.system->getRequestorId(this)),
+        deviceRequestorId(999),
+        system(p.system),
+        stats(this)
     {
-        DPRINTF(GPUPTWalker, "Walker::Walker %p\n", this);
+        fatal_if(pwcFetchBytes < sizeof(uint64_t),
+                 "VegaPagetableWalker pwc_fetch_bytes must be at least 8 "
+                 "bytes");
+        fatal_if(pwcFetchBytes % sizeof(uint64_t) != 0,
+                 "VegaPagetableWalker pwc_fetch_bytes must be divisible "
+                 "by 8");
+        fatal_if((pwcFetchBytes & (pwcFetchBytes - 1)) != 0,
+                 "VegaPagetableWalker pwc_fetch_bytes must be a power "
+                 "of two");
+        fatal_if(pwcFetchBytes > MaxPwcFetchBytes,
+                 "VegaPagetableWalker pwc_fetch_bytes exceeds supported "
+                 "maximum of 1024 bytes");
+        fatal_if(pwcFetchBytes > system->cacheLineSize(),
+                 "VegaPagetableWalker pwc_fetch_bytes (%lu) exceeds the "
+                 "system cache-line size (%lu)", pwcFetchBytes,
+                 system->cacheLineSize());
+        fatal_if(pwcFetchEntries > MaxPwcFetchEntries,
+                 "VegaPagetableWalker pwc_fetch_entries exceeds supported "
+                 "maximum");
+
+        DPRINTF(GPUPTWalker, "Walker::Walker %p fetchBytes %lu "
+                "fetchEntries %u\n", this, pwcFetchBytes, pwcFetchEntries);
     }
 };
 

From f319f1cb96117b3a8d0a41ec412816912f15fe0e Mon Sep 17 00:00:00 2001
From: Basem Mohammed <95645899+Basemism@users.noreply.github.com>
Date: Wed, 2 Sep 2026 11:44:53 +0200
Subject: [PATCH 2/4] arch-amdgpu: improve large-page TLB coalescing and
 prediction

Track Vega TLB lookup outcomes by page size and extend lookups to the large
page sizes the walker can install. Probe additional page sizes during a walk
and clamp page sizes used as coalescer keys so a giant mapping cannot
serialize an unbounded address region.

Add adaptive predictors for page size and coalescing distance, together with
statistics for their decisions and outcomes. Preserve request completion and
retry ordering while allowing coalescing granularity to follow observed TLB
behavior.
---
 configs/common/GPUTLBConfig.py           |   8 +
 src/arch/amdgpu/vega/VegaGPUTLB.py       |   5 +
 src/arch/amdgpu/vega/pagetable_walker.cc |  11 +-
 src/arch/amdgpu/vega/tlb.cc              |  71 ++++++-
 src/arch/amdgpu/vega/tlb.hh              |  35 ++-
 src/arch/amdgpu/vega/tlb_coalescer.cc    | 257 ++++++++++++++++++-----
 src/arch/amdgpu/vega/tlb_coalescer.hh    |  57 ++++-
 7 files changed, 376 insertions(+), 68 deletions(-)

diff --git a/configs/common/GPUTLBConfig.py b/configs/common/GPUTLBConfig.py
index 646abc747b3..267ce7d2ce7 100644
--- a/configs/common/GPUTLBConfig.py
+++ b/configs/common/GPUTLBConfig.py
@@ -214,6 +214,14 @@ def config_tlb_hierarchy(
                         system.%s_tlb[%d].cpu_side_ports[0]"
                     % (name, index, name, index)
                 )
+                # Give each Vega coalescer a handle to the TLB it feeds so the
+                # L3 coalescer can read that TLB's line-coalescing predictor.
+                if full_system:
+                    exec(
+                        "system.%s_coalescer[%d].downstream_tlb = \
+                            system.%s_tlb[%d]"
+                        % (name, index, name, index)
+                    )
 
     # Connect the cpuSidePort of all the coalescers in level 1
     # < Modify here if you want a different configuration >
diff --git a/src/arch/amdgpu/vega/VegaGPUTLB.py b/src/arch/amdgpu/vega/VegaGPUTLB.py
index d7b0ba2fe05..11a2c6d7a9f 100644
--- a/src/arch/amdgpu/vega/VegaGPUTLB.py
+++ b/src/arch/amdgpu/vega/VegaGPUTLB.py
@@ -124,3 +124,8 @@ class VegaTLBCoalescer(ClockedObject):
     cpu_side_ports = VectorResponsePort("Port on side closer to CPU/CU")
     mem_side_ports = VectorRequestPort("Port on side closer to memory")
     disableCoalescing = Param.Bool(False, "Dispable Coalescing")
+    downstream_tlb = Param.VegaGPUTLB(
+        NULL,
+        "The TLB this coalescer feeds; used by the last-level coalescer to read the "
+        "line-coalescing predictor",
+    )
diff --git a/src/arch/amdgpu/vega/pagetable_walker.cc b/src/arch/amdgpu/vega/pagetable_walker.cc
index 0cbd3656917..da0f24df219 100644
--- a/src/arch/amdgpu/vega/pagetable_walker.cc
+++ b/src/arch/amdgpu/vega/pagetable_walker.cc
@@ -223,12 +223,7 @@ Walker::WalkerState::startWalk()
             entry.paddr = entry.pte.ppn << PageShift;
             entry.paddr += entry.vaddr & mask(entry.logBytes);
 
-            // Insert to TLB
-            assert(walker);
-            assert(walker->tlb);
-            walker->tlb->insert(entry.vaddr, entry);
-
-            // Send translation return event
+            // Send translation return event. The TLB allocates the returned entry in the normal miss-return path.
             walker->walkerResponse(this, entry, tlbPkt);
         }
     }
@@ -648,7 +643,9 @@ Walker::recvReqRetry()
 void
 Walker::walkerResponse(WalkerState *state, VegaTlbEntry &entry, PacketPtr pkt)
 {
-    tlb->walkerResponse(entry, pkt);
+    // Propagate whether the final PTE came from the PWC/PWC2 (a PWC hit) or
+    // from memory (a real miss) so the TLB can update its line predictor.
+    tlb->walkerResponse(entry, pkt, state->pendingFromPwc);
 
     delete state;
 }
diff --git a/src/arch/amdgpu/vega/tlb.cc b/src/arch/amdgpu/vega/tlb.cc
index d21a7ced2c9..f5158f21a8d 100644
--- a/src/arch/amdgpu/vega/tlb.cc
+++ b/src/arch/amdgpu/vega/tlb.cc
@@ -167,6 +167,15 @@ GpuTLB::insert(Addr vpn, VegaTlbEntry &entry)
 
     int set = getSet(entry.vaddr, entry.logBytes);
 
+    // Reuse an existing same-size page entry; duplicates can break shootdown.
+    auto existing = lookupIt(entry.vaddr, entry.logBytes, true);
+    if (existing != entryList[set].end()) {
+        DPRINTF(GPUTLB, "Reused existing %#lx -> %#lx of size %#lx in "
+                "set %d\n", (*existing)->vaddr, (*existing)->paddr,
+                (*existing)->size(), set);
+        return *existing;
+    }
+
     if (!freeList[set].empty()) {
         newEntry = freeList[set].front();
         freeList[set].pop_front();
@@ -176,8 +185,18 @@ GpuTLB::insert(Addr vpn, VegaTlbEntry &entry)
     }
 
     *newEntry = entry;
+
+    // Store a canonical page base for range checks.
+    newEntry->vaddr = newEntry->vaddr & ~mask(newEntry->logBytes);
+
     entryList[set].push_front(newEntry);
 
+    if (entry.logBytes == VegaISA::PageShift) {
+        stats.inserts4K++;
+    } else if (entry.logBytes == 21) {
+        stats.inserts2M++;
+    }
+
     DPRINTF(GPUTLB, "Inserted %#lx -> %#lx of size %#lx into set %d\n",
             newEntry->vaddr, newEntry->paddr, entry.size(), set);
 
@@ -197,12 +216,21 @@ GpuTLB::lookupIt(Addr va, unsigned int ps, bool update_lru)
     for (; entry != entryList[set].end(); ++entry) {
         int page_size = (*entry)->size();
 
-        if ((*entry)->vaddr <= va && (*entry)->vaddr + page_size > va &&
+        // Compare against the page-aligned base of the entry, not its raw
+        // vaddr. getSet() already aligns (it shifts va by the page size), but
+        // the range check below does not: if an entry was filled with a vaddr
+        // that is not aligned to its own page size, the raw comparison both
+        // rejects in-page addresses below (*entry)->vaddr and lets the range
+        // (*entry)->vaddr + page_size spill past the page boundary. Aligning
+        // here makes lookups correct regardless of how the fill was aligned.
+        Addr entry_base = (*entry)->vaddr & ~mask((*entry)->logBytes);
+
+        if (entry_base <= va && entry_base + page_size > va &&
             ps == (*entry)->logBytes) {
             DPRINTF(GPUTLB,
                     "Matched vaddr %#x to entry starting at %#x "
                     "with size %#x.\n",
-                    va, (*entry)->vaddr, page_size);
+                    va, entry_base, page_size);
 
             if (update_lru) {
                 entryList[set].push_front(*entry);
@@ -277,11 +305,10 @@ GpuTLB::tlbLookup(const RequestPtr &req, bool update_stats)
         return NULL;
     }
     Addr vaddr = req->getVaddr();
-    Addr alignedVaddr = pageAlign(vaddr);
     DPRINTF(GPUTLB, "TLB Lookup for vaddr %#x.\n", vaddr);
 
     // update LRU stack on a hit
-    VegaTlbEntry *entry = lookup(alignedVaddr, true);
+    VegaTlbEntry *entry = lookup(vaddr, true);
 
     if (!update_stats) {
         // functional tlb access for memory initialization
@@ -296,6 +323,11 @@ GpuTLB::tlbLookup(const RequestPtr &req, bool update_stats)
         stats.localNumTLBMisses++;
     } else {
         stats.localNumTLBHits++;
+        if (entry->logBytes == VegaISA::PageShift) {
+            stats.localHits4K++;
+        } else if (entry->logBytes == 21) {
+            stats.localHits2M++;
+        }
     }
 
     return entry;
@@ -358,6 +390,9 @@ GpuTLB::issueTLBLookup(PacketPtr pkt)
     if (entry || pkt->req->hasNoAddr()) {
         // Put the entry in SenderState
         lookup_outcome = TLB_HIT;
+        // A genuine hit in this TLB's entry array biases the line-coalescing
+        // predictor away from line prefetch (we are not missing here).
+        noteTlbHit();
         if (pkt->req->hasNoAddr()) {
             sender_state->tlbEntry =
                 new VegaTlbEntry(1 /* VMID */, 0, 0, 0, 0);
@@ -442,11 +477,18 @@ GpuTLB::pagingProtectionChecks(PacketPtr pkt, VegaTlbEntry *tlb_entry,
 }
 
 void
-GpuTLB::walkerResponse(VegaTlbEntry &entry, PacketPtr pkt)
+GpuTLB::walkerResponse(VegaTlbEntry &entry, PacketPtr pkt, bool from_pwc)
 {
     DPRINTF(GPUTLB, "WalkerResponse for %#lx. Entry: (%#lx, %#lx, %#lx)\n",
             pkt->req->getVaddr(), entry.vaddr, entry.paddr, entry.size());
 
+    // A walk that had to reach memory is a real TLB miss and biases the
+    // line-coalescing predictor toward line prefetch. A walk satisfied by the
+    // PWC/PWC2 deliberately does not count (neither hit nor miss).
+    if (!from_pwc) {
+        noteTlbMiss();
+    }
+
     Addr virt_page_addr = roundDown(pkt->req->getVaddr(), VegaISA::PageBytes);
 
     Addr page_addr = entry.pte.ppn << VegaISA::PageShift;
@@ -458,6 +500,12 @@ GpuTLB::walkerResponse(VegaTlbEntry &entry, PacketPtr pkt)
         safe_cast<GpuTranslationState *>(pkt->senderState);
     sender_state->tlbEntry = new VegaTlbEntry(entry);
 
+    if (entry.logBytes == VegaISA::PageShift) {
+        stats.walkerReturns4K++;
+    } else if (entry.logBytes == 21) {
+        stats.walkerReturns2M++;
+    }
+
     handleTranslationReturn(virt_page_addr, TLB_MISS, pkt);
 }
 
@@ -623,8 +671,7 @@ GpuTLB::translationReturn(Addr virtPageAddr, tlbOutcome outcome, PacketPtr pkt)
             TLBEvent *tlb_event = translationReturnEvent[virtPageAddr];
             assert(tlb_event);
             tlb_event->updateOutcome(PAGE_WALK);
-            schedule(tlb_event,
-                     curTick() + cyclesToTicks(Cycles(missLatency2)));
+            schedule(tlb_event, curTick());
         }
     } else if (outcome == PAGE_WALK) {
         if (update_stats) {
@@ -996,8 +1043,14 @@ GpuTLB::VegaTLBStats::VegaTLBStats(statistics::Group *parent)
       ADD_STAT(accessCycles, "Cycles spent accessing this TLB level"),
       ADD_STAT(pageTableCycles, "Cycles spent accessing the page table"),
       ADD_STAT(localCycles, "Number of cycles spent in queue for all "
-                            "incoming reqs"),
-      ADD_STAT(localLatency, "Avg. latency over incoming coalesced reqs")
+                             "incoming reqs"),
+      ADD_STAT(localLatency, "Avg. latency over incoming coalesced reqs"),
+      ADD_STAT(inserts4K, "Number of 4 KiB entries inserted"),
+      ADD_STAT(inserts2M, "Number of 2 MiB entries inserted"),
+      ADD_STAT(localHits4K, "Number of local hits on 4 KiB entries"),
+      ADD_STAT(localHits2M, "Number of local hits on 2 MiB entries"),
+       ADD_STAT(walkerReturns4K, "Number of page walks returning 4 KiB entries"),
+       ADD_STAT(walkerReturns2M, "Number of page walks returning 2 MiB entries")
 {
     localTLBMissRate = 100 * localNumTLBMisses / localNumTLBAccesses;
     globalTLBMissRate = 100 * globalNumTLBMisses / globalNumTLBAccesses;
diff --git a/src/arch/amdgpu/vega/tlb.hh b/src/arch/amdgpu/vega/tlb.hh
index 1c8ad0416e5..9b8dc823b56 100644
--- a/src/arch/amdgpu/vega/tlb.hh
+++ b/src/arch/amdgpu/vega/tlb.hh
@@ -41,6 +41,7 @@
 #include "arch/generic/mmu.hh"
 #include "base/statistics.hh"
 #include "base/trace.hh"
+#include "base/sat_counter.hh"
 #include "mem/packet.hh"
 #include "mem/port.hh"
 #include "params/VegaGPUTLB.hh"
@@ -115,9 +116,10 @@ class GpuTLB : public ClockedObject
 
     int getSet(Addr va, unsigned int page_shift);
 
-    // List of possible page size, 4k and 2m for now
-    const std::array<unsigned int, 2> logPageShiftList = {VegaISA::PageShift,
-                                                          21};
+    // Page sizes probed on lookup (4K, 2M, 1G, 512G). Must match the sizes
+    // the walker can install, else large-page entries are never found.
+    const std::array<unsigned int, 4> logPageShiftList = {VegaISA::PageShift,
+                                                          21, 30, 39};
 
     int size;
     int assoc;
@@ -193,6 +195,14 @@ class GpuTLB : public ClockedObject
         // from the perspective of this TLB
         statistics::Scalar localCycles;
         statistics::Formula localLatency;
+
+        // Page-size breakdowns for diagnosing TLB reuse/capacity behavior.
+        statistics::Scalar inserts4K;
+        statistics::Scalar inserts2M;
+        statistics::Scalar localHits4K;
+        statistics::Scalar localHits2M;
+        statistics::Scalar walkerReturns4K;
+        statistics::Scalar walkerReturns2M;
     } stats;
 
     VegaTlbEntry *insert(Addr vpn, VegaTlbEntry &entry);
@@ -210,7 +220,24 @@ class GpuTLB : public ClockedObject
     };
     VegaTlbEntry *tlbLookup(const RequestPtr &req, bool update_stats);
 
-    void walkerResponse(VegaTlbEntry &entry, PacketPtr pkt);
+    // Line-coalescing predictor: a 4-bit saturating counter,
+    // meaningful only at the L3 TLB. A real TLB miss (walk that reached memory)
+    // decrements it toward "expect an L3 miss"; a TLB hit increments it. A PWC
+    // or PWC2 hit deliberately does not count as a hit and leaves it unchanged.
+    // When the MSB is clear we expect an L3 miss and it is worthwhile for the
+    // L3 coalescer to line-coalesce (prefetch a full PWC2 line).
+    static constexpr unsigned LinePredBits = 4;
+    SatCounter8 lineCounter{LinePredBits, 0};
+    void noteTlbHit() { lineCounter++; }
+    void noteTlbMiss() { lineCounter--; }
+    bool shouldLineCoalesce() const
+    {
+        const uint8_t msb = 1u << (LinePredBits - 1);
+        return !(static_cast<uint8_t>(lineCounter) & msb);
+    }
+
+    void walkerResponse(VegaTlbEntry &entry, PacketPtr pkt,
+                        bool from_pwc = false);
     void handleTranslationReturn(Addr addr, tlbOutcome outcome, PacketPtr pkt);
 
     void handleFuncTranslationReturn(PacketPtr pkt, tlbOutcome outcome);
diff --git a/src/arch/amdgpu/vega/tlb_coalescer.cc b/src/arch/amdgpu/vega/tlb_coalescer.cc
index 0a41732124d..9760876d992 100644
--- a/src/arch/amdgpu/vega/tlb_coalescer.cc
+++ b/src/arch/amdgpu/vega/tlb_coalescer.cc
@@ -31,6 +31,7 @@
 
 #include "arch/amdgpu/vega/tlb_coalescer.hh"
 
+#include <algorithm>
 #include <cstring>
 
 #include "arch/amdgpu/common/gpu_translation_state.hh"
@@ -53,6 +54,10 @@ VegaTLBCoalescer::VegaTLBCoalescer(const VegaTLBCoalescerParams &p)
       cleanupEvent([this] { processCleanupEvent(); },
                    "Cleanup issuedTranslationsTable hashmap", false,
                    Event::Maximum_Pri),
+      // Start the size predictor saturated toward 2 MiB (MSB clear), matching
+      // the previous hardcoded default_pgSize speculation.
+      sizePredictor(SizePredBits, 0),
+      downstreamTLB(p.downstream_tlb),
       tlb_level(p.tlb_level),
       maxDownstream(p.maxDownstream),
       numDownstream(0)
@@ -71,6 +76,85 @@ VegaTLBCoalescer::VegaTLBCoalescer(const VegaTLBCoalescerParams &p)
 
     default_pgSize = p.default_pgSize;
     potentialPagesize.insert(default_pgSize);
+    // Always consider the 4 KiB boundary as well. The issue-side reissue path
+    // keys mispredicted requests in issuedTranslationsTable at 4 KiB, so both
+    // the outstanding-page block check and updatePhysAddresses must probe the
+    // 4 KiB boundary even before any 4 KiB translation has returned.
+    potentialPagesize.insert(VegaISA::PageBytes);
+}
+
+void
+VegaTLBCoalescer::updateSizePredictor(Addr returned_pgsize)
+{
+    // 2 MiB (or larger) return -> bias toward 2 MiB (decrement);
+    // anything smaller -> bias toward 4 KiB (increment).
+    if (returned_pgsize >= default_pgSize) {
+        sizePredictor--;
+    } else {
+        sizePredictor++;
+    }
+}
+
+Addr
+VegaTLBCoalescer::predictedPageSize() const
+{
+    // MSB set -> predict 4 KiB, otherwise the large (default) page size.
+    const uint8_t msb = 1u << (SizePredBits - 1);
+    return (static_cast<uint8_t>(sizePredictor) & msb) ? VegaISA::PageBytes
+                                                       : default_pgSize;
+}
+
+Addr
+VegaTLBCoalescer::newGroupPageSize()
+{
+    Addr pg_size = predictedPageSize();
+    if (lineCoalesceEnabled()) {
+        pg_size *= LineGroupPages;
+        // The line granule becomes a key in issuedTranslationsTable, so the
+        // return-side finder and the outstanding-page block check must probe
+        // it too.
+        potentialPagesize.insert(pg_size);
+    }
+    return pg_size;
+}
+
+size_t
+VegaTLBCoalescer::coalescerFIFOEntries() const
+{
+    size_t entries = 0;
+    for (const auto &tick_entry : coalescerFIFO) {
+        entries += tick_entry.second.size();
+    }
+    return entries;
+}
+
+size_t
+VegaTLBCoalescer::issuedTranslationPackets() const
+{
+    size_t packets = 0;
+    for (const auto &entry : issuedTranslationsTable) {
+        packets += entry.second.size();
+    }
+    return packets;
+}
+
+void
+VegaTLBCoalescer::updateOccupancyStats()
+{
+    const size_t fifo_entries = coalescerFIFOEntries();
+    if (fifo_entries > fifoMaxEntries.value()) {
+        fifoMaxEntries = fifo_entries;
+    }
+
+    const size_t issued_entries = issuedTranslationsTable.size();
+    if (issued_entries > issuedTranslationsMax.value()) {
+        issuedTranslationsMax = issued_entries;
+    }
+
+    const size_t issued_packets = issuedTranslationPackets();
+    if (issued_packets > issuedTranslationPacketsMax.value()) {
+        issuedTranslationPacketsMax = issued_packets;
+    }
 }
 
 Port &
@@ -161,60 +245,79 @@ VegaTLBCoalescer::updatePhysAddresses(PacketPtr pkt)
         *safe_cast<VegaISA::VegaTlbEntry *>(sender_state->tlbEntry);
     Addr first_entry_vaddr = tlb_entry.vaddr;
     Addr first_entry_paddr = tlb_entry.paddr;
-    int page_size = tlb_entry.size();
+    Addr page_size = tlb_entry.size();
+
+    // Clamp giant walker sizes before they become key granules; true
+    // page_size is still used for the range check and paddr computation.
+    potentialPagesize.insert(clampCoalescePageSize(page_size));
 
-    potentialPagesize.insert(page_size);
+    // Train future page-size speculation from the actual return size.
+    updateSizePredictor(page_size);
 
-    Addr virt_page_addr;
+    bool uncacheable = tlb_entry.uncacheable();
+    int first_hit_level = sender_state->hitLevel;
+    bool is_system = pkt->req->systemReq();
+
+    // Save before responding; pkt may be recycled after sendTimingResp().
+    const Addr ret_vaddr = pkt->req->getVaddr();
 
-    // Find coalesced translation request.
+    // Find the outstanding bucket that actually contains this packet.
+    // Overlapping 2 MiB and 4 KiB buckets can both be live, so the key alone
+    // is ambiguous; packet identity disambiguates it.
+    Addr virt_page_addr = 0;
+    auto table_it = issuedTranslationsTable.end();
     for (auto pgsize_seen : potentialPagesize) {
-        virt_page_addr = roundDown(pkt->req->getVaddr(), pgsize_seen);
-        if (issuedTranslationsTable.count(virt_page_addr) != 0) {
+        Addr loc_virt_page_addr = roundDown(ret_vaddr, pgsize_seen);
+        auto it = issuedTranslationsTable.find(loc_virt_page_addr);
+        if (it == issuedTranslationsTable.end()) {
+            continue;
+        }
+        if (std::find(it->second.begin(), it->second.end(), pkt) !=
+            it->second.end()) {
+            virt_page_addr = loc_virt_page_addr;
+            table_it = it;
             break;
         }
     }
 
-    DPRINTF(GPUTLB, "Update phys. addr. for %d \
-            coalesced reqs for page %#x\n",
-            issuedTranslationsTable[virt_page_addr].size(), virt_page_addr);
+    // A returned packet must belong to one outstanding bucket.
+    assert(table_it != issuedTranslationsTable.end());
 
-    bool uncacheable = tlb_entry.uncacheable();
-    int first_hit_level = sender_state->hitLevel;
-    bool is_system = pkt->req->systemReq();
+    // Copy the list since sends/reissues may mutate the table.
+    std::vector<PacketPtr> coalesced_pkts = table_it->second;
+
+    DPRINTF(GPUTLB, "Update phys. addr. for %d coalesced reqs for "
+            "page %#x\n", coalesced_pkts.size(), virt_page_addr);
 
-    for (int i = 0; i < issuedTranslationsTable[virt_page_addr].size(); ++i) {
-        PacketPtr local_pkt = issuedTranslationsTable[virt_page_addr][i];
+    for (int i = 0; i < coalesced_pkts.size(); ++i) {
+        PacketPtr local_pkt = coalesced_pkts[i];
 
         Addr local_pkt_vaddr = local_pkt->req->getVaddr();
 
-        // check if the pending req's vaddr matches the returned page,
-        // if not, reissue pending req as a 4k page
+        // Reissue packets outside the returned page at a clamped granule;
+        // range check uses the true page_size.
         if (!(first_entry_vaddr <= local_pkt_vaddr &&
               local_pkt_vaddr < first_entry_vaddr + page_size)) {
-            reissue_pkt_helper(local_pkt);
+            reissue_pkt_helper(local_pkt, clampCoalescePageSize(page_size));
             continue;
         }
 
-        GpuTranslationState *sender_state =
+        GpuTranslationState *local_sender_state =
             safe_cast<GpuTranslationState *>(local_pkt->senderState);
 
         // we are sending the packet back, so pop the reqCnt associated
-        // with this level in the TLB hierarchy
-        if (!sender_state->isPrefetch) {
-            sender_state->reqCnt.pop_back();
+        // with this level in the TLB hiearchy
+        if (!local_sender_state->isPrefetch) {
+            local_sender_state->reqCnt.pop_back();
             localCycles += curCycle();
         }
 
-        /*
-         * Only the first packet from this coalesced request has been
-         * translated. Grab the translated phys. page addr and update the
-         * physical addresses of the remaining packets with the appropriate
-         * page offsets.
-         */
-        if (i) {
+        // Only the returned packet already has its physical address.
+        // Every other coalesced packet needs to be filled in from the returned page + offset.
+        // Use pointer identity instead of the loop index so this stays correct regardless of the packet's position in the bucket.
+        if (local_pkt != pkt) {
             Addr paddr = first_entry_paddr +
-                         (local_pkt->req->getVaddr() & (page_size - 1));
+                         (local_pkt_vaddr & (page_size - 1));
             local_pkt->req->setPaddr(paddr);
 
             if (uncacheable) {
@@ -223,26 +326,26 @@ VegaTLBCoalescer::updatePhysAddresses(PacketPtr pkt)
 
             // update senderState->tlbEntry, so we can insert
             // the correct TLBEentry in the TLBs above.
-
-            // auto p = sender_state->tc->getProcessPtr();
-            if (sender_state->tlbEntry == NULL) {
+            if (local_sender_state->tlbEntry == NULL) {
                 // not set by lower(l2) coalescer
-                sender_state->tlbEntry = new VegaISA::VegaTlbEntry(
-                    1 /* VMID TODO */, first_entry_vaddr, first_entry_paddr,
-                    tlb_entry.logBytes, tlb_entry.pte);
+                local_sender_state->tlbEntry =
+                    new VegaISA::VegaTlbEntry(
+                        1 /* VMID TODO */, first_entry_vaddr,
+                        first_entry_paddr, tlb_entry.logBytes,
+                        tlb_entry.pte);
             }
 
             // update the hitLevel for all uncoalesced reqs
             // so that each packet knows where it hit
             // (used for statistics in the CUs)
-            sender_state->hitLevel = first_hit_level;
+            local_sender_state->hitLevel = first_hit_level;
         }
 
         // Copy PTE system bit information to coalesced requests
         local_pkt->req->setSystemReq(is_system);
 
-        ResponsePort *return_port = sender_state->ports.back();
-        sender_state->ports.pop_back();
+        ResponsePort *return_port = local_sender_state->ports.back();
+        local_sender_state->ports.pop_back();
 
         // Translation is done - Convert to a response pkt if necessary and
         // send the translation back
@@ -265,9 +368,10 @@ VegaTLBCoalescer::updatePhysAddresses(PacketPtr pkt)
     }
 }
 
-// re-coalesce packet to 4k pages
+// Re-coalesce a packet after page-size speculation failed.
+// 2 MiB returns reissue remaining siblings at 2 MiB; otherwise use 4 KiB.
 void
-VegaTLBCoalescer::reissue_pkt_helper(PacketPtr pkt)
+VegaTLBCoalescer::reissue_pkt_helper(PacketPtr pkt, Addr reissue_pgsize)
 {
     // first packet of a coalesced request
     PacketPtr first_packet = nullptr;
@@ -279,6 +383,12 @@ VegaTLBCoalescer::reissue_pkt_helper(PacketPtr pkt)
     GpuTranslationState *sender_state =
         safe_cast<GpuTranslationState *>(pkt->senderState);
 
+    // Never key a group larger than the coalescer's max granule.
+    reissue_pgsize = clampCoalescePageSize(reissue_pgsize);
+
+    // Ensure return-side lookup probes this reissue size.
+    potentialPagesize.insert(reissue_pgsize);
+
     DPRINTF(GPUTLB, "Trying to re-issue req at tick: %llu, addr: %#x\n",
             sender_state->issueTime, pkt->req->getVaddr());
 
@@ -295,11 +405,11 @@ VegaTLBCoalescer::reissue_pkt_helper(PacketPtr pkt)
     // coalesced request with the same tick_index
     for (int i = 0; i < coalescedReq_cnt; ++i) {
         first_packet = coalescerFIFO[tick_index][i].first[0];
-        if (coalescerFIFO[tick_index][i].second != VegaISA::PageBytes) {
+        if (coalescerFIFO[tick_index][i].second != reissue_pgsize) {
             continue;
         }
 
-        if (canCoalesce(pkt, first_packet, VegaISA::PageBytes)) {
+        if (canCoalesce(pkt, first_packet, reissue_pgsize)) {
             coalescerFIFO[tick_index][i].first.push_back(pkt);
 
             DPRINTF(GPUTLB, "Coalesced re-issued req %i \
@@ -318,7 +428,7 @@ VegaTLBCoalescer::reissue_pkt_helper(PacketPtr pkt)
         std::vector<PacketPtr> new_array;
         new_array.push_back(pkt);
         coalescerFIFO[tick_index].push_back(
-            std::make_pair(new_array, VegaISA::PageBytes));
+            std::make_pair(new_array, reissue_pgsize));
 
         DPRINTF(GPUTLB,
                 "coalescerFIFO[%d] now has %d coalesced reqs after "
@@ -423,7 +533,7 @@ VegaTLBCoalescer::CpuSidePort::recvTimingReq(PacketPtr pkt)
         std::vector<PacketPtr> new_array;
         new_array.push_back(pkt);
         coalescer->coalescerFIFO[tick_index].push_back(
-            std::make_pair(new_array, coalescer->default_pgSize));
+            std::make_pair(new_array, coalescer->newGroupPageSize()));
 
         DPRINTF(GPUTLB,
                 "coalescerFIFO[%d] now has %d coalesced reqs after "
@@ -507,10 +617,12 @@ VegaTLBCoalescer::MemSidePort::recvTimingResp(PacketPtr pkt)
 void
 VegaTLBCoalescer::MemSidePort::recvReqRetry()
 {
-    // we've received a retry. Schedule a probeTLBEvent
+    coalescer->retryEvents++;
+
+    // we've receeived a retry. Schedule a probeTLBEvent
     if (!coalescer->probeTLBEvent.scheduled()) {
         coalescer->schedule(coalescer->probeTLBEvent,
-                            curTick() + coalescer->clockPeriod());
+                curTick() + coalescer->clockPeriod());
     }
 }
 
@@ -546,6 +658,8 @@ VegaTLBCoalescer::processProbeTLBEvent()
 
     if ((tlb_level == 1) && (availDownstreamSlots() == 0)) {
         DPRINTF(GPUTLB, "IssueProbeEvent - no downstream slots, bail out\n");
+        downstreamSlotBlocked++;
+        updateOccupancyStats();
         return;
     }
 
@@ -574,7 +688,7 @@ VegaTLBCoalescer::processProbeTLBEvent()
                                             iter->second[vector_index].second);
 
             // is there another outstanding request for the same page addr?
-            // consider all possible page size
+            // consider all possible page sizes already observed by the walker.
             int pending_reqs = 0;
             for (auto i_pgsize : potentialPagesize) {
                 pending_reqs += issuedTranslationsTable.count(
@@ -587,6 +701,11 @@ VegaTLBCoalescer::processProbeTLBEvent()
                         "page %#x\n",
                         virt_page_addr);
 
+                pendingBlockedProbes++;
+                pendingBlockedPackets +=
+                    iter->second[vector_index].first.size();
+                updateOccupancyStats();
+
                 ++vector_index;
                 continue;
             }
@@ -596,6 +715,9 @@ VegaTLBCoalescer::processProbeTLBEvent()
                 DPRINTF(GPUTLB, "Failed to send TLB request for page %#x",
                         virt_page_addr);
 
+                sendTimingReqFailed++;
+                updateOccupancyStats();
+
                 // No need for a retries queue since we are already
                 // buffering the coalesced request in coalescerFIFO.
                 // Arka:: No point trying to send other requests to TLB at
@@ -636,6 +758,8 @@ VegaTLBCoalescer::processProbeTLBEvent()
                 // copy coalescedReq to issuedTranslationsTable
                 issuedTranslationsTable[virt_page_addr] =
                     iter->second[vector_index].first;
+                probesIssued++;
+                updateOccupancyStats();
 
                 // erase the entry of this coalesced req
                 iter->second.erase(iter->second.begin() + vector_index);
@@ -679,14 +803,23 @@ VegaTLBCoalescer::processProbeTLBEvent()
 void
 VegaTLBCoalescer::processCleanupEvent()
 {
+    bool cleaned = false;
+
     while (!cleanupQueue.empty()) {
         Addr cleanup_addr = cleanupQueue.front();
         cleanupQueue.pop();
         issuedTranslationsTable.erase(cleanup_addr);
+        cleanupEntries++;
+        cleaned = true;
+        updateOccupancyStats();
 
         DPRINTF(GPUTLB, "Cleanup - Delete coalescer entry with key %#x\n",
                 cleanup_addr);
     }
+
+    if (cleaned && !coalescerFIFO.empty() && !probeTLBEvent.scheduled()) {
+        schedule(probeTLBEvent, cyclesToTicks(curCycle() + Cycles(1)));
+    }
 }
 
 void
@@ -709,6 +842,36 @@ VegaTLBCoalescer::regStats()
     localCycles.name(name() + ".local_cycles")
         .desc("Number of cycles spent in queue for all incoming reqs");
 
+    pendingBlockedProbes.name(name() + ".pending_blocked_probes")
+        .desc("Probe attempts blocked by an overlapping outstanding translation");
+
+    pendingBlockedPackets.name(name() + ".pending_blocked_packets")
+        .desc("Buffered packets in probe attempts blocked by an overlapping outstanding translation");
+
+    downstreamSlotBlocked.name(name() + ".downstream_slot_blocked")
+        .desc("Probe event invocations blocked by downstream slot exhaustion");
+
+    sendTimingReqFailed.name(name() + ".send_timing_req_failed")
+        .desc("TLB probe sends rejected by the downstream port");
+
+    probesIssued.name(name() + ".probes_issued")
+        .desc("Coalesced translation probes issued downstream");
+
+    cleanupEntries.name(name() + ".cleanup_entries")
+        .desc("Outstanding translation entries cleaned up after response");
+
+    retryEvents.name(name() + ".retry_events")
+        .desc("Request retry callbacks from downstream TLB port");
+
+    fifoMaxEntries.name(name() + ".fifo_max_entries")
+        .desc("Maximum number of coalesced requests buffered in the FIFO");
+
+    issuedTranslationsMax.name(name() + ".issued_translations_max")
+        .desc("Maximum number of outstanding translation table entries");
+
+    issuedTranslationPacketsMax.name(name() + ".issued_translation_packets_max")
+        .desc("Maximum number of packets represented by outstanding translations");
+
     localLatency.name(name() + ".local_latency")
         .desc("Avg. latency over all incoming pkts");
 
diff --git a/src/arch/amdgpu/vega/tlb_coalescer.hh b/src/arch/amdgpu/vega/tlb_coalescer.hh
index c44ee09c65f..4f7c912f66d 100644
--- a/src/arch/amdgpu/vega/tlb_coalescer.hh
+++ b/src/arch/amdgpu/vega/tlb_coalescer.hh
@@ -39,6 +39,7 @@
 #include <vector>
 
 #include "arch/amdgpu/vega/tlb.hh"
+#include "base/sat_counter.hh"
 #include "base/statistics.hh"
 #include "mem/port.hh"
 #include "mem/request.hh"
@@ -132,7 +133,22 @@ class VegaTLBCoalescer : public ClockedObject
     // latency of a request to be completed
     statistics::Formula latency;
 
+    // Diagnostics for outstanding-translation backpressure.
+    statistics::Scalar pendingBlockedProbes;
+    statistics::Scalar pendingBlockedPackets;
+    statistics::Scalar downstreamSlotBlocked;
+    statistics::Scalar sendTimingReqFailed;
+    statistics::Scalar probesIssued;
+    statistics::Scalar cleanupEntries;
+    statistics::Scalar retryEvents;
+    statistics::Scalar fifoMaxEntries;
+    statistics::Scalar issuedTranslationsMax;
+    statistics::Scalar issuedTranslationPacketsMax;
+
     bool canCoalesce(PacketPtr pkt1, PacketPtr pkt2, Addr pagebytes);
+    void updateOccupancyStats();
+    size_t coalescerFIFOEntries() const;
+    size_t issuedTranslationPackets() const;
     void updatePhysAddresses(PacketPtr pkt);
     void regStats() override;
 
@@ -220,11 +236,50 @@ class VegaTLBCoalescer : public ClockedObject
     /// in order to free memory and do the required clean-up
     EventFunctionWrapper cleanupEvent;
 
-    void reissue_pkt_helper(PacketPtr pkt);
+    void reissue_pkt_helper(PacketPtr pkt,
+                            Addr reissue_pgsize = VegaISA::PageBytes);
 
     Addr default_pgSize = 1ULL << 21;
     std::set<Addr> potentialPagesize;
 
+    // Per-coalescer page-size predictor.
+    // MSB set => predict 4 KiB, otherwise predict 2 MiB..
+    static constexpr unsigned SizePredBits = 4;
+    SatCounter8 sizePredictor;
+    // Update the size predictor from a returned entry's page size.
+    void updateSizePredictor(Addr returned_pgsize);
+    // The page size to speculate for a newly opened coalesced group.
+    Addr predictedPageSize() const;
+
+    // Line coalescing support at the last GPU TLB level.
+    // Uses the downstream TLB's line predictor to decide whether to open
+    // a full PWC2-line group and reissue unresolved neighbours. The GPU TLB
+    // hierarchy is currently two levels (L1, L2), so this fires at L2.
+    static constexpr unsigned LineGroupPages = 16;
+    VegaISA::GpuTLB *downstreamTLB = nullptr;
+    bool lineCoalesceEnabled() const
+    {
+        return tlb_level == 2 && downstreamTLB &&
+               downstreamTLB->shouldLineCoalesce();
+    }
+    // Assumed page granule for a freshly opened coalesced group: the predicted
+    // page size, widened to a full line when line-coalescing is enabled.
+    Addr newGroupPageSize();
+
+    // Largest granule used as an issuedTranslationsTable key. Giant walker
+    // pages (1 GiB+) must be clamped to this, else the block check collapses
+    // a whole region onto one key and serializes every request in it.
+    Addr maxCoalescePageSize() const
+    {
+        return lineCoalesceEnabled() ? default_pgSize * LineGroupPages
+                                     : default_pgSize;
+    }
+    Addr clampCoalescePageSize(Addr pg_size) const
+    {
+        Addr cap = maxCoalescePageSize();
+        return pg_size > cap ? cap : pg_size;
+    }
+
     int tlb_level;
     int maxDownstream;
     unsigned int numDownstream;

From 86308701529c39065321221e1b2f9ae84fa7423e Mon Sep 17 00:00:00 2001
From: Basem Mohammed <95645899+Basemism@users.noreply.github.com>
Date: Fri, 11 Sep 2026 10:59:53 +0200
Subject: [PATCH 3/4] style: format page-walk caching changes

---
 configs/common/GPUTLBConfig.py           |  6 +++---
 src/arch/amdgpu/vega/VegaGPUTLB.py       |  4 ++--
 src/arch/amdgpu/vega/pagetable_walker.cc |  3 ++-
 src/arch/amdgpu/vega/pagetable_walker.hh |  2 +-
 src/arch/amdgpu/vega/tlb.cc              |  6 ++++--
 src/arch/amdgpu/vega/tlb.hh              | 13 +++++++------
 src/arch/amdgpu/vega/tlb_coalescer.cc    | 18 ++++++++++++------
 7 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/configs/common/GPUTLBConfig.py b/configs/common/GPUTLBConfig.py
index 267ce7d2ce7..5a43d09442c 100644
--- a/configs/common/GPUTLBConfig.py
+++ b/configs/common/GPUTLBConfig.py
@@ -291,9 +291,9 @@ def config_tlb_hierarchy(
             l2_coalescer_index += 1
 
     # L2 <-> L3
-    # system.l2_tlb[0].mem_side_ports[0] = system.l3_coalescer[0].cpu_side_ports[
-    #    0
-    # ]
+    # system.l2_tlb[0].mem_side_ports[0] = (
+    #     system.l3_coalescer[0].cpu_side_ports[0]
+    # )
 
     # L3 TLB Vega page table walker to memory for full system only
     if full_system:
diff --git a/src/arch/amdgpu/vega/VegaGPUTLB.py b/src/arch/amdgpu/vega/VegaGPUTLB.py
index 11a2c6d7a9f..ceb3b573e35 100644
--- a/src/arch/amdgpu/vega/VegaGPUTLB.py
+++ b/src/arch/amdgpu/vega/VegaGPUTLB.py
@@ -126,6 +126,6 @@ class VegaTLBCoalescer(ClockedObject):
     disableCoalescing = Param.Bool(False, "Dispable Coalescing")
     downstream_tlb = Param.VegaGPUTLB(
         NULL,
-        "The TLB this coalescer feeds; used by the last-level coalescer to read the "
-        "line-coalescing predictor",
+        "The TLB this coalescer feeds; used by the last-level coalescer "
+        "to read the line-coalescing predictor",
     )
diff --git a/src/arch/amdgpu/vega/pagetable_walker.cc b/src/arch/amdgpu/vega/pagetable_walker.cc
index da0f24df219..2186133fce2 100644
--- a/src/arch/amdgpu/vega/pagetable_walker.cc
+++ b/src/arch/amdgpu/vega/pagetable_walker.cc
@@ -223,7 +223,8 @@ Walker::WalkerState::startWalk()
             entry.paddr = entry.pte.ppn << PageShift;
             entry.paddr += entry.vaddr & mask(entry.logBytes);
 
-            // Send translation return event. The TLB allocates the returned entry in the normal miss-return path.
+            // Send translation return event. The TLB allocates the returned
+            // entry in the normal miss-return path.
             walker->walkerResponse(this, entry, tlbPkt);
         }
     }
diff --git a/src/arch/amdgpu/vega/pagetable_walker.hh b/src/arch/amdgpu/vega/pagetable_walker.hh
index 5d34569504c..cd33e0caa1d 100644
--- a/src/arch/amdgpu/vega/pagetable_walker.hh
+++ b/src/arch/amdgpu/vega/pagetable_walker.hh
@@ -299,7 +299,7 @@ class Walker : public ClockedObject
         pwcFetchEntries(p.pwc_fetch_bytes / sizeof(uint64_t)),
         funcState(this, nullptr, true),
         enable_pwc(p.enable_pwc),
-	tlb(nullptr),
+        tlb(nullptr),
         requestorId(p.system->getRequestorId(this)),
         deviceRequestorId(999),
         system(p.system),
diff --git a/src/arch/amdgpu/vega/tlb.cc b/src/arch/amdgpu/vega/tlb.cc
index f5158f21a8d..df5f93a10c8 100644
--- a/src/arch/amdgpu/vega/tlb.cc
+++ b/src/arch/amdgpu/vega/tlb.cc
@@ -1049,8 +1049,10 @@ GpuTLB::VegaTLBStats::VegaTLBStats(statistics::Group *parent)
       ADD_STAT(inserts2M, "Number of 2 MiB entries inserted"),
       ADD_STAT(localHits4K, "Number of local hits on 4 KiB entries"),
       ADD_STAT(localHits2M, "Number of local hits on 2 MiB entries"),
-       ADD_STAT(walkerReturns4K, "Number of page walks returning 4 KiB entries"),
-       ADD_STAT(walkerReturns2M, "Number of page walks returning 2 MiB entries")
+      ADD_STAT(walkerReturns4K,
+               "Number of page walks returning 4 KiB entries"),
+      ADD_STAT(walkerReturns2M,
+               "Number of page walks returning 2 MiB entries")
 {
     localTLBMissRate = 100 * localNumTLBMisses / localNumTLBAccesses;
     globalTLBMissRate = 100 * globalNumTLBMisses / globalNumTLBAccesses;
diff --git a/src/arch/amdgpu/vega/tlb.hh b/src/arch/amdgpu/vega/tlb.hh
index 9b8dc823b56..37de5f88095 100644
--- a/src/arch/amdgpu/vega/tlb.hh
+++ b/src/arch/amdgpu/vega/tlb.hh
@@ -39,9 +39,9 @@
 
 #include "arch/amdgpu/vega/pagetable.hh"
 #include "arch/generic/mmu.hh"
+#include "base/sat_counter.hh"
 #include "base/statistics.hh"
 #include "base/trace.hh"
-#include "base/sat_counter.hh"
 #include "mem/packet.hh"
 #include "mem/port.hh"
 #include "params/VegaGPUTLB.hh"
@@ -221,11 +221,12 @@ class GpuTLB : public ClockedObject
     VegaTlbEntry *tlbLookup(const RequestPtr &req, bool update_stats);
 
     // Line-coalescing predictor: a 4-bit saturating counter,
-    // meaningful only at the L3 TLB. A real TLB miss (walk that reached memory)
-    // decrements it toward "expect an L3 miss"; a TLB hit increments it. A PWC
-    // or PWC2 hit deliberately does not count as a hit and leaves it unchanged.
-    // When the MSB is clear we expect an L3 miss and it is worthwhile for the
-    // L3 coalescer to line-coalesce (prefetch a full PWC2 line).
+    // meaningful only at the L3 TLB. A real TLB miss (walk that reached
+    // memory) decrements it toward "expect an L3 miss"; a TLB hit increments
+    // it. A PWC or PWC2 hit deliberately does not count as a hit and leaves
+    // it unchanged. When the MSB is clear we expect an L3 miss and it is
+    // worthwhile for the L3 coalescer to line-coalesce (prefetch a full PWC2
+    // line).
     static constexpr unsigned LinePredBits = 4;
     SatCounter8 lineCounter{LinePredBits, 0};
     void noteTlbHit() { lineCounter++; }
diff --git a/src/arch/amdgpu/vega/tlb_coalescer.cc b/src/arch/amdgpu/vega/tlb_coalescer.cc
index 9760876d992..9a5772f8574 100644
--- a/src/arch/amdgpu/vega/tlb_coalescer.cc
+++ b/src/arch/amdgpu/vega/tlb_coalescer.cc
@@ -313,8 +313,10 @@ VegaTLBCoalescer::updatePhysAddresses(PacketPtr pkt)
         }
 
         // Only the returned packet already has its physical address.
-        // Every other coalesced packet needs to be filled in from the returned page + offset.
-        // Use pointer identity instead of the loop index so this stays correct regardless of the packet's position in the bucket.
+        // Every other coalesced packet needs to be filled in from the
+        // returned page + offset. Use pointer identity instead of the loop
+        // index so this stays correct regardless of the packet's position in
+        // the bucket.
         if (local_pkt != pkt) {
             Addr paddr = first_entry_paddr +
                          (local_pkt_vaddr & (page_size - 1));
@@ -843,10 +845,12 @@ VegaTLBCoalescer::regStats()
         .desc("Number of cycles spent in queue for all incoming reqs");
 
     pendingBlockedProbes.name(name() + ".pending_blocked_probes")
-        .desc("Probe attempts blocked by an overlapping outstanding translation");
+        .desc("Probe attempts blocked by an overlapping outstanding "
+              "translation");
 
     pendingBlockedPackets.name(name() + ".pending_blocked_packets")
-        .desc("Buffered packets in probe attempts blocked by an overlapping outstanding translation");
+        .desc("Buffered packets in probe attempts blocked by an overlapping "
+              "outstanding translation");
 
     downstreamSlotBlocked.name(name() + ".downstream_slot_blocked")
         .desc("Probe event invocations blocked by downstream slot exhaustion");
@@ -869,8 +873,10 @@ VegaTLBCoalescer::regStats()
     issuedTranslationsMax.name(name() + ".issued_translations_max")
         .desc("Maximum number of outstanding translation table entries");
 
-    issuedTranslationPacketsMax.name(name() + ".issued_translation_packets_max")
-        .desc("Maximum number of packets represented by outstanding translations");
+    issuedTranslationPacketsMax.name(
+        name() + ".issued_translation_packets_max")
+        .desc("Maximum number of packets represented by outstanding "
+              "translations");
 
     localLatency.name(name() + ".local_latency")
         .desc("Avg. latency over all incoming pkts");

From dcff9a1ce65ee9ff06bf74f26438c67bd4c27867 Mon Sep 17 00:00:00 2001
From: Basem Mohammed <95645899+Basemism@users.noreply.github.com>
Date: Fri, 11 Sep 2026 11:02:31 +0200
Subject: [PATCH 4/4] configs: clarify final-level GPU TLB comments

---
 configs/common/GPUTLBConfig.py | 11 +++--------
 src/arch/amdgpu/vega/tlb.hh    | 11 +++++------
 2 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/configs/common/GPUTLBConfig.py b/configs/common/GPUTLBConfig.py
index 5a43d09442c..68cb6efa47e 100644
--- a/configs/common/GPUTLBConfig.py
+++ b/configs/common/GPUTLBConfig.py
@@ -214,8 +214,8 @@ def config_tlb_hierarchy(
                         system.%s_tlb[%d].cpu_side_ports[0]"
                     % (name, index, name, index)
                 )
-                # Give each Vega coalescer a handle to the TLB it feeds so the
-                # L3 coalescer can read that TLB's line-coalescing predictor.
+                # Give each Vega coalescer a handle to the TLB it feeds. The
+                # final-level coalescer uses that TLB's line predictor.
                 if full_system:
                     exec(
                         "system.%s_coalescer[%d].downstream_tlb = \
@@ -290,12 +290,7 @@ def config_tlb_hierarchy(
             )
             l2_coalescer_index += 1
 
-    # L2 <-> L3
-    # system.l2_tlb[0].mem_side_ports[0] = (
-    #     system.l3_coalescer[0].cpu_side_ports[0]
-    # )
-
-    # L3 TLB Vega page table walker to memory for full system only
+    # Final GPU-TLB page-table walkers connect to memory in full-system mode.
     if full_system:
         for TLB_type in L2:
             name = TLB_type["name"]
diff --git a/src/arch/amdgpu/vega/tlb.hh b/src/arch/amdgpu/vega/tlb.hh
index 37de5f88095..f2043ffa0bc 100644
--- a/src/arch/amdgpu/vega/tlb.hh
+++ b/src/arch/amdgpu/vega/tlb.hh
@@ -221,12 +221,11 @@ class GpuTLB : public ClockedObject
     VegaTlbEntry *tlbLookup(const RequestPtr &req, bool update_stats);
 
     // Line-coalescing predictor: a 4-bit saturating counter,
-    // meaningful only at the L3 TLB. A real TLB miss (walk that reached
-    // memory) decrements it toward "expect an L3 miss"; a TLB hit increments
-    // it. A PWC or PWC2 hit deliberately does not count as a hit and leaves
-    // it unchanged. When the MSB is clear we expect an L3 miss and it is
-    // worthwhile for the L3 coalescer to line-coalesce (prefetch a full PWC2
-    // line).
+    // meaningful only at the final GPU-TLB level. A real TLB miss (walk that
+    // reached memory) decrements it toward "expect a TLB miss"; a TLB hit
+    // increments it. A PWC or PWC2 hit deliberately does not count as a hit
+    // and leaves it unchanged. When the MSB is clear, line coalescing is
+    // worthwhile (prefetch a full PWC2 line).
     static constexpr unsigned LinePredBits = 4;
     SatCounter8 lineCounter{LinePredBits, 0};
     void noteTlbHit() { lineCounter++; }