[FEA]AST-Aware Unified Memory Prefetching to Mitigate Page Fault Storms
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
### Problem Description: The Page Fault Bottleneck
In Unified Memory (UM) scenarios, `cudf` operations on datasets larger than VRAM often suffer from severe latency due to **"Page Fault Storms"**. When a kernel accesses host-resident pages, the GPU SMs (Streaming Multiprocessors) stall, waiting for the interconnect (PCIe/NVLink) to serve the page migration.
Currently, `cudf` relies heavily on the driver's heuristic or explicit `rmm` pooling. However, the library typically possesses semantic knowledge of future access patterns (via the AST/computation graph) that the low-level CUDA driver lacks.
### Proposed Architecture: Semantic Prefetching Layer
I propose investigating a **Semantic Prefetching Layer** that analyzes the `cudf` Expression Graph before kernel launch to issue proactive memory hints.
Instead of relying on demand-paging, we can utilize:
1. **`cudaMemPrefetchAsync`**: To move columns involved in the *next* operation to the GPU immediately, overlapping with the current CPU-side query planning.
2. **`cudaMemAdvise` Hints**:
* Mark join keys as `cudaMemAdviseSetPreferredLocation` (GPU).
* Mark scan-only columns as `cudaMemAdviseSetReadMostly`.
Impact & Justification
Latency Masking: By issuing prefetches during the query planning phase (CPU time), we hide the PCIe transfer latency.
TLB Efficiency: Reducing on-the-fly page migrations reduces pressure on the GPU MMU (Memory Management Unit) and TLB misses.
Out-of-Core Scaling: Enables smoother degradation for workloads slightly larger than VRAM, preventing hard stalls.
I believe this aligns with the roadmap for better Managed Memory support in libcudf.
#### Technical Implementation Concept
We can introduce a `MemoryOptimizer` visitor that walks the `libcudf` AST.
```cpp
// Pseudo-code for AST-Aware Prefetcher
void optimize_memory_placement(const aggregation_request& req, cudaStream_t stream) {
// 1. Identify "Hot" Columns (e.g., GroupBy Keys)
for (auto& col : req.keys) {
// We know these pages will be accessed randomly and heavily.
// Force prefetch to VRAM to avoid TLB Thrashing.
cudaMemPrefetchAsync(col.data(), col.size(), current_device, stream);
// Hint the driver that this data should stay on GPU
cudaMemAdvise(col.data(), col.size(), cudaMemAdviseSetPreferredLocation, current_device);
}
// 2. Identify "Cold/Streaming" Columns (e.g., Value columns in simple sum)
for (auto& col : req.values) {
// These might be accessed linearly. Relax residency requirements if VRAM is tight.
if (is_vram_pressure_high()) {
cudaMemAdvise(col.data(), col.size(), cudaMemAdviseSetAccessedBy, current_device);
}
}
}
Contributor guide
Assessment
This issue has not been assessed yet.