Mojo Matrix Operations
Matrices are two-dimensional grids of numbers. They appear everywhere in AI — weight tables in neural networks, image pixel grids, transformation operations in graphics, and covariance tables in statistics. Mojo lets you implement matrix operations with full control over memory layout and hardware utilization, reaching speeds that match specialized libraries.
Matrix Representation
A 3×4 matrix has 3 rows and 4 columns:
col0 col1 col2 col3
row0 [ 1 2 3 4 ]
row1 [ 5 6 7 8 ]
row2 [ 9 10 11 12 ]
Row-major storage in memory (C order):
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
←── row0 ──→←── row1 ──→←──── row2 ────→
Element [row, col] is at index: row × num_cols + col
Element [1, 2] = index 1×4 + 2 = 6 → value 7 ✓
Matrix Struct
from memory import UnsafePointer
struct Matrix:
var data: UnsafePointer[Float32]
var rows: Int
var cols: Int
fn __init__(inout self, rows: Int, cols: Int):
self.rows = rows
self.cols = cols
self.data = UnsafePointer[Float32].alloc(rows * cols)
for i in range(rows * cols):
self.data.init_pointee_copy(0.0)
fn __del__(owned self):
for i in range(self.rows * self.cols):
(self.data + i).destroy_pointee()
self.data.free()
fn get(self, row: Int, col: Int) -> Float32:
return self.data[row * self.cols + col]
fn set(inout self, row: Int, col: Int, value: Float32):
self.data[row * self.cols + col] = value
fn print_matrix(self):
for r in range(self.rows):
for c in range(self.cols):
print(self.get(r, c), end="\t")
print("")
Matrix Addition
Element-wise addition: add corresponding elements from two matrices of the same shape.
fn mat_add(a: Matrix, b: Matrix, result: inout Matrix):
for r in range(a.rows):
for c in range(a.cols):
result.set(r, c, a.get(r, c) + b.get(r, c))
fn main():
var a = Matrix(2, 3)
var b = Matrix(2, 3)
var result = Matrix(2, 3)
a.set(0, 0, 1.0); a.set(0, 1, 2.0); a.set(0, 2, 3.0)
a.set(1, 0, 4.0); a.set(1, 1, 5.0); a.set(1, 2, 6.0)
b.set(0, 0, 10.0); b.set(0, 1, 20.0); b.set(0, 2, 30.0)
b.set(1, 0, 40.0); b.set(1, 1, 50.0); b.set(1, 2, 60.0)
mat_add(a, b, result)
result.print_matrix()
Output:
11.0 22.0 33.0 44.0 55.0 66.0
Scalar Multiplication
fn mat_scale(m: Matrix, factor: Float32, result: inout Matrix):
for r in range(m.rows):
for c in range(m.cols):
result.set(r, c, m.get(r, c) * factor)
Matrix Transpose
Transposing swaps rows and columns. A 3×2 matrix becomes a 2×3 matrix. Element [r, c] moves to position [c, r].
fn transpose(m: Matrix, result: inout Matrix):
# result must be (m.cols × m.rows)
for r in range(m.rows):
for c in range(m.cols):
result.set(c, r, m.get(r, c))
Transpose diagram:
Original (2×3): Transposed (3×2):
[ 1 2 3 ] [ 1 4 ]
[ 4 5 6 ] [ 2 5 ]
[ 3 6 ]
[r, c] → [c, r]
Matrix Multiplication
Matrix multiplication combines two matrices into a third. To multiply A (M×K) by B (K×N), each output element C[r, c] is the dot product of row r of A with column c of B.
fn matmul(a: Matrix, b: Matrix, c: inout Matrix):
# a: M×K, b: K×N, c: M×N
for r in range(a.rows):
for col in range(b.cols):
var sum: Float32 = 0.0
for k in range(a.cols):
sum += a.get(r, k) * b.get(k, col)
c.set(r, col, sum)
Matrix multiply (2×3) × (3×2) → (2×2):
A: B: C = A×B:
[ 1 2 3 ] [ 7 8 ] [ 1×7+2×9+3×11 1×8+2×10+3×12 ]
[ 4 5 6 ] [ 9 10 ] = [ 4×7+5×9+6×11 4×8+5×10+6×12 ]
[ 11 12 ]
C[0,0] = 1×7 + 2×9 + 3×11 = 7 + 18 + 33 = 58
C[0,1] = 1×8 + 2×10 + 3×12 = 8 + 20 + 36 = 64
C[1,0] = 4×7 + 5×9 + 6×11 = 28 + 45 + 66 = 139
C[1,1] = 4×8 + 5×10 + 6×12 = 32 + 50 + 72 = 154
Result: [ 58 64 ]
[139 154 ]
Vectorized Matrix Multiply
The inner dot-product loop is the hot path in matrix multiplication. Vectorize it with SIMD for a direct performance gain.
fn matmul_simd(a: Matrix, b: Matrix, c: inout Matrix):
alias width = 8
for r in range(a.rows):
for col in range(b.cols):
var acc = SIMD[DType.float32, width].splat(0.0)
var k = 0
while k + width <= a.cols:
var av = SIMD[DType.float32, width].load(a.data + r * a.cols + k)
var bv = SIMD[DType.float32, width].load(b.data + k * b.cols + col)
acc += av * bv
k += width
var sum = acc.reduce_add()
while k < a.cols:
sum += a.get(r, k) * b.get(k, col)
k += 1
c.set(r, col, sum)
Identity Matrix
The identity matrix is the matrix equivalent of the number 1. Multiplying any matrix by the identity matrix returns the original matrix unchanged. The identity has 1s on the main diagonal and 0s everywhere else.
fn identity(n: Int, result: inout Matrix):
for r in range(n):
for c in range(n):
result.set(r, c, 1.0 if r == c else 0.0)
3×3 Identity matrix: [ 1 0 0 ] [ 0 1 0 ] [ 0 0 1 ] A × I = A (for any matrix A with matching dimensions)
Element-wise Operations
Unlike matrix multiplication, element-wise operations apply a function to each corresponding pair of elements independently. These appear everywhere in AI — applying activation functions, scaling layers, computing losses.
fn element_multiply(a: Matrix, b: Matrix, result: inout Matrix):
for r in range(a.rows):
for c in range(a.cols):
result.set(r, c, a.get(r, c) * b.get(r, c))
Element-wise multiply (Hadamard product): [ 2 3 ] ⊙ [ 5 6 ] = [ 10 18 ] [ 4 5 ] [ 7 8 ] [ 28 40 ] Notation: ⊙ (not the same as matrix multiply ×)
Performance Summary
Operation | Complexity | Key optimization ----------------------|------------|------------------------------- Matrix add | O(M×N) | SIMD row scan Matrix transpose | O(M×N) | Cache-friendly access pattern Matrix multiply | O(M×K×N) | SIMD inner loop + tiling Element-wise ops | O(M×N) | SIMD + vectorize Row/column reduction | O(M×N) | SIMD reduce_add per row
Key Takeaways
Matrices store two-dimensional data in a flat array using row-major order. Element [r, c] lives at index r × num_cols + c. Matrix addition and element-wise multiplication operate independently on each position. Transpose swaps row and column indices. Matrix multiplication computes a dot product between each row of the left matrix and each column of the right — it is the central operation in neural networks. Vectorize the inner loop with SIMD and tile the outer loops for cache efficiency to reach peak performance on large matrices.