rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 46 Packages Projects Releases Wiki Activity

perf(backend-cpu): parallelize blocked gemm over row blocks with rayon

Browse Source 
Each task owns a disjoint BLOCK_SIZE-row slice of the result; the inner
blocked kernel is unchanged. Needed for dg-gnn HetGAT training throughput
(node-level [M,64]x[64,64] matmuls dominated single-threaded step time).

Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
Omar Sobh
2026-08-20 11:37:43 -07:00
co-authored by Claude Fable 5
parent 35b2b2cdf4
commit 67c47898fa
1 changed files with 22 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/core/rtx-backend-cpu/src/ops/gemm.rs
+22 -15
View File
@@ -20,26 +20,33 @@ pub fn matmul<E: CpuFloat>(
// Block size for cache efficiency // Block size for cache efficiency
const BLOCK_SIZE: usize = 64; const BLOCK_SIZE: usize = 64;
// Blocked matrix multiplication // Blocked matrix multiplication, parallelized over row blocks: each rayon
for i_block in (0..m).step_by(BLOCK_SIZE) { // task owns a disjoint `BLOCK_SIZE`-row slice of the result, so the inner
for j_block in (0..n).step_by(BLOCK_SIZE) { // blocked kernel is unchanged and no synchronization is needed.
for k_block in (0..k).step_by(BLOCK_SIZE) { result
let i_end = (i_block + BLOCK_SIZE).min(m); .par_chunks_mut(BLOCK_SIZE * n)
let j_end = (j_block + BLOCK_SIZE).min(n); .enumerate()
let k_end = (k_block + BLOCK_SIZE).min(k); .for_each(|(bi, res_rows)| {
let i_block = bi * BLOCK_SIZE;
let i_end = (i_block + BLOCK_SIZE).min(m);
for j_block in (0..n).step_by(BLOCK_SIZE) {
for k_block in (0..k).step_by(BLOCK_SIZE) {
let j_end = (j_block + BLOCK_SIZE).min(n);
let k_end = (k_block + BLOCK_SIZE).min(k);
for i in i_block..i_end { for i in i_block..i_end {
for j in j_block..j_end { let row = &mut res_rows[(i - i_block) * n..(i - i_block) * n + n];
let mut sum = result[i * n + j]; for j in j_block..j_end {
for kk in k_block..k_end { let mut sum = row[j];
sum = sum + lhs.data[i * k + kk] * rhs.data[kk * n + j]; for kk in k_block..k_end {
sum = sum + lhs.data[i * k + kk] * rhs.data[kk * n + j];
}
row[j] = sum;
} }
result[i * n + j] = sum;
} }
} }
} }
} });
}
CpuTensorPrimitive::new(result, [m, n], lhs.device.clone()) CpuTensorPrimitive::new(result, [m, n], lhs.device.clone())
} }