Skip to content

Take the hot loops off the shape that blocks vectorization #510

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
godofecht merged 3 commits into from
Sep 28, 2026
Merged

Take the hot loops off the shape that blocks vectorization #510

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ab91797
Take the hot loops off the shape that blocks vectorization
godofecht Sep 28, 2026
8ba0c12
Take four more rows off the loops that were holding them back
godofecht Sep 28, 2026
6f4d538
Take the three generalized linear models off it too
godofecht Sep 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 43 additions & 24 deletions lib/scikit/cluster.flow
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,13 @@ import "lib/scikit/matrix.flow"
import "lib/scikit/prng.flow"
import "lib/scikit/blas.flow"

extern {
# clang lowers both to a single instruction, so a loop written with them
# holds no branch and the vectorizer can take it.
function fmaxf(a: f32, b: f32) -> f32
function fminf(a: f32, b: f32) -> f32
}

export struct KMeans {
centroids: ptr<ptr<f32> >,
n_clusters: i32,
Expand Down Expand Up @@ -2007,17 +2014,34 @@ export function affinity_propagation_fit(X: Matrix, damping: f32, max_iter: i32,
}
}

for k in 0 to n {
let mut max_other: f32 = best
if k == best_k { max_other = second }
let v: f32 = take * (si3[k] - max_other) + keep * ri2[k]
# Every element of the row takes the row maximum, and the two
# elements that do not are corrected afterwards. Deciding it per
# element put a branch on k against best_k and another on k
# against i inside a loop that is otherwise a fused multiply add
# and a max, which is the difference between a vector loop and a
# scalar one.
let old_best: f32 = ri2[best_k]
# Written as a while loop with a constant step. A `for k in 0 to n`
# emits an increment chosen by a runtime test of the bounds, which
# is not a stride the LLVM vectorizer will accept, and this loop
# was scalar with it.
let mut k: i32 = 0
while k < n {
let v: f32 = take * (si3[k] - best) + keep * ri2[k]
ri2[k] = v
if k == i {
rkk[i] = v
} elif v > 0.0 {
col_pos[k] = col_pos[k] + v
}
col_pos[k] = col_pos[k] + fmaxf(v, 0.0)
k = k + 1
}

# The largest entry of the row competes against the runner up.
let fixed: f32 = take * (si3[best_k] - second) + keep * old_best
col_pos[best_k] = col_pos[best_k] + fmaxf(fixed, 0.0) - fmaxf(ri2[best_k], 0.0)
ri2[best_k] = fixed

# The diagonal is the self responsibility, and it is not part of
# the column sum.
rkk[i] = ri2[i]
col_pos[i] = col_pos[i] - fmaxf(ri2[i], 0.0)
}

# The availability update. The term that depends only on the column is
Expand All @@ -2029,21 +2053,16 @@ export function affinity_propagation_fit(X: Matrix, damping: f32, max_iter: i32,
for i in 0 to n {
let ai2: ptr<f32> = A + i * n
let ri4: ptr<f32> = R + i * n
for k in 0 to i {
let mut val: f32 = base[k]
let r_ik: f32 = ri4[k]
if r_ik > 0.0 { val = val - r_ik }
if val > 0.0 { val = 0.0 }
ai2[k] = take * val + keep * ai2[k]
}
ai2[i] = take * base[i] + keep * ai2[i]
for k in i + 1 to n {
let mut val2: f32 = base[k]
let r_ik2: f32 = ri4[k]
if r_ik2 > 0.0 { val2 = val2 - r_ik2 }
if val2 > 0.0 { val2 = 0.0 }
ai2[k] = take * val2 + keep * ai2[k]
}
let old_diag: f32 = ai2[i]
let mut k2: i32 = 0
while k2 < n {
let val: f32 = fminf(base[k2] - fmaxf(ri4[k2], 0.0), 0.0)
ai2[k2] = take * val + keep * ai2[k2]
k2 = k2 + 1
}
# The diagonal keeps the whole column sum rather than the clipped
# one, and is written after the loop for the reason above.
ai2[i] = take * base[i] + keep * old_diag
}

# Convergence is the set of exemplars holding still, which is what
Expand Down
11 changes: 9 additions & 2 deletions lib/scikit/linear.flow
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7252,8 +7252,12 @@ function _glm_gradient_fit(X: Matrix, y: ptr<f32>, weights: ptr<f32>, bias_init:

for iter in 0 to max_iter {
blas_matvec(xdata, n, p, weights, eta, 1.0, 0.0)
# While loops with a constant step. Flow emits a range loop with a
# stride chosen at runtime, which LLVM will not vectorize. Flow issue
# #1004.
let mut grad_b: f32 = 0.0
for i in 0 to n {
let mut i: i32 = 0
while i < n {
let e: f32 = eta[i] + bias
let mut mu: f32 = 0.0
if kind == 0 { mu = _poisson_inverse_link(e) }
Expand All @@ -7262,11 +7266,14 @@ function _glm_gradient_fit(X: Matrix, y: ptr<f32>, weights: ptr<f32>, bias_init:
let r: f32 = y[i] - mu
resid[i] = r
grad_b = grad_b + r
i = i + 1
}
blas_matvec_trans(xdata, n, p, resid, grad_w, 1.0, 0.0)
for j in 0 to p {
let mut j: i32 = 0
while j < p {
let g: f32 = grad_w[j] * inv_n - alpha * weights[j]
weights[j] = weights[j] + lr * g
j = j + 1
}
bias = bias + lr * grad_b * inv_n
}
Expand Down
199 changes: 93 additions & 106 deletions lib/scikit/manifold.flow
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@ import "lib/scikit/blas.flow"
extern {
function rand() -> i32
function srand(seed: u32) -> void
# clang lowers this to a single instruction, so a loop written with it
# holds no branch and the vectorizer can take it.
function fminf(a: f32, b: f32) -> f32
}

export struct TSNE {
Expand Down Expand Up @@ -306,125 +309,132 @@ export struct Isomap {
fitted: bool
}

# Modified Gram-Schmidt on the columns of a row-major n by k block.
# A column that collapses is left at zero, and the remaining columns stay
# orthonormal.
function _orthonormalize_block(V: ptr<f32>, n: i32, k: i32) -> void {
for c in 0 to k {
for prev in 0 to c {
let mut dot: f32 = 0.0
for i in 0 to n { dot = dot + V[i * k + c] * V[i * k + prev] }
for i in 0 to n { V[i * k + c] = V[i * k + c] - dot * V[i * k + prev] }
}
let mut norm: f32 = 0.0
for i in 0 to n {
let v: f32 = V[i * k + c]
norm = norm + v * v
}
norm = sqrt((norm) as f64) as f32
if norm > 0.0000000001 {
let inv: f32 = 1.0 / norm
for i in 0 to n { V[i * k + c] = V[i * k + c] * inv }
}
}
}

export function isomap_fit(X: Matrix, n_components: i32, n_neighbors: i32) -> Isomap {
let n: i32 = X.rows
let p: i32 = X.cols
let kc: i32 = n_components

let D: ptr<ptr<f32> > = malloc((n as i64) * 8) as ptr<ptr<f32> >
# Distances, then geodesics, then the double centering, all in contiguous
# blocks. Every inner loop below is a while loop with a constant step,
# because Flow emits a range loop with a stride chosen at runtime and LLVM
# will not vectorize that. Flow issue #1004.
let xdata: ptr<f32> = X.data
let geo: ptr<f32> = array_new_f32(n * n)
for i in 0 to n {
D[i] = array_new_f32(n)
for j in 0 to n {
let ri: ptr<f32> = xdata + i * p
let gi: ptr<f32> = geo + i * n
for j in i + 1 to n {
let rj: ptr<f32> = xdata + j * p
let mut s: f32 = 0.0
for k in 0 to p {
let d: f32 = matrix_at(X, i, k) - matrix_at(X, j, k)
let mut m: i32 = 0
while m < p {
let d: f32 = ri[m] - rj[m]
s = s + d * d
m = m + 1
}
D[i][j] = sqrt((s as f64)) as f32
let dist: f32 = sqrt((s) as f64) as f32
gi[j] = dist
geo[j * n + i] = dist
}
}

let geo: ptr<ptr<f32> > = malloc((n as i64) * 8) as ptr<ptr<f32> >
for i in 0 to n {
geo[i] = array_copy_f32(D[i], n)
}

# Floyd-Warshall. The term that depends only on i and k is lifted out of
# the inner loop, and the comparison is a minimum, so the loop holds no
# branch and vectorizes.
for k in 0 to n {
let rowk: ptr<f32> = geo + k * n
for i in 0 to n {
for j in 0 to n {
if geo[i][k] + geo[k][j] < geo[i][j] {
geo[i][j] = geo[i][k] + geo[k][j]
}
let rowi: ptr<f32> = geo + i * n
let dik: f32 = rowi[k]
let mut j: i32 = 0
while j < n {
rowi[j] = fminf(rowi[j], dik + rowk[j])
j = j + 1
}
}
}

let mut mean_d: f32 = 0.0
for i in 0 to n {
for j in 0 to n {
geo[i][j] = geo[i][j] * geo[i][j]
mean_d = mean_d + geo[i][j]
let gi2: ptr<f32> = geo + i * n
let mut j2: i32 = 0
while j2 < n {
let v: f32 = gi2[j2] * gi2[j2]
gi2[j2] = v
mean_d = mean_d + v
j2 = j2 + 1
}
}
mean_d = mean_d / ((n * n) as f32)

let B: ptr<ptr<f32> > = malloc((n as i64) * 8) as ptr<ptr<f32> >
let B: ptr<f32> = array_new_f32(n * n)
for i in 0 to n {
B[i] = array_new_f32(n)
for j in 0 to n {
B[i][j] = -0.5 * (geo[i][j] - mean_d)
let gi3: ptr<f32> = geo + i * n
let bi: ptr<f32> = B + i * n
let mut j3: i32 = 0
while j3 < n {
bi[j3] = -0.5 * (gi3[j3] - mean_d)
j3 = j3 + 1
}
}
array_free_f32(geo)

let embedding: ptr<f32> = array_new_f32(n * n_components)
# Track found eigenvectors for deflation
let found_vecs: ptr<ptr<f32> > = malloc((n_components as i64) * 8) as ptr<ptr<f32> >
for c in 0 to n_components {
let v: ptr<f32> = array_new_f32(n)
srand(42 + c)
for i in 0 to n { v[i] = ((rand() % 1000) as f32) / 500.0 - 1.0 }

# Orthogonalize against previously found eigenvectors
for prev in 0 to c {
let mut dot: f32 = 0.0
for i in 0 to n { dot = dot + v[i] * found_vecs[prev][i] }
for i in 0 to n { v[i] = v[i] - dot * found_vecs[prev][i] }
}

let mut iter: i32 = 0
while iter < 200 {
let Bv: ptr<f32> = array_new_f32(n)
for i in 0 to n {
let mut s: f32 = 0.0
for j in 0 to n { s = s + B[i][j] * v[j] }
Bv[i] = s
}

# Deflate against previous eigenvectors
for prev in 0 to c {
let mut dot: f32 = 0.0
for i in 0 to n { dot = dot + Bv[i] * found_vecs[prev][i] }
for i in 0 to n { Bv[i] = Bv[i] - dot * found_vecs[prev][i] }
}

let mut norm: f32 = 0.0
for i in 0 to n { norm = norm + Bv[i] * Bv[i] }
norm = sqrt((norm as f64)) as f32
if norm < 0.0000000001 { norm = 0.0000000001 }

# Convergence check
let mut max_change: f32 = 0.0
for i in 0 to n {
let new_v: f32 = Bv[i] / norm
let diff: f32 = new_v - v[i]
if diff < 0.0 { diff = -diff }
if diff > max_change { max_change = diff }
v[i] = new_v
}
array_free_f32(Bv)
# Block power iteration. One sgemm per iteration advances every component,
# where the old loop ran a scalar matvec per component per iteration and
# allocated its result buffer inside the iteration.
let V: ptr<f32> = array_new_f32(n * kc)
let Z: ptr<f32> = array_new_f32(n * kc)
srand(42)
for i in 0 to n * kc { V[i] = ((rand() % 1000) as f32) / 500.0 - 1.0 }
_orthonormalize_block(V, n, kc)

if max_change < 0.000001 { break }
iter = iter + 1
}
let mut iter: i32 = 0
while iter < 200 {
blas_matmat(B, V, Z, n, kc, n, 1.0, 0.0)
_orthonormalize_block(Z, n, kc)

for i in 0 to n {
embedding[i * n_components + c] = v[i]
let mut max_change: f32 = 0.0
for i in 0 to n * kc {
let mut diff: f32 = Z[i] - V[i]
if diff < 0.0 { diff = 0.0 - diff }
if diff > max_change { max_change = diff }
V[i] = Z[i]
}
found_vecs[c] = v
if max_change < 0.000001 { break }
iter = iter + 1
}

for c in 0 to n_components { array_free_f32(found_vecs[c]) }
free(found_vecs as ptr<void>)

for i in 0 to n { array_free_f32(D[i]); array_free_f32(geo[i]); array_free_f32(B[i]) }
free(D as ptr<void>)
free(geo as ptr<void>)
free(B as ptr<void>)
array_free_f32(B)
array_free_f32(Z)

return Isomap {
embedding: embedding,
embedding: V,
n_samples: n,
n_features: p,
n_components: n_components,
n_components: kc,
n_neighbors: n_neighbors,
fitted: true
}
Expand Down Expand Up @@ -456,29 +466,6 @@ export struct LLE {
fitted: bool
}

# Modified Gram-Schmidt on the columns of a row-major n by k block.
# A column that collapses is left at zero, and the remaining columns stay
# orthonormal.
function _orthonormalize_block(V: ptr<f32>, n: i32, k: i32) -> void {
for c in 0 to k {
for prev in 0 to c {
let mut dot: f32 = 0.0
for i in 0 to n { dot = dot + V[i * k + c] * V[i * k + prev] }
for i in 0 to n { V[i * k + c] = V[i * k + c] - dot * V[i * k + prev] }
}
let mut norm: f32 = 0.0
for i in 0 to n {
let v: f32 = V[i * k + c]
norm = norm + v * v
}
norm = sqrt((norm) as f64) as f32
if norm > 0.0000000001 {
let inv: f32 = 1.0 / norm
for i in 0 to n { V[i * k + c] = V[i * k + c] * inv }
}
}
}

export function lle_fit(X: Matrix, n_components: i32, n_neighbors: i32) -> LLE {
let n: i32 = X.rows
let p: i32 = X.cols
Expand Down
Loading
Loading