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: 10-100x faster for row-major!

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: 16 bytes
# Moving to next column: 4 bytes

# 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: ~50-100 cycles
# Remote memory: ~200-400 cycles (4-8x slower)

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: larger batches typically more efficient (better cache usage)

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