Skip to content

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

-

  • [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