Tensors and Matmul

The idea in one minute#

A tensor is a block of numbers with a shape. In memory it is always the same thing: one flat []float32 plus the shape that says how to interpret it. Element [i][j][k] of a tensor of shape (A, B, C) lives at flat index i·B·C + j·C + k. Reshaping, transposing and slicing change the bookkeeping, not the data.

Almost all the time in a neural network goes into one operation, matrix multiplication. A naive triple loop in Go manages a GFLOP/s or two. The same arithmetic with the loops reordered, the inner loop unrolled, and rows split across goroutines is ten or more times faster on an ordinary multi-core machine — and that gap is modules III–V applied.

An analogy#

A cinema’s seats. Physically there is one long list of seats in the booking system. “Row 7, seat 12” is just a way of computing a position in that list: 7 × seats_per_row + 12. Renumbering the rows, or deciding to read the hall column by column, changes the formula — not where the seats are bolted down.

A picture#

flowchart TB
  subgraph T["Tensor: shape (2, 3), strides (3, 1)"]
    direction LR
    SH["Shape []int{2, 3}"]
    DT["Data []float32:  a  b  c  d  e  f"]
  end
  DT --> V1["view as (2,3):<br/>a b c<br/>d e f"]
  DT --> V2["reshape to (3,2): same data<br/>a b<br/>c d<br/>e f"]
  DT --> V3["transpose: strides (1, 3)<br/>a d<br/>b e<br/>c f"]
  subgraph MM["C = A x B, the inner loop"]
    direction LR
    IJK["i, j, k order:<br/>walks B DOWN a column<br/>one cache miss per step"]
    IKJ["i, k, j order:<br/>walks B and C ALONG rows<br/>sequential, prefetched"]
  end
  class SH neutral
  class DT memory
  class V1,V2,V3 compute
  class IJK warn
  class IKJ queue

How it really works#

Storage#

type Tensor struct {
    Shape []int
    Data  []float32      // row-major: the last index varies fastest
}

• float32, not float64: half the memory and bandwidth, and it is what models use.
• Row-major order (as in C and NumPy): a matrix’s rows are contiguous.
• Strides say how far to step in Data per index: for shape (A, B, C) they are (B·C, C, 1). A transpose can be expressed by permuting strides without moving data; so can slices and broadcasts. Such views are free to create and can be slow to traverse (they break sequential access), so libraries materialize a contiguous copy before heavy math.
• No pointers inside: a gigabyte of weights costs the garbage collector nothing (III.04).

Shapes you will meet#

Broadcasting stretches a smaller shape to match a larger one — adding a bias of shape (d) to activations of shape (batch, seq, d). In hand-written Go it is simply an inner loop over the last dimension.

Matrix multiplication#

For A (m×k) and B (k×n): C[i][j] = Σₚ A[i][p] · B[p][j]. That is 2·m·n·k floating-point operations. The work is fixed; the speed is not.

Steps 1–4 are a few dozen lines of Go: typically several times faster on one core, and well over ten times with the cores of a laptop. Steps 5–6 are another large factor, and are why production systems call tuned libraries or GPUs: a modern CPU core reaches tens of GFLOP/s with AVX-512, a data-center GPU thousands.

Matrix times vector#

Generating one token multiplies each weight matrix by a single vector (GEMV). Every weight is read once and used for one multiply and one add: about a quarter of a FLOP per byte. No reordering can make that compute-bound; it runs at the speed memory can be read. This is the whole reason batching matters in LLM serving — a batch turns matrix-vector into matrix-matrix and reuses each weight many times per read (Inference Engineering I.08).

Allocation discipline#

The rules from III.05, in tensor form:

• Every operation takes its destination: MatMul(dst, a, b), not c := MatMul(a, b).
• Allocate activations once per model or per worker and reuse them for every token.
• In-place where the math allows: activations, normalization, residual adds.

Numerical notes#

• float32 carries about 7 decimal digits. Summing many values in a different order gives a slightly different result, so a parallel or unrolled sum will not match a sequential one bit for bit. Compare with a tolerance, never with ==.
• Go does not fuse multiply-add unless you write math.FMA (float64) — results are reproducible across platforms at a small cost in speed.

Code#

// matmul.go — a tensor, and matrix multiplication from naive to parallel, with GFLOP/s.
package main

import (
	"fmt"
	"math"
	"math/rand"
	"runtime"
	"sync"
	"time"
)

type Tensor struct {
	Shape []int
	Data  []float32
}

func New(shape ...int) *Tensor {
	n := 1
	for _, s := range shape {
		n *= s
	}
	return &Tensor{Shape: shape, Data: make([]float32, n)}
}

// At computes the flat index from row-major strides.
func (t *Tensor) At(idx ...int) float32 {
	off, stride := 0, 1
	for d := len(t.Shape) - 1; d >= 0; d-- {
		off += idx[d] * stride
		stride *= t.Shape[d]
	}
	return t.Data[off]
}

