Memory Layout & Cache Efficiency¶
Overview¶
How data is organized in memory dramatically affects performance:
- Contiguous vs scattered: Linear scan is 10-100x faster than random access
- Cache lines: 64-byte chunks loaded together from RAM
- NUMA: Multi-socket systems have local vs remote memory
- Impact on ML: Matrix multiply can be 10x slower with bad memory layout
Optimizing memory layout is one of the biggest performance gains available.
-
Memory Hierarchy¶
CPU Cache Structure¶
- ┌─ L1 Cache (32KB, ~4 cycles)
- L2 Cache (256KB, ~10 cycles)
- L3 Cache (8MB, ~40 cycles)
- RAM (16GB+, ~200 cycles)
- Disk (SSD/HDD, ~10M+ cycles)
Each cache miss is 1-1000x slower depending on level!
Cache Lines and Prefetching¶
import numpy as np
import time
# Contiguous (cache-friendly)
arr_row = np.zeros((1000, 1000))
start = time.time()
for i in range(1000):
for j in range(1000):
arr_row[i, j] += 1 # Sequential access - cache prefetch works!
row_time = time.time() - start
# Non-contiguous (cache-unfriendly)
arr_col = np.zeros((1000, 1000))
start = time.time()
for j in range(1000):
for i in range(1000):
arr_col[i, j] += 1 # Random memory access - cache misses!
col_time = time.time() - start
print(f"Row-major: {row_time:.3f}s")
print(f"Col-major: {col_time:.3f}s")
print(f"Speedup: {col_time / row_time:.1f}x")
# Typical result
Contiguity in NumPy and PyTorch¶
NumPy Flags¶
import numpy as np
# C-contiguous (row-major) - default
arr_c = np.zeros((1000, 1000), order='C')
print(f"C-contiguous: {arr_c.flags['C_CONTIGUOUS']}") # True
print(f"F-contiguous: {arr_c.flags['F_CONTIGUOUS']}") # False
# Fortran-contiguous (column-major)
arr_f = np.zeros((1000, 1000), order='F')
print(f"C-contiguous: {arr_f.flags['C_CONTIGUOUS']}") # False
print(f"F-contiguous: {arr_f.flags['F_CONTIGUOUS']}") # True
# Neither (after transpose or slice)
arr_t = arr_c.T
print(f"C-contiguous: {arr_t.flags['C_CONTIGUOUS']}") # False
print(f"F-contiguous: {arr_t.flags['F_CONTIGUOUS']}") # False (unless by luck)
# Fix non-contiguous
arr_t_fixed = np.ascontiguousarray(arr_t)
print(f"C-contiguous after fix: {arr_t_fixed.flags['C_CONTIGUOUS']}") # True
PyTorch Contiguity¶
import torch
import time
# Create tensor
x = torch.randn(1000, 1000)
print(f"Is contiguous: {x.is_contiguous()}") # True
# Transpose breaks contiguity
y = x.T
print(f"Transposed is contiguous: {y.is_contiguous()}") # False
# Some operations require contiguous
try:
# May fail or be slow on non-contiguous
z = torch.matmul(y, y) # Inefficient!
except:
pass
# Fix by making contiguous
y_cont = y.contiguous()
z = torch.matmul(y_cont, y_cont) # Fast!
# Benchmark
x = torch.randn(1000, 1000, device='cuda')
y = x.T
# Non-contiguous multiply (may be slow)
start = time.time()
for _ in range(100):
_ = torch.matmul(y, y)
nc_time = time.time() - start
# Contiguous multiply (fast)
y_cont = y.contiguous()
start = time.time()
for _ in range(100):
_ = torch.matmul(y_cont, y_cont)
c_time = time.time() - start
print(f"Non-contiguous: {nc_time:.3f}s")
print(f"Contiguous: {c_time:.3f}s")
Stride and Memory Layout¶
Understanding Strides¶
import numpy as np
# 2D array
arr = np.arange(12).reshape(3, 4)
print(arr)
# [[ 0 1 2 3]
# [ 4 5 6 7]
# [ 8 9 10 11]]
# Strides = bytes to move to next element
print(f"Strides: {arr.strides}") # (16, 4) = (4 elements * 4 bytes, 1 element * 4 bytes)
# Moving to next row
# Moving to next column
# Transpose has different strides
arr_t = arr.T
print(f"Transposed strides: {arr_t.strides}") # (4, 16) - reversed!
# Slice creates unusual strides
arr_slice = arr[::2,::2] # Every other element
print(f"Sliced strides: {arr_slice.strides}") # (32, 8) - skipped stride!
PyTorch Stride Details¶
import torch
# 2D tensor
t = torch.arange(12).reshape(3, 4)
print(f"Strides: {t.stride()}") # (4, 1)
# Transpose
t_t = t.T
print(f"Transposed strides: {t_t.stride()}") # (1, 4)
# Slice (view doesn't copy data)
t_slice = t[::2,::2]
print(f"Sliced strides: {t_slice.stride()}") # (8, 2)
# Contiguous tensors have stride (rows*cols, 1) for row-major
# Non-contiguous have other patterns
NUMA and Multi-Socket Systems¶
NUMA Awareness¶
import numpy as np
import torch
# On NUMA systems (multiple sockets), accessing remote memory is slow
# Local memory
# Remote memory
def allocate_local():
"""Allocate on local socket."""
# NumPy allocation goes to whatever socket allocated it
arr = np.zeros((10000, 10000))
return arr
def allocate_remote():
"""Simulate remote allocation by accessing from different socket."""
# Would need numactl or hwloc to actually control this
arr = np.zeros((10000, 10000))
return arr
# For optimal performance on NUMA:
# 1. Allocate data on local socket
# 2. Process on same socket
# 3. Use numactl to bind process to socket (Linux only)
Thread Affinity¶
import os
import numpy as np
import threading
# Bind thread to specific CPU cores (reduces context switches)
def set_thread_affinity(cpu_ids):
"""Bind thread to specific CPUs (Linux only)."""
os.sched_setaffinity(0, cpu_ids)
def thread_worker(worker_id, data):
"""Worker bound to specific CPU."""
# Bind to CPU cores (e.g., cores 0-3 for worker 0)
cores_per_worker = 4
cores = list(range(worker_id * cores_per_worker,
(worker_id + 1) * cores_per_worker))
try:
set_thread_affinity(set(cores))
except:
pass # Not supported on all systems
# Process data (stays on same cache, better performance)
result = np.sum(data)
return result
# Multi-threaded processing with affinity
num_workers = 4
threads = []
data = np.random.randn(1000000)
for i in range(num_workers):
t = threading.Thread(target=thread_worker, args=(i, data))
t.start()
threads.append(t)
for t in threads:
t.join()
-
Optimizing Matrix Multiply¶
Memory Layout Impact¶
import numpy as np
import time
def benchmark_matmul(size=1000, order='C'):
"""Benchmark matrix multiply with different layouts."""
# Create matrices
if order == 'C':
a = np.zeros((size, size), order='C')
b = np.zeros((size, size), order='C')
else:
a = np.zeros((size, size), order='F')
b = np.zeros((size, size), order='F')
# Fill with values
a[:] = np.random.randn(size, size)
b[:] = np.random.randn(size, size)
# Benchmark
start = time.time()
c = a @ b
elapsed = time.time() - start
return elapsed
# Benchmark different layouts
c_time = benchmark_matmul(1000, 'C')
f_time = benchmark_matmul(1000, 'F')
print(f"C-order: {c_time:.3f}s")
print(f"F-order: {f_time:.3f}s")
print(f"Speedup: {f_time / c_time:.1f}x")
# C-order typically faster (CPU cache optimization)
GPU Memory Layout¶
import torch
import time
# GPU memory is optimized for certain layouts
x_gpu = torch.randn(1000, 1000, device='cuda')
y_gpu = torch.randn(1000, 1000, device='cuda')
# Contiguous (fast)
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(100):
_ = torch.matmul(x_gpu, y_gpu)
end.record()
torch.cuda.synchronize()
contiguous_time = start.elapsed_time(end)
# Non-contiguous (slower)
x_nc = x_gpu.T.T # Transpose twice to break contiguity
start.record()
for _ in range(100):
_ = torch.matmul(x_nc, x_nc)
end.record()
torch.cuda.synchronize()
noncontiguous_time = start.elapsed_time(end)
print(f"Contiguous: {contiguous_time:.1f}ms")
print(f"Non-contiguous: {noncontiguous_time:.1f}ms")
print(f"Speedup: {noncontiguous_time / contiguous_time:.1f}x")
-
Batch Processing and Memory Locality¶
Batch Size Impact¶
import torch
import time
model = torch.nn.Linear(1000, 500)
model.eval()
def benchmark_batch(batch_size):
"""Benchmark inference with different batch sizes."""
x = torch.randn(batch_size, 1000)
start = time.time()
for _ in range(100):
with torch.no_grad():
y = model(x)
elapsed = time.time() - start
# Throughput: samples per second
throughput = (batch_size * 100) / elapsed
return throughput
# Benchmark different batch sizes
for batch_size in [1, 8, 32, 128, 512]:
throughput = benchmark_batch(batch_size)
print(f"Batch {batch_size:3d}: {throughput:.0f} samples/sec")
# Result
Prefetching Pattern¶
import torch
import torch.nn as nn
class PrefetchingDataLoader:
"""Prefetch next batch while GPU processes current batch."""
def __init__(self, data_loader, device='cuda'):
self.data_loader = data_loader
self.device = device
def __iter__(self):
# Prefetch first batch
try:
next_batch_x, next_batch_y = next(iter(self.data_loader))
next_batch_x = next_batch_x.to(self.device, non_blocking=True)
next_batch_y = next_batch_y.to(self.device, non_blocking=True)
except StopIteration:
return
# Iterate with prefetching
for batch_x, batch_y in self.data_loader:
# Process previous batch while prefetching
batch_x = next_batch_x
batch_y = next_batch_y
# Prefetch next batch
next_batch_x, next_batch_y = next(iter(self.data_loader))
next_batch_x = next_batch_x.to(self.device, non_blocking=True)
next_batch_y = next_batch_y.to(self.device, non_blocking=True)
yield batch_x, batch_y
# Usage
# for batch_x, batch_y in PrefetchingDataLoader(train_loader):
# predictions = model(batch_x)
Optimization Checklist¶
def optimize_memory_layout(tensor_or_array):
"""Check and optimize memory layout."""
# For PyTorch
if isinstance(tensor_or_array, torch.Tensor):
if not tensor_or_array.is_contiguous():
# Make contiguous
tensor_or_array = tensor_or_array.contiguous()
# Check stride pattern
print(f"Strides: {tensor_or_array.stride()}")
# Check on optimal device
print(f"Device: {tensor_or_array.device}")
# For NumPy
else:
if not tensor_or_array.flags['C_CONTIGUOUS']:
# Make C-contiguous
tensor_or_array = np.ascontiguousarray(tensor_or_array)
# Check stride
print(f"Strides: {tensor_or_array.strides}")
print(f"Flags: C={tensor_or_array.flags['C_CONTIGUOUS']}, F={tensor_or_array.flags['F_CONTIGUOUS']}")
return tensor_or_array
Summary: Memory Optimization Rules¶
| Pattern | Impact | Implementation |
|---|---|---|
| Contiguity | 10-100x | .contiguous(), avoid transpose |
| Cache locality | 5-10x | Sequential access, batch ops |
| Stride optimization | 2-5x | Avoid slicing, reshape instead |
| Batch size | 2-4x | Process multiple samples |
| NUMA awareness | 2-8x | numactl (Linux) |
| Prefetching | 1.5-2x | Overlap loading/compute |
-
Related Topics¶
- [01 Reference Counting & Garbage Collection](/05-py3/05-memory-&-performance/(01-reference-counting-garbage-collection/) - Memory allocation
- 03 Global Interpreter Lock (Gil) - Threading and memory
- 04 Multithreading Vs Multiprocessing - Shared memory
- [06 Memory Profiling & Optimization](/05-py3/05-memory-&-performance/(06-memory-profiling-optimization/) - Profiling tools