CHAPTER 9
How can matrix multiplication reuse its inputs?
Trace one output as a row meets a column, then compute the same matrix with fewer repeated input loads.
Start with the operation
PyTorch
import torch
n = 8
a = torch.tensor([[(r + k) % 5 - 2 for k in range(n)]
for r in range(n)], dtype=torch.float32, device="cuda")
b = torch.tensor([[(k + c) % 3 - 1 for c in range(n)]
for k in range(n)], dtype=torch.float32, device="cuda")
c = a @ b
print(c.cpu())A matrix is a two-dimensional tensor. In a neural-network layer, matrix multiplication combines many inputs with learned weights. Here @ multiplies two 8 × 8 matrices. It produces a new 8 × 8 matrix, C.
The nested Python lists build one row at a time. % gives the remainder after division, producing small, repeating input values; dtype=torch.float32 stores them as FP32. A’s first row is [-2, -1, 0, 1, 2, -2, -1, 0]; B’s first column is [-1, 0, 1, -1, 0, 1, -1, 0]. You choose the matrices and operation. PyTorch selects the implementation, allocates the result and launches the GPU work; this code does not choose a tile size.
Run the examples below in your browser with no installation or GPU required. The browser simulates the calculation; it does not execute PyTorch or GPU code. To run the PyTorch snippet itself, use a Python environment with PyTorch and an NVIDIA GPU.
Idea 1 of 2
Build one output from a row and a column
Definitions
- Dot product
- Multiply matching entries of two lists and add those products.
- Accumulation
- Keep a running sum as more products are added.
- K index
- The position shared by a row of A and a column of B; it identifies which pair to multiply.
For output C[r,c], take row r of A and column c of B. Multiply A[r,k] × B[k,c] for every k, then add the products. This is a dot product. Both inputs here are 4 × 4, a smaller instance of the opening example.
The trace begins at C[0,0]. With tile size one, each step adds one product. After all four K positions, that output is complete. This trace stays with C[0,0]; the complete output matrix is available below. “Next” is an inspection order, not GPU scheduling.
Model assumptions: the small integer products and sums are exact. Each output separately loads its inputs; caches, concurrent execution and timings are absent.
Try it
Predict C[0,0] by multiplying row 0 of A by column 0 of B. Run, then step through the first four products. Change matrix dimension to 8 and explain why each output now needs more products.
Inspect the PyTorch example for these settings
PyTorch
import torch
n = 4
a = torch.tensor([[(r + k) % 5 - 2 for k in range(n)]
for r in range(n)], dtype=torch.float32, device="cuda")
b = torch.tensor([[(k + c) % 3 - 1 for c in range(n)]
for k in range(n)], dtype=torch.float32, device="cuda")
c = a @ b
print(c.cpu())Predict the result, then press Run.
Inspecting C[0,0]: K = 0. Select the next step to add the next pair to this output.
| K | A row value | B column value | Product |
|---|---|---|---|
| 0 | -2 | -1 | 2 |
| 1 | -1 | 0 | 0 |
| 2 | 0 | 1 | 0 |
| 3 | 1 | -1 | -1 |
- Selected A input
- -2
- Selected B input
- -1
- Selected product
- 2
- Partial C after this step
- 2
- Final C for this pair
- 1
Inspect the complete inputs and final C
| Row / column | 0 | 1 | 2 | 3 |
|---|---|---|---|---|
| 0 | -2 | -1 | 0 | 1 |
| 1 | -1 | 0 | 1 | 2 |
| 2 | 0 | 1 | 2 | -2 |
| 3 | 1 | 2 | -2 | -1 |
| Row / column | 0 | 1 | 2 | 3 |
|---|---|---|---|---|
| 0 | -1 | 0 | 1 | -1 |
| 1 | 0 | 1 | -1 | 0 |
| 2 | 1 | -1 | 0 | 1 |
| 3 | -1 | 0 | 1 | -1 |
| Row / column | 0 | 1 | 2 | 3 |
|---|---|---|---|---|
| 0 | 1 | -1 | 0 | 1 |
| 1 | 0 | -1 | 1 | 0 |
| 2 | 4 | -1 | -3 | 4 |
| 3 | -2 | 4 | -2 | -2 |
Other outputs use different row–column pairs, but many pairs contain the same A or B value. The next experiment keeps these inputs and results while changing how often inputs are loaded.
Idea 2 of 2
Load a tile, use it for several outputs
Definitions
- Tile
- A small rectangular part of a matrix, computed or loaded as a group.
- K phase
- One group of K positions whose products are added to the current output tile.
Choose a square output tile of width T. In one K phase, load a T × T tile of A and a T × T tile of B. The block keeps these in its shared memory, the workspace shared by its threads. Each loaded A value contributes to T output columns; each loaded B value contributes to T output rows.
The trace follows the first output tile. It retains its partial sums as later K phases arrive; the complete output matrix is available below. Input tiles can be replaced only after their users finish; a real shared-memory implementation must establish that ordering between threads. Here T = 1 is the separate-output algorithm from the first experiment.
Model assumptions: each input tile is loaded once per output tile and K phase. The count covers explicit input loads, excludes caches, and is not a count of physical memory transactions or elapsed time. The displayed PyTorch code changes with matrix dimension; tile selection changes only the browser algorithm because ordinary @ leaves implementation selection to PyTorch.
Try it
For the fixed 4 × 4 inputs, predict whether tile width 2 changes C or the input-load count. Run and follow the first two K phases. Change only tile width to 1, then 4. Compare final C and explain which outputs share each loaded input.
Inspect the PyTorch example for these settings
PyTorch
import torch
n = 4
a = torch.tensor([[(r + k) % 5 - 2 for k in range(n)]
for r in range(n)], dtype=torch.float32, device="cuda")
b = torch.tensor([[(k + c) % 3 - 1 for c in range(n)]
for k in range(n)], dtype=torch.float32, device="cuda")
c = a @ b
print(c.cpu())Predict the result, then press Run.
C tile (0,0), K 0–1. These input tiles contribute to the selected partial output tile.
| Row / column | 0 | 1 |
|---|---|---|
| 0 | -2 | -1 |
| 1 | -1 | 0 |
| Row / column | 0 | 1 |
|---|---|---|
| 0 | -1 | 0 |
| 1 | 0 | 1 |
| Row / column | 0 | 1 |
|---|---|---|
| 0 | 2 | -1 |
| 1 | 1 | 0 |
- Inputs loaded in this phase
- 8
- Outputs sharing each loaded input
- 2
- Input loads for complete C
- 64
- Separate-output input loads
- 128
The complete algorithm computes the same 4 × 4 C while changing explicit input reuse. Shared input storage for one block is 32 bytes in this model. The partial tile above includes all K phases up to this step.
Compare the complete final C
| Row / column | 0 | 1 | 2 | 3 |
|---|---|---|---|---|
| 0 | 1 | -1 | 0 | 1 |
| 1 | 0 | -1 | 1 | 0 |
| 2 | 4 | -1 | -3 | 4 |
| 3 | -2 | 4 | -2 | -2 |
For each output tile, a phase loads 2T² inputs. There are (N/T)² output tiles and N/T phases, giving 2N³/T input loads. The multiplication still computes the same N² dot products. This is a reuse argument; it does not establish how much faster a real kernel will run.
Optional: why not make every tile larger?
Shared-memory tiles and per-thread registers consume finite storage. A register is storage private to a thread. A larger tile can reuse more inputs while consuming more storage per block, leaving room for fewer resident blocks (blocks whose resources are currently allocated on an SM). Actual choices also depend on hardware, layouts and the implementation. The focused experiments above count reuse; they do not predict a resource optimum or a speedup.
Try a new case
Apply it
For the same 8 × 8 inputs, compare tile widths 1 and 4. Predict how many input values each algorithm explicitly loads and whether C changes. Explain the sharing before revealing the comparison.
Compare your answer
Run the example on your desktop
Download this chapter’s opening PyTorch example. See the complete example index and setup appendix for CPU PyTorch, the supported GPGPU-Sim programs and optional NVIDIA GPU execution.
The browser still requires no installation. Downloads are generated from the displayed source; CPU checks and simulator checks do not establish GPU performance.
References
Use these official references to check the concepts and code on this page.
- CUDA Best Practices: tiled matrix multiplication (new tab)
Follow input tiles that are reused by multiple threads.
- PyTorch: torch.matmul (new tab)
Connect the low-level example to the two-dimensional tensor operation.
These browser exercises do not run an installed PyTorch or CUDA toolkit. For optional local execution, choose documentation matching your installed tools. Links open online in a new tab; the questions and answer comparisons below work offline.
Use the documentation
Use the references to explain each answer, then apply it to this page’s example. No installation is needed.