Menu
About
Articles
Papers
Resume
Speaking
Uses
Consulting
Memory Layout Optimization in PyTorch Compilation
Visualization of memory layout optimization techniques in PyTorch kernel fusion. The diagram shows how standard execution leads to memory access inefficiency while optimized layouts improve data locality and reduce memory latency.
Memory Optimization via Kernel Fusion (e.g., torch.compile)
Before Compilation: Eager Execution (Separate Kernels)
GPU Memory Hierarchy
Global Memory (High Capacity, High Latency)
L2 Cache
L1 Cache / Shared Memory
Registers (Fastest, Lowest Latency)
Kernel 1 (Linear)
Kernel 2 (ReLU)
Kernel 3 (BatchNorm)
Kernel 4 (Conv)
Load Input
Store Interm. A
Load Interm. A
Store Interm. B
Load Interm. B
Store Interm. C
Load Interm. C
Store Output
Bottleneck: Frequent Global Memory reads/writes for intermediate results.
Bottleneck: High overhead from launching many small kernels.
After Compilation: Optimized Execution (Fused Kernels)
GPU Memory Hierarchy
Global Memory
L2 Cache
L1 Cache / Shared Memory
Registers
Fused Kernel A
(Linear)
(ReLU)
Fused Kernel B
(BatchNorm)
(Conv)
Load Input
Interm. A'
(in Regs/L1)
Store Interm. X
Load Interm. X
Interm. B'
(in Regs/L1)
Store Output
Benefit: Intermediate results (A', B') kept in fast memory, avoiding Global Memory writes/reads.
Benefit: Fewer kernel launches (2 instead of 4), reducing CPU-GPU synchronization overhead.
Mastodon