// 0. The textbook version on a slice of slices.
func matmulNaive(a, b [][]float32, m, n, k int) [][]float32 {
	c := make([][]float32, m)
	for i := range c {
		c[i] = make([]float32, n)
		for j := 0; j < n; j++ {
			var s float32
			for p := 0; p < k; p++ {
				s += a[i][p] * b[p][j] // b[p][j]: a different row every step
			}
			c[i][j] = s
		}
	}
	return c
}

// rowKernel computes rows [lo, hi) of C = A x B in i,p,j order with a 4-way unrolled axpy.
func rowKernel(dst, a, b []float32, n, k, lo, hi int) {
	for i := lo; i < hi; i++ {
		ci := dst[i*n : (i+1)*n]
		for j := range ci {
			ci[j] = 0
		}
		ai := a[i*k : (i+1)*k]
		for p, av := range ai {
			bp := b[p*n : (p+1)*n]
			bp = bp[:len(ci)] // bounds-check hint
			j := 0
			for ; j+4 <= len(ci); j += 4 {
				ci[j] += av * bp[j]
				ci[j+1] += av * bp[j+1]
				ci[j+2] += av * bp[j+2]
				ci[j+3] += av * bp[j+3]
			}
			for ; j < len(ci); j++ {
				ci[j] += av * bp[j]
			}
		}
	}
}

// MatMul writes A (m x k) times B (k x n) into dst (m x n), using `workers` goroutines.
func MatMul(dst, a, b *Tensor, workers int) {
	m, k, n := a.Shape[0], a.Shape[1], b.Shape[1]
	if workers <= 1 {
		rowKernel(dst.Data, a.Data, b.Data, n, k, 0, m)
		return
	}
	var wg sync.WaitGroup
	chunk := (m + workers - 1) / workers
	for lo := 0; lo < m; lo += chunk {
		hi := min(lo+chunk, m)
		wg.Add(1)
		go func() { // each goroutine writes its own rows of dst: no sharing, no lock
			defer wg.Done()
			rowKernel(dst.Data, a.Data, b.Data, n, k, lo, hi)
		}()
	}
	wg.Wait()
}

func main() {
	t := New(2, 3, 4)
	for i := range t.Data {
		t.Data[i] = float32(i)
	}
	fmt.Println("shape", t.Shape, "element [1][2][3] =", t.At(1, 2, 3), "(flat index 1*12 + 2*4 + 3 = 23)")

	const n = 384
	rng := rand.New(rand.NewSource(1))
	a, b, c := New(n, n), New(n, n), New(n, n)
	a2, b2 := make([][]float32, n), make([][]float32, n)
	for i := 0; i < n; i++ {
		a2[i], b2[i] = make([]float32, n), make([]float32, n)
		for j := 0; j < n; j++ {
			x, y := rng.Float32()-0.5, rng.Float32()-0.5
			a.Data[i*n+j], b.Data[i*n+j] = x, y
			a2[i][j], b2[i][j] = x, y
		}
	}
	flops := 2 * float64(n) * float64(n) * float64(n)
	report := func(name string, f func()) {
		best := time.Hour
		for r := 0; r < 3; r++ {
			start := time.Now()
			f()
			if d := time.Since(start); d < best {
				best = d
			}
		}
		fmt.Printf("%-34s %7.1f ms  %6.2f GFLOP/s\n", name, float64(best.Microseconds())/1000, flops/best.Seconds()/1e9)
	}

	fmt.Printf("\n%dx%d matrix multiply\n", n, n)
	var ref [][]float32
	report("naive i,j,p on [][]float32", func() { ref = matmulNaive(a2, b2, n, n, n) })
	report("flat, i,p,j, unrolled, 1 thread", func() { MatMul(c, a, b, 1) })
	workers := runtime.GOMAXPROCS(0)
	report(fmt.Sprintf("same, %d goroutines", workers), func() { MatMul(c, a, b, workers) })

	// Same result? Compare with a tolerance: summation order differs.
	maxDiff := 0.0
	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			maxDiff = math.Max(maxDiff, math.Abs(float64(ref[i][j]-c.Data[i*n+j])))
		}
	}
	fmt.Printf("largest difference between versions: %.2e\n", maxDiff)
	fmt.Println("(zero here: both add the k products in the same order; reorder the sum and it would differ by rounding)")
}

Remember this#

• A tensor is a flat []float32 plus a shape; the index is computed from row-major strides.
• Matrix multiply is 2·m·n·k FLOPs; loop order, unrolling and parallel rows give 10× or more in pure Go.
• Matrix-vector work (one token at a time) is memory-bound; batching is what fixes it.
• Operations write into a destination you own. Compare floats with a tolerance.

Try it#

1. Run matmul.go. What speed-up did each step give? Is the parallel speed-up close to your core count? If not, what is the limit (V.03)?
2. Add blocking: process j in tiles of 64 columns. Does it help at 384? At 1,536?
3. Implement MatVec(dst, w, x) and measure GFLOP/s for a 4,096 × 4,096 matrix. Compare with the matmul figure and explain the difference.

Check yourself#

1. How is the flat index of element [i][j][k] computed?
2. Why is the i, p, j loop order faster than i, j, p?
3. Why can parallel workers write the output without a lock?