Module 17: Acceleration

A Python matrix multiplication can spend much of its runtime interpreting loops. This module compares BLAS-backed multiplication with explicit tiling, compares two NumPy GELU implementations, and rewrites Module 09’s convolution loops as one matrix multiply. These experiments expose execution overhead and memory reuse, while production kernel fusion provides a comparison beyond the NumPy implementation.

NoteModule Info

OPTIMIZATION TIER | Difficulty: ●●●○ | Time: 5-7 hours | Prerequisites: 01-14

Prerequisites: Modules 01-14 means you need:

  • Tensor operations (Module 01) for understanding data structures
  • Neural network layers (Module 03) for knowing what to accelerate
  • The Conv2d layer (Module 09), which your im2col convolution must match
  • Training loops (Module 08) for understanding the performance context
  • Profiling tools (Module 14) for measuring acceleration gains

If you can multiply matrices and understand why matrix multiplication is expensive, you’re ready.

Audio overview (AI-generated)
Open in Binder → Runs in your browser with nothing to install; the session is discarded when you leave, so download your notebook to keep it.
Lecture slides AI-generated · opens an in-page viewer
🔥 Slide Deck · AI-generated
1 / -
Loading slides...
Download Fullscreen

Overview

Matrix multiplication and element-wise activations stress different parts of a processor. Matrix multiplication can reuse each input many times; activations often do little arithmetic per array element.

You will implement vectorized and tiled multiplication, then compare grouped and staged GELU expressions. NumPy still creates intermediate arrays in the grouped expression. Last, you will lower a convolution to a single matrix multiply with im2col, write its backward pass with col2im so it can train, and measure what that costs in memory. Measure each implementation on your hardware; this module does not guarantee a particular speedup.

Commands

# first time
tito module start 17

# later sessions
tito module resume 17

# when your tests pass
tito module complete 17

Your notebook is modules/17_acceleration/acceleration.ipynb.

Learning objectives

TipBy completing this module, you will:
  • Implement a cache-aware tiled matrix multiplication and measure it against the supplied version that calls optimized BLAS libraries for maximum throughput
  • Distinguish grouped NumPy expressions from native kernel fusion that can avoid intermediate arrays
  • Lower a convolution to one matrix multiply with im2col, write its backward pass with col2im, and measure its speed against Module 09’s loops and its memory cost
  • Understand the roofline model and arithmetic intensity to predict performance bottlenecks
  • Analyze production acceleration strategies for different deployment scenarios (edge, cloud, GPU)

What you’ll build

Figure 1: NumPy acceleration experiments. Compare a BLAS-backed matrix multiply, explicit tiled multiplication, and staged or grouped GELU expressions. Check numerical agreement before timing; NumPy expressions can allocate temporaries and do not guarantee a single fused kernel.

The pattern you’ll enable:

# Fast matrix operations using BLAS
output = vectorized_matmul(x, weights)  # Compare against Python loops

# Grouped activation expression
activated = fused_gelu(output)  # NumPy still creates temporary arrays

# Convolution as one matrix multiply
features = im2col_conv2d(images, conv.weight, conv.bias, padding=1)  # same output as conv(images)

# ...and one you can train through
out = Im2colConv2dFunction.apply(images, conv.weight, conv.bias, stride=1, padding=1)
loss.backward()  # gradients match Conv2d's, via two matmuls and col2im

What you’re not building yet

To keep this module focused, you will not implement:

  • GPU kernels (that requires CUDA programming, covered in production frameworks)
  • Custom CPU assembly (BLAS libraries already provide this)
  • Automatic kernel fusion (compilers like XLA do this automatically)
  • Multi-threading control (NumPy handles this via OpenBLAS/MKL)

You are building the understanding. Hardware-specific implementations come later.

What you write

The notebook arrives with the surrounding code already written and explained. You write 7 functions, each marked # YOUR CODE HERE and followed by a test cell:

