Rust · Parallelism

Parallelism

8.0.0-rc.1 · b1f866ab

NumKong does not own a thread pool. The parallel feature adds host-side orchestration helpers via ForkUnion, not a hidden scheduler.

use numkong::{DotsPackedMatrix, Tensor};
use forkunion::{ThreadPool, Topology};

let a = Tensor::<f32>::full(&[4096, 768], 1.0).unwrap();
let b = Tensor::<f32>::full(&[8192, 768], 1.0).unwrap();
let topology = Topology::new().unwrap();
let mut pool = ThreadPool::spawn(&topology, 4).unwrap();

// GEMM-like: rows of A partitioned across threads, one shared packed B
let b_packed = DotsPackedMatrix::new(&b).unwrap();
let c = a.dots_packed_parallel(&b_packed, &mut pool).unwrap();
assert_eq!(c.shape(), &[4096, 8192]);

// SYRK-like: row windows of one square output partitioned across threads
let gram = a.dots_symmetric_parallel(&mut pool).unwrap();
assert_eq!(gram.shape(), &[4096, 4096]);

Packing splits the same way: DotsPackedMatrix::new_parallel by columns, AttentionPackedMatrix::new_parallel by segments and heads, and ShapesPackedMatrix::new_parallel by clouds. ShapesPackedMatrix::chamfers_parallel and its three siblings split the row clouds, each thread solving in its own workspace, and SqEuclideans3D::sqeuclideans3d_parallel splits the row points.

Rayon or a manual thread pool can still work if the rest of your application already depends on them.