Why AI Collapses to Two Primitives

Look at any compiled AI graph and you'll see two ops

Open a Triton kernel, an XLA HLO dump, or PyTorch's torch.compile output and skim the operations. Convolutions, attention, MLPs, RNNs, GNNs — every learned layer ultimately decomposes into combinations of two primitives:

Primitive	Math	Where you actually see it
Dot Product	`y = Σ xᵢ · wᵢ`	Attention scores, cosine similarity in vector search, output of any single neuron
Matrix Multiply (GEMM)	`C = α·A·B + β·C`	Dense / Linear layers, Q-K-V projections, MLP blocks, embedding lookups when batched

Convolutions become GEMMs after im2col or implicit-gemm. Softmax is element-wise math + a reduction. Layer normalization is two reductions and an element-wise scale. The 'real' compute, the part where 90% of FLOPs are spent, is GEMM.

That's why hardware vendors and library authors pour decades of work into one shape — fast GEMM. Tensor Cores on NVIDIA, the matrix coprocessor on Apple, AMX on M1+, AMD Matrix Cores. They're all bigger, faster GEMM accelerators bolted onto the GPU.

Once you internalize this, model design feels different. 'Should I add another head?' becomes 'how does this change my GEMM shape?' Padding, batching, hidden-size choices — they all collapse to whether you're feeding the GEMM library shapes it can vectorize cleanly.

Code

The two primitives in PyTorch — what every model is made of·python

import torch

# Dot product — single neuron output, attention score
x = torch.randn(1024)
w = torch.randn(1024)
y = torch.dot(x, w)         # scalar

# Matrix multiply — every Linear layer, Q/K/V projection, MLP
A = torch.randn(64, 1024)   # 64 sequence positions × 1024 hidden
B = torch.randn(1024, 4096) # weight matrix, hidden → ffn
C = A @ B                   # (64, 4096)

# torch.compile lowers these to fused, vendor-optimized kernels.
# Under the hood: cuBLAS on NVIDIA, MPS on Apple, oneDNN on Intel.

Same idea in MLX (Apple) — same two primitives, different runtime·python

import mlx.core as mx

x = mx.random.normal((1024,))
w = mx.random.normal((1024,))
y = mx.matmul(x, w)         # scalar via the same matmul kernel

A = mx.random.normal((64, 1024))
B = mx.random.normal((1024, 4096))
C = A @ B                   # MLX dispatches to MPSMatrixMultiplication
mx.eval(C)                  # MLX is lazy; eval forces execution.

Exercise

In any PyTorch script you have, set torch._dynamo.config.verbose = True and run with torch.compile on a small model (a 2-layer MLP is enough). Skim the trace output looking for the words 'matmul' / 'mm' / 'addmm'. Notice how often they appear vs other op names — that's why GEMM gets all the optimization budget.

Why AI Collapses to Two Primitives

Look at any compiled AI graph and you'll see two ops

Code

External links

Exercise

Progress

Comments 0