fused_gelu
Compact GELU expression that avoids retaining intermediate Tensor copies.
tiled_matmul
Cache-aware matrix multiplication using tiling (also called blocking).
im2col
Unroll every convolution patch of a batch of images into one row of a matrix.
im2col_conv2d
2D convolution computed as one matrix multiply over an im2col patch matrix.
col2im
Add every row of a patch-gradient matrix back into the image it came from.
Im2colConv2dFunction.forward
Convolve with one matrix multiply and save what backward needs.
Im2colConv2dFunction.backward
Turn the output gradient into gradients for the input, weight, and bias.

How you know it works

tito module complete stops at the first step that fails:

  1. the unit tests inside your notebook run;
  2. your code is exported into tinytorch.perf.acceleration;
  3. the integration tests run against that exported package, together with the modules before it;
  4. the module is recorded as done, and tito module status shows it.

Unit tests in your notebook (8). Each prints a ✅ line when it passes.

  • Vectorized Matrix Multiplication
  • Fused GELU
  • Kernel Fusion Performance Impact
  • Tiled Matrix Multiplication
  • im2col
  • im2col Convolution
  • col2im
  • im2col Convolution Gradients

Integration tests after export (16).

  • tests/17_acceleration/test_acceleration_core.py
  • tests/17_acceleration/test_acceleration_integration.py

Completing this module unlocks Milestone 07, Custom Kernels (2024) (tito milestone run 07).

When it fails

A bare NotImplementedError with no message means a cell reached a function you have not written yet: the notebook ships each one as # YOUR CODE HERE followed by raise NotImplementedError(). The messages below are ones this module actually prints when an implementation is present but wrong.

Tiled and vectorized results should match
From test_unit_tiled_matmul. Each tile product overwrote the output block. Accumulate across the k tiles: C[i0:i1, j0:j1] += A_tile @ B_tile.
Columns must be ordered (channel, kernel row, kernel column)
From test_unit_im2col. Your patches hold the right numbers in the wrong order, usually because the channel axis ended up last. Transpose the (N, C, k, k, out_h, out_w) array with (0, 4, 5, 1, 2, 3) so each row is flattened in the same order as weight.reshape(out_ch, -1).
Differs from Conv2d by up to ...
From test_unit_im2col_conv2d. The shapes worked out but the values did not. Check that the flattened weight is transposed to (C·k·k, out_ch) before the multiply, and that the result is reshaped to (N, out_h, out_w, out_ch) and then transposed to put the channels second.
Overlapping patches must add up
From test_unit_col2im. Your col2im assigns each patch gradient with = instead of adding it with +=, so a pixel covered by several patches keeps only the last one’s gradient.
weight gradient differs from Conv2d by up to ...
From test_unit_im2col_conv2d_function. self.cols.T @ G has shape (C·k·k, out_ch); transpose it before reshaping to (out_ch, C, k, k). If the input gradient is the one that differs, check that grad_output was moved to (N, out_h, out_w, out_ch) before flattening.

Finished? Read why

The reasoning behind this module (why it is built this way, what it costs, and how production frameworks differ) is the chapter Acceleration: Kernel Fusion and Cache Tiling in the companion book, TinyTorch: From Tensors to Transformers (PDF). The book prints complete reference implementations, so read it after you finish the module, not while you are working on it.

What’s next

You have compared implementations of the same operations. Next, reuse results to avoid repeated computation.

NoteUp next: Module 18, Memoization

Acceleration shrinks the cost of every call. Memoization eliminates the call. Module 18 builds KV caching for transformer generation, so each new token reuses the keys and values of the tokens before it. The two techniques compose: a fused kernel that gets called zero times is the cheapest kernel of all.

Next: Module 18: Memoization

How later modules use this one

Table 1: How acceleration composes with memoization, benchmarking, and capstone.
Module What It Does Your Acceleration In Action
18: Memoization Cache repeated computations Vectorized operations complement KV reuse
19: Benchmarking Systematic performance measurement Benchmark(...).run_latency_benchmark() times the models you accelerated
20: Capstone Complete optimized model Acceleration throughout model pipeline
Back to top