GPU Lab / interactive book
GPU mapSource ↗

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())
Desktop setup

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())
Desktop setup

Predict the result, then press Run.

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())
Desktop setup

Predict the result, then press Run.

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
Tile width 1 loads 1024 input values; tile width 4 loads 256. Each loaded input can contribute to 4 outputs in a phase. Both produce the same C, including C[0,0] = 0. This comparison assumes no cache reuse and predicts no elapsed time.

References

Use these official references to check the concepts and code on this page.

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.

Install GPU Lab

Give GPU Lab its own icon and window. All lessons and simulators work offline once the course is saved.

Use your browser’s install option:

If your browser has no install option, try Chrome or Edge.

App status

Version

Checking for updates

Offline access

Not saved for offline use

Save all lessons, browser simulators, diagrams and example files. Desktop execution needs separate PyTorch or CUDA/GPGPU-Sim tools; their installations are not part of this offline download. External links need internet access.