daniel543x/effectcraft
1
0 Fork 0
mirror of https://github.com/storytold/effectcraft.git synced 2026-10-31 15:48:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

S2.62: GPU encoder (wgpu): Hiera + neck on any GPU, matches PyTorch to 7e-6

Browse Source 
Strided batched matmul with bias, GELU/ReLU and residual in the epilogue, layer
norm, softmax, window partition, pooling, patch embedding, RoPE and neck kernels
(WGSL, all tested against CPU definitions). GpuEncoder keeps activations on the
GPU and returns the 64x64 level and the decoder's skip features. On the RX 7900
GRE one frame encodes in 0.29 s (debug build) against 2.1 s on the CPU in release;
errors 6.8e-6 vs PyTorch, 1.2e-5 vs the CPU path.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Kristian MyramandClaude Sonnet 5.5 committed 2026-10-30 08:43:31 +01:00
1 parent 921bb00ba6
commit 5c7eb3c2ae
12 files changed
+1490 -16
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+1
View File
@@ -1836,6 +1836,7 @@ dependencies = [
"candle-core",
"candle-nn",
"effectcraft-segment",
"log",
"pollster 0.4.0",
"rayon",
"wgpu",
crates/sam2/Cargo.toml
+2 -3
View File
@@ -12,14 +12,13 @@ workspace = true
[dependencies]
effectcraft-segment = { workspace = true }
rayon = { workspace = true }
log = { workspace = true }
# Native only (D1): candle is not built for wasm32, where this crate is empty.
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
candle-core = "=0.9.2"
candle-nn = "=0.9.2"
# The GPU matmul spike (examples/gpu_matmul.rs) only.
[dev-dependencies]
# The optional GPU path (wgpu compute, any vendor); the CPU path stays the reference.
wgpu = { workspace = true }
pollster = { workspace = true }
bytemuck = { workspace = true }
crates/sam2/src/gpu/encoder.rs
+286
View File
@@ -0,0 +1,286 @@
// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2
// implementation (Apache-2.0, see NOTICE).
//! The vision encoder on the GPU: patch embedding, the Hiera blocks and the neck, with the decoder's two
//! skip convolutions. Activations stay on the GPU between the kernels; only the three tensors the rest of
//! the model needs come back.
use std::collections::HashMap;
use std::sync::{Arc, Mutex, PoisonError};
use candle_core::{Device, Tensor};
use super::{Act, Buf, GLinear, GNorm, Gpu, Matmul, Strides};
use crate::hiera::{BlockPlan, Hiera, HieraConfig};
use crate::weights::Weights;
use crate::{Error, Result};
struct GBlock {
plan: BlockPlan,
norm1: GNorm,
qkv: GLinear,
proj: GLinear,
norm2: GNorm,
mlp_in: GLinear,
mlp_out: GLinear,
block_proj: Option<GLinear>,
}
/// What the rest of the model takes from the encoder (CPU tensors, NCHW).
pub struct GpuFeatures {
/// The lowest-resolution neck level `[1, 256, 64, 64]`.
pub low: Tensor,
/// The decoder's skip features: `conv_s0` of the highest level `[1, 32, 256, 256]` and `conv_s1` of the
/// next `[1, 64, 128, 128]`.
pub high: (Tensor, Tensor),
}
pub struct GpuEncoder {
gpu: Arc<Gpu>,
cfg: HieraConfig,
patch_w: Buf,
patch_b: Buf,
blocks: Vec<GBlock>,
stage_ends: Vec<usize>,
neck: Vec<GLinear>,
top_down: Vec<usize>,
skip0: GLinear,
skip1: GLinear,
/// The position embedding of each grid size, uploaded on first use.
pos: Mutex<HashMap<(usize, usize), Buf>>,
/// The CPU backbone, which owns the position embedding tables.
cpu: Arc<Hiera>,
eps: f32,
}
impl GpuEncoder {
/// Upload the backbone, the neck and the skip convolutions. `cpu` supplies the position embedding.
pub fn load(gpu: Arc<Gpu>, w: &Weights, cpu: Arc<Hiera>) -> Result<GpuEncoder> {
let cfg = cpu.config().clone();
let flat = |t: Tensor| -> Result<Vec<f32>> { Ok(t.flatten_all()?.to_vec1::<f32>()?) };
let hidden = cfg.embed_dims.first().copied().unwrap_or(96);
let mut blocks = vec![];
for (total, plan) in BlockPlan::all(&cfg)?.into_iter().enumerate() {
let p = format!("vision_encoder.backbone.blocks.{total}");
blocks.push(GBlock {
plan,
norm1: GNorm::load(&gpu, w, &format!("{p}.layer_norm1"), plan.dim)?,
qkv: GLinear::load(&gpu, w, &format!("{p}.attn.qkv"), plan.dim, plan.dim_out * 3)?,
proj: GLinear::load(&gpu, w, &format!("{p}.attn.proj"), plan.dim_out, plan.dim_out)?,
norm2: GNorm::load(&gpu, w, &format!("{p}.layer_norm2"), plan.dim_out)?,
mlp_in: GLinear::load(&gpu, w, &format!("{p}.mlp.proj_in"), plan.dim_out, plan.dim_out * cfg.mlp_ratio)?,
mlp_out: GLinear::load(&gpu, w, &format!("{p}.mlp.proj_out"), plan.dim_out * cfg.mlp_ratio, plan.dim_out)?,
block_proj: (plan.dim != plan.dim_out).then(|| GLinear::load(&gpu, w, &format!("{p}.proj"), plan.dim, plan.dim_out)).transpose()?,
});
}
let mut stage_ends = vec![];
let mut n = 0;
for b in &cfg.blocks {
n += b;
stage_ends.push(n.saturating_sub(1));
}
let channels: Vec<usize> = cfg.embed_dims.iter().rev().copied().collect();
let neck = channels
.iter()
.enumerate()
.map(|(j, c)| GLinear::load_conv1x1(&gpu, w, &format!("vision_encoder.neck.convs.{j}"), *c, 256))
.collect::<Result<Vec<_>>>()?;
Ok(GpuEncoder {
patch_w: gpu.upload(&flat(w.get("vision_encoder.backbone.patch_embed.projection.weight", &[hidden, 3, 7, 7])?)?),
patch_b: gpu.upload(&flat(w.get("vision_encoder.backbone.patch_embed.projection.bias", &[hidden])?)?),
skip0: GLinear::load_conv1x1(&gpu, w, "mask_decoder.conv_s0", 256, 32)?,
skip1: GLinear::load_conv1x1(&gpu, w, "mask_decoder.conv_s1", 256, 64)?,
gpu,
eps: cfg.eps as f32,
cfg,
blocks,
stage_ends,
neck,
top_down: vec![2, 3],
pos: Mutex::default(),
cpu,
})
}
fn linear(&self, x: &Buf, rows: usize, l: &GLinear, out: &Buf, act: Act, residual: Option<&Buf>) -> Result<()> {
self.gpu.linear(x, rows, l, out, act, residual)
}
fn norm(&self, x: &Buf, rows: usize, dim: usize, n: &GNorm) -> Result<Buf> {
let out = self.gpu.zeros(rows * dim)?;
self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps)?;
Ok(out)
}
fn pos_buffer(&self, h: usize, w: usize) -> Result<Buf> {
let mut map = self.pos.lock().unwrap_or_else(PoisonError::into_inner);
if let Some(b) = map.get(&(h, w)) {
return Ok(b.clone());
}
let table = self.cpu.pos(h, w)?.flatten_all()?.to_vec1::<f32>()?;
let buf = self.gpu.upload(&table);
map.insert((h, w), buf.clone());
Ok(buf)
}
/// One Hiera block on `[h * w, dim]` tokens; returns the new tokens and grid size.
fn block(&self, b: &GBlock, x: &Buf, h: usize, w: usize) -> Result<(Buf, usize, usize)> {
let g = &self.gpu;
let BlockPlan { dim, dim_out, heads, window, query_stride } = b.plan;
let n = h * w;
let normed = self.norm(x, n, dim, &b.norm1)?;
// the skip connection: from the block's input, or (stage change) its projection, pooled with the queries
let (oh, ow) = query_stride.map_or((h, w), |s| (h / s, w / s));
if query_stride.is_some_and(|s| s != 2) {
return Err(Error::Input("only a query stride of 2 is supported on the GPU".into()));
}
let residual = match (&b.block_proj, query_stride) {
(Some(p), Some(_)) => {
let proj = g.zeros(n * dim_out)?;
self.linear(&normed, n, p, &proj, Act::None, None)?;
let pooled = g.zeros(oh * ow * dim_out)?;
g.maxpool2(&proj, &pooled, 1, h, w, dim_out, (n * dim_out, w * dim_out, dim_out))?;
pooled
}
(Some(p), None) => {
let proj = g.zeros(n * dim_out)?;
self.linear(&normed, n, p, &proj, Act::None, None)?;
proj
}
(None, _) => x.clone(),
};
// windows (or the whole grid for global attention)
let (hw, nwin, t, win_h, win_w) = if window > 0 {
let (nwx, nwy) = (w.div_ceil(window), h.div_ceil(window));
let buf = g.zeros(nwx * nwy * window * window * dim)?;
g.window_partition(&normed, &buf, h, w, dim, window)?;
(buf, nwx * nwy, window * window, window, window)
} else {
(normed, 1, n, h, w)
};
let tok = nwin * t;
let qkv = g.zeros(tok * 3 * dim_out)?;
self.linear(&hw, tok, &b.qkv, &qkv, Act::None, None)?;
// queries, pooled when the stage changes
let (q_buf, q_rs, tq) = match query_stride {
Some(s) => {
let (qh, qw) = (win_h / s, win_w / s);
let pooled = g.zeros(nwin * qh * qw * dim_out)?;
g.maxpool2(&qkv, &pooled, nwin, win_h, win_w, dim_out, (t * 3 * dim_out, win_w * 3 * dim_out, 3 * dim_out))?;
(pooled, dim_out, qh * qw)
}
None => (qkv.clone(), 3 * dim_out, t),
};
let hd = dim_out / heads.max(1);
let scale = (hd as f32).powf(-0.5);
let scores = g.zeros(nwin * tq * t)?;
let attn = g.zeros(nwin * tq * dim_out)?;
for hh in 0..heads {
g.matmul(&Matmul {
m: tq,
n: t,
k: hd,
batch: nwin,
a: (&q_buf, Strides { offset: hh * hd, batch: tq * q_rs, row: q_rs, col: 1 }),
b: (&qkv, Strides { offset: dim_out + hh * hd, batch: t * 3 * dim_out, row: 1, col: 3 * dim_out }),
bias: None,
c: (&scores, Strides { offset: 0, batch: tq * t, row: t, col: 1 }),
residual: None,
alpha: scale,
act: Act::None,
})?;
g.softmax(&scores, nwin * tq, t)?;
g.matmul(&Matmul {
m: tq,
n: hd,
k: t,
batch: nwin,
a: (&scores, Strides { offset: 0, batch: tq * t, row: t, col: 1 }),
b: (&qkv, Strides { offset: 2 * dim_out + hh * hd, batch: t * 3 * dim_out, row: 3 * dim_out, col: 1 }),
bias: None,
c: (&attn, Strides { offset: hh * hd, batch: tq * dim_out, row: dim_out, col: 1 }),
residual: None,
alpha: 1.0,
act: Act::None,
})?;
}
let projected = g.zeros(nwin * tq * dim_out)?;
self.linear(&attn, nwin * tq, &b.proj, &projected, Act::None, None)?;
let attn_img = if window > 0 {
let ws = window / query_stride.unwrap_or(1);
let img = g.zeros(oh * ow * dim_out)?;
g.window_unpartition(&projected, &img, oh, ow, dim_out, ws)?;
img
} else {
projected
};
let n2 = oh * ow;
let x2 = g.zeros(n2 * dim_out)?;
g.add(&residual, &attn_img, &x2, n2 * dim_out)?;
let normed2 = self.norm(&x2, n2, dim_out, &b.norm2)?;
let hidden = g.zeros(n2 * b.mlp_in.out)?;
self.linear(&normed2, n2, &b.mlp_in, &hidden, Act::Gelu, None)?;
let out = g.zeros(n2 * dim_out)?;
self.linear(&hidden, n2, &b.mlp_out, &out, Act::None, Some(&x2))?;
Ok((out, oh, ow))
}
/// Encode normalised `[1, 3, H, W]` pixels (`H` and `W` multiples of 32).
pub fn encode(&self, pixels: &Tensor) -> Result<GpuFeatures> {
let g = &self.gpu;
let (_, c, ph, pw) = pixels.dims4()?;
let window0 = self.cfg.windows.first().copied().unwrap_or(8);
if c != 3 || ph == 0 || pw == 0 || ph % 4 != 0 || pw % 4 != 0 || (ph / 4) % window0 != 0 || (pw / 4) % window0 != 0 {
return Err(Error::Input(format!("image size {pw}x{ph} does not fit the encoder")));
}
let hidden = self.cfg.embed_dims.first().copied().unwrap_or(96);
let px = g.upload(&pixels.contiguous()?.flatten_all()?.to_vec1::<f32>()?);
let (mut h, mut w) = (ph / 4, pw / 4);
let tokens = g.zeros(h * w * hidden)?;
g.patch_embed(&px, &self.patch_w, &self.patch_b, &tokens, ph, pw, hidden)?;
let mut x = g.zeros(h * w * hidden)?;
g.add(&tokens, &self.pos_buffer(h, w)?, &x, h * w * hidden)?;
let mut stages: Vec<(Buf, usize, usize, usize)> = vec![];
for (i, b) in self.blocks.iter().enumerate() {
let (nx, nh, nw) = self.block(b, &x, h, w)?;
(x, h, w) = (nx, nh, nw);
if self.stage_ends.contains(&i) {
stages.push((x.clone(), h, w, b.plan.dim_out));
}
}
// the neck, from the lowest resolution up: lateral 1x1 convolutions, and the top-down sum on the levels that have it
let n = self.neck.len();
if stages.len() != n {
return Err(Error::Compute("the backbone returned the wrong number of stages".into()));
}
let mut levels: Vec<(Buf, usize, usize)> = vec![];
let mut prev: Option<(Buf, usize, usize)> = None;
for i in (0..n).rev() {
let (buf, sh, sw, _) = &stages[i];
let conv = self.neck.get(n - 1 - i).ok_or_else(|| Error::Compute("missing neck convolution".into()))?;
let lateral = g.zeros(sh * sw * 256)?;
self.linear(buf, sh * sw, conv, &lateral, Act::None, None)?;
let cur = match (&prev, self.top_down.contains(&i) && i != n - 1) {
(Some((p, _, _)), true) => {
let sum = g.zeros(sh * sw * 256)?;
g.upsample_add(&lateral, p, &sum, *sh, *sw, 256)?;
sum
}
_ => lateral,
};
prev = Some((cur.clone(), *sh, *sw));
levels.push((cur, *sh, *sw));
}
// levels: lowest resolution first. The decoder needs the third (64 x 64) and the skip convolutions of the last two.
let [_, (l2, h2, w2), (l1, h1, w1), (l0, h0, w0)] = <[_; 4]>::try_from(levels).map_err(|_| Error::Compute("the neck must have four levels".into()))?;
let s0 = g.zeros(h0 * w0 * 32)?;
self.linear(&l0, h0 * w0, &self.skip0, &s0, Act::None, None)?;
let s1 = g.zeros(h1 * w1 * 64)?;
self.linear(&l1, h1 * w1, &self.skip1, &s1, Act::None, None)?;
g.flush()?;
let nchw = |buf: &Buf, hh: usize, ww: usize, ch: usize| -> Result<Tensor> {
let v = g.download(buf)?;
Ok(Tensor::from_vec(v, (1, hh, ww, ch), &Device::Cpu)?.permute((0, 3, 1, 2))?.contiguous()?)
};
Ok(GpuFeatures { low: nchw(&l2, h2, w2, 256)?, high: (nchw(&s0, h0, w0, 32)?, nchw(&s1, h1, w1, 64)?) })
}
}
crates/sam2/src/gpu/kernels.wgsl
+98
View File
@@ -0,0 +1,98 @@
// SAM 2 compute kernels (f32). Original work for EffectCraft; the maths follows the Hugging Face
// `transformers` SAM 2 implementation (Apache-2.0, see NOTICE).
//
// Every kernel reads its parameters from one uniform: 24 u32 words (floats are bit-cast).
struct P { v: array<vec4<u32>, 6> }
@group(0) @binding(0) var<storage, read> in0: array<f32>;
@group(0) @binding(1) var<storage, read> in1: array<f32>;
@group(0) @binding(2) var<storage, read> in2: array<f32>;
@group(0) @binding(3) var<storage, read_write> out0: array<f32>;
@group(0) @binding(4) var<storage, read> in3: array<f32>;
@group(0) @binding(5) var<uniform> p: P;
fn u(i: u32) -> u32 { return p.v[i / 4u][i % 4u]; }
fn f(i: u32) -> f32 { return bitcast<f32>(u(i)); }
// erf, Abramowitz and Stegun 7.1.26 (abs error 1.5e-7).
fn erf_(x: f32) -> f32 {
let s = select(1.0, -1.0, x < 0.0);
let a = abs(x);
let t = 1.0 / (1.0 + 0.3275911 * a);
let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-a * a);
return s * y;
}
fn gelu_(x: f32) -> f32 { return (erf_(x * 0.70710678) + 1.0) * 0.5 * x; }
// ---------------------------------------------------------------- strided batched matmul
// C[b] = act(alpha * A[b] (m x k) . B[b] (k x n) + bias) (+ residual), with free strides for A and B.
// words: 0 m, 1 n, 2 k, 3 batch, 4 a_off, 5 a_bs, 6 a_rs, 7 a_cs, 8 b_off, 9 b_bs, 10 b_rs, 11 b_cs,
// 12 c_off, 13 c_bs, 14 c_rs, 15 flags (bit0 bias, bits 1-2 act: 1 gelu 2 relu, bit3 residual),
// 16 alpha (f32)
// in0 = A, in1 = B, in2 = bias (n), out0 = C, in3 = residual (same layout as C)
var<workgroup> sa: array<f32, 1024>;
var<workgroup> sb: array<f32, 1024>;
@compute @workgroup_size(16, 16)
fn matmul(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let m = u(0u); let n = u(1u); let k = u(2u);
let a_base = u(4u) + wg.z * u(5u);
let b_base = u(8u) + wg.z * u(9u);
let a_rs = u(6u); let a_cs = u(7u); let b_rs = u(10u); let b_cs = u(11u);
let row0 = wg.y * 64u;
let col0 = wg.x * 64u;
let tid = li.y * 16u + li.x;
let a_k_contig = a_cs == 1u;
let b_n_contig = b_cs == 1u;
var acc: array<f32, 16>;
for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; }
for (var kt = 0u; kt < k; kt += 16u) {
for (var i = 0u; i < 4u; i++) {
let idx = tid + i * 256u;
var r: u32; var c: u32;
if (a_k_contig) { r = idx / 16u; c = idx % 16u; } else { r = idx % 64u; c = idx / 64u; }
let gr = row0 + r; let gc = kt + c;
var va = 0.0;
if (gr < m && gc < k) { va = in0[a_base + gr * a_rs + gc * a_cs]; }
sa[r * 16u + c] = va;
var r2: u32; var c2: u32;
if (b_n_contig) { r2 = idx / 64u; c2 = idx % 64u; } else { r2 = idx % 16u; c2 = idx / 16u; }
let gr2 = kt + r2; let gc2 = col0 + c2;
var vb = 0.0;
if (gr2 < k && gc2 < n) { vb = in1[b_base + gr2 * b_rs + gc2 * b_cs]; }
sb[r2 * 64u + c2] = vb;
}
workgroupBarrier();
for (var kk = 0u; kk < 16u; kk++) {
var ar: array<f32, 4>;
var br: array<f32, 4>;
for (var i = 0u; i < 4u; i++) {
ar[i] = sa[(li.y * 4u + i) * 16u + kk];
br[i] = sb[kk * 64u + li.x * 4u + i];
}
for (var i = 0u; i < 4u; i++) {
for (var j = 0u; j < 4u; j++) { acc[i * 4u + j] += ar[i] * br[j]; }
}
}
workgroupBarrier();
}
let flags = u(15u);
let alpha = f(16u);
let c_base = u(12u) + wg.z * u(13u);
let c_rs = u(14u);
let act = (flags >> 1u) & 3u;
for (var i = 0u; i < 4u; i++) {
for (var j = 0u; j < 4u; j++) {
let gr = row0 + li.y * 4u + i;
let gc = col0 + li.x * 4u + j;
if (gr < m && gc < n) {
var v = acc[i * 4u + j] * alpha;
if ((flags & 1u) != 0u) { v += in2[gc]; }
if (act == 1u) { v = gelu_(v); } else if (act == 2u) { v = max(v, 0.0); }
let idx = c_base + gr * c_rs + gc;
if ((flags & 8u) != 0u) { v += in3[idx]; }
out0[idx] = v;
}
}
}
}
crates/sam2/src/gpu/kernels_ops.wgsl
+173
View File
@@ -0,0 +1,173 @@
// ---------------------------------------------------------------- helpers
// A 1-D index over n items spread on a 2-D grid of 256-thread groups (a dimension holds 65535 groups).
fn gid1(wg: vec3<u32>, nw: vec3<u32>, li: vec3<u32>) -> u32 { return (wg.y * nw.x + wg.x) * 256u + li.x; }
var<workgroup> red: array<f32, 64>;
// Sum of `red` over the workgroup (64 threads), for a caller that has filled it; all threads get the result.
fn reduce_sum(li: u32) -> f32 {
workgroupBarrier();
for (var st = 32u; st > 0u; st = st / 2u) {
if (li < st) { red[li] += red[li + st]; }
workgroupBarrier();
}
let r = red[0];
workgroupBarrier();
return r;
}
fn reduce_max(li: u32) -> f32 {
workgroupBarrier();
for (var st = 32u; st > 0u; st = st / 2u) {
if (li < st) { red[li] = max(red[li], red[li + st]); }
workgroupBarrier();
}
let r = red[0];
workgroupBarrier();
return r;
}
// ---------------------------------------------------------------- layer norm over the last dimension
// words: 0 rows, 1 cols, 2 eps (f32). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row.
@compute @workgroup_size(64)
fn layernorm(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let row = wg.y * nw.x + wg.x;
let rows = u(0u); let cols = u(1u);
if (row >= rows) { return; }
let base = row * cols;
var s = 0.0;
for (var c = li.x; c < cols; c += 64u) { s += in0[base + c]; }
red[li.x] = s;
let mean = reduce_sum(li.x) / f32(cols);
var q = 0.0;
for (var c = li.x; c < cols; c += 64u) { let d = in0[base + c] - mean; q += d * d; }
red[li.x] = q;
let inv = inverseSqrt(reduce_sum(li.x) / f32(cols) + f(2u));
for (var c = li.x; c < cols; c += 64u) { out0[base + c] = (in0[base + c] - mean) * inv * in1[c] + in2[c]; }
}
// ---------------------------------------------------------------- softmax over rows, in place on out0
// words: 0 rows, 1 cols. One workgroup per row.
@compute @workgroup_size(64)
fn softmax(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let row = wg.y * nw.x + wg.x;
let rows = u(0u); let cols = u(1u);
if (row >= rows) { return; }
let base = row * cols;
var mx = -3.0e38;
for (var c = li.x; c < cols; c += 64u) { mx = max(mx, out0[base + c]); }
red[li.x] = mx;
let m = reduce_max(li.x);
var s = 0.0;
for (var c = li.x; c < cols; c += 64u) { let e = exp(out0[base + c] - m); out0[base + c] = e; s += e; }
red[li.x] = s;
let inv = 1.0 / reduce_sum(li.x);
for (var c = li.x; c < cols; c += 64u) { out0[base + c] = out0[base + c] * inv; }
}
// ---------------------------------------------------------------- elementwise add: out0 = in0 + in1 (n items)
@compute @workgroup_size(256)
fn add(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let i = gid1(wg, nw, li);
if (i < u(0u)) { out0[i] = in0[i] + in1[i]; }
}
// ---------------------------------------------------------------- windows
// words: 0 n (output elements), 1 H, 2 W, 3 C, 4 ws, 5 windows across (nwx)
// window_partition: [H, W, C] -> [nwin, ws, ws, C] (zero padded); window_unpartition: the inverse.
@compute @workgroup_size(256)
fn window_partition(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let i = gid1(wg, nw, li);
if (i >= u(0u)) { return; }
let H = u(1u); let W = u(2u); let C = u(3u); let ws = u(4u); let nwx = u(5u);
let c = i % C; var t = i / C;
let x = t % ws; t = t / ws;
let y = t % ws; let win = t / ws;
let gy = (win / nwx) * ws + y;
let gx = (win % nwx) * ws + x;
var v = 0.0;
if (gy < H && gx < W) { v = in0[(gy * W + gx) * C + c]; }
out0[i] = v;
}
@compute @workgroup_size(256)
fn window_unpartition(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let i = gid1(wg, nw, li);
if (i >= u(0u)) { return; }
let W = u(2u); let C = u(3u); let ws = u(4u); let nwx = u(5u);
let c = i % C; let t = i / C;
let gx = t % W; let gy = t / W;
let win = (gy / ws) * nwx + gx / ws;
out0[i] = in0[((win * ws + gy % ws) * ws + gx % ws) * C + c];
}
// ---------------------------------------------------------------- 2 x 2 max pool over [nb, h, w, c] tokens
// words: 0 n (output elements), 1 nb, 2 h, 3 w, 4 c, 5 src image stride, 6 src row stride, 7 src token stride
// (so it can read the q third of a qkv buffer in place). Output [nb, h/2, w/2, c] contiguous.
@compute @workgroup_size(256)
fn maxpool2(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let i = gid1(wg, nw, li);
if (i >= u(0u)) { return; }
let h = u(2u); let w = u(3u); let c = u(4u);
let oh = h / 2u; let ow = w / 2u;
let ch = i % c; var t = i / c;
let ox = t % ow; t = t / ow;
let oy = t % oh; let img = t / oh;
let base = img * u(5u) + (oy * 2u) * u(6u) + (ox * 2u) * u(7u) + ch;
out0[i] = max(max(in0[base], in0[base + u(7u)]), max(in0[base + u(6u)], in0[base + u(6u) + u(7u)]));
}
// ---------------------------------------------------------------- patch embedding: 7 x 7 stride 4 pad 3 convolution
// words: 0 n (output elements), 1 H, 2 W, 3 Co, 4 Ho, 5 Wo. in0 pixels [3, H, W], in1 weight [Co, 3, 7, 7],
// in2 bias [Co]. Output [Ho, Wo, Co].
@compute @workgroup_size(256)
fn patch_embed(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let i = gid1(wg, nw, li);
if (i >= u(0u)) { return; }
let H = u(1u); let W = u(2u); let Co = u(3u); let Wo = u(5u);
let co = i % Co; let t = i / Co;
let ox = t % Wo; let oy = t / Wo;
var acc = in2[co];
for (var ci = 0u; ci < 3u; ci++) {
for (var ky = 0u; ky < 7u; ky++) {
let iy = i32(oy * 4u + ky) - 3;
if (iy < 0 || iy >= i32(H)) { continue; }
for (var kx = 0u; kx < 7u; kx++) {
let ix = i32(ox * 4u + kx) - 3;
if (ix < 0 || ix >= i32(W)) { continue; }
acc += in0[ci * H * W + u32(iy) * W + u32(ix)] * in1[((co * 3u + ci) * 7u + ky) * 7u + kx];
}
}
}
out0[i] = acc;
}
// ---------------------------------------------------------------- rotary position encoding on [rows, cols]
// words: 0 n (= rows * cols), 1 cols, 2 table rows, 3 rows rotated (the rest is copied). in0 x, in1 cos, in2 sin
// (tables [table rows, cols]); pairs (x0, x1) -> (x0 c - x1 s, x1 c + x0 s).
@compute @workgroup_size(256)
fn rope(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let i = gid1(wg, nw, li);
if (i >= u(0u)) { return; }
let cols = u(1u);
let row = i / cols; let col = i % cols;
if (row >= u(3u)) { out0[i] = in0[i]; return; }
let t = (row % u(2u)) * cols + col;
let x = in0[i];
var pair = 0.0;
if ((col & 1u) == 1u) { pair = in0[i - 1u]; } else { pair = in0[i + 1u]; }
if ((col & 1u) == 0u) { out0[i] = x * in1[t] - pair * in2[t]; } else { out0[i] = x * in1[t] + pair * in2[t]; }
}
// ---------------------------------------------------------------- neck: lateral + nearest 2x upsample of prev
// words: 0 n (output elements), 1 h, 2 w, 3 c. in0 lateral [h, w, c], in1 prev [h/2, w/2, c].
@compute @workgroup_size(256)
fn upsample_add(@builtin(workgroup_id) wg: vec3<u32>, @builtin(num_workgroups) nw: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
let i = gid1(wg, nw, li);
if (i >= u(0u)) { return; }
let w = u(2u); let c = u(3u);
let ch = i % c; let t = i / c;
let x = t % w; let y = t / w;
out0[i] = in0[i] + in1[(((y / 2u) * (w / 2u)) + x / 2u) * c + ch];
}
crates/sam2/src/gpu/mod.rs
+376
View File
@@ -0,0 +1,376 @@
// Original work for EffectCraft (wgpu compute kernels); the maths follows the Hugging Face
// `transformers` SAM 2 implementation (Apache-2.0, see NOTICE).
//! An optional GPU path for the heavy parts of SAM 2 (the Hiera backbone, the neck and the memory
//! attention), as wgpu compute kernels: it runs on AMD, NVIDIA and Intel alike, where candle has
//! no backend for the first. The CPU path stays the reference: every kernel is tested against it,
//! and any GPU failure (no adapter, an out-of-memory error, a lost device) is an `Err` the caller
//! answers by running the CPU path.
use std::collections::HashMap;
use std::sync::{Arc, Mutex, PoisonError};
use wgpu::util::DeviceExt;
use crate::{Error, Result};
const SHADER: &str = concat!(include_str!("kernels.wgsl"), include_str!("kernels_ops.wgsl"));
/// Words in a kernel's parameter block.
const PARAM_WORDS: usize = 24;
/// A buffer of f32 on the GPU (cheap to clone: it is a handle).
#[derive(Clone)]
pub struct Buf {
pub(crate) buffer: wgpu::Buffer,
/// Elements (f32).
pub len: usize,
}
/// One strided matrix operand: `(offset, batch stride, row stride, column stride)` in elements.
#[derive(Clone, Copy, Debug)]
pub struct Strides {
pub offset: usize,
pub batch: usize,
pub row: usize,
pub col: usize,
}
impl Strides {
/// A row-major `[rows, cols]` matrix (one batch).
pub fn row_major(cols: usize) -> Strides {
Strides { offset: 0, batch: 0, row: cols, col: 1 }
}
}
/// What to do to a matrix product after the multiplication.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub enum Act {
#[default]
None,
Gelu,
Relu,
}
/// A batched matrix product `C[b] = act(alpha * A[b] . B[b] + bias) (+ residual)`.
pub struct Matmul<'a> {
pub m: usize,
pub n: usize,
pub k: usize,
pub batch: usize,
pub a: (&'a Buf, Strides),
pub b: (&'a Buf, Strides),
pub bias: Option<&'a Buf>,
pub c: (&'a Buf, Strides),
pub residual: Option<&'a Buf>,
pub alpha: f32,
pub act: Act,
}
/// The GPU device and its compiled kernels.
pub struct Gpu {
device: wgpu::Device,
queue: wgpu::Queue,
module: wgpu::ShaderModule,
pipelines: Mutex<HashMap<&'static str, Arc<wgpu::ComputePipeline>>>,
encoder: Mutex<Option<wgpu::CommandEncoder>>,
error: Arc<Mutex<Option<String>>>,
adapter: String,
}
fn words(v: &[u32]) -> [u32; PARAM_WORDS] {
let mut out = [0u32; PARAM_WORDS];
for (o, x) in out.iter_mut().zip(v) {
*o = *x;
}
out
}
fn u32_of(v: usize, what: &str) -> Result<u32> {
u32::try_from(v).map_err(|_| Error::Compute(format!("{what} {v} does not fit a GPU index")))
}
impl Gpu {
/// A device on the best real GPU (a software adapter is refused: the CPU path is faster).
pub fn new() -> Result<Gpu> {
let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
let adapter = pollster::block_on(
instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }),
)
.map_err(|e| Error::Compute(format!("no GPU adapter: {e}")))?;
let info = adapter.get_info();
if info.device_type == wgpu::DeviceType::Cpu {
return Err(Error::Compute(format!("only a software adapter ({}) is available", info.name)));
}
let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
label: Some("effectcraft sam2"),
required_limits: adapter.limits(),
..Default::default()
}))
.map_err(|e| Error::Compute(format!("no GPU device: {e}")))?;
let error: Arc<Mutex<Option<String>>> = Arc::default();
let sink = error.clone();
device.on_uncaptured_error(Arc::new(move |e| {
log::error!("sam2 gpu: {e}");
sink.lock().unwrap_or_else(PoisonError::into_inner).get_or_insert_with(|| e.to_string());
}));
let sink = error.clone();
device.set_device_lost_callback(move |reason, message| {
sink.lock().unwrap_or_else(PoisonError::into_inner).get_or_insert_with(|| format!("GPU device lost ({reason:?}): {message}"));
});
let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("sam2 kernels"), source: wgpu::ShaderSource::Wgsl(SHADER.into()) });
Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), error, adapter: info.name })
}
pub fn adapter_name(&self) -> &str {
&self.adapter
}
fn take_error(&self) -> Result<()> {
match self.error.lock().unwrap_or_else(PoisonError::into_inner).take() {
Some(e) => Err(Error::Compute(format!("GPU: {e}"))),
None => Ok(()),
}
}
pub fn upload(&self, data: &[f32]) -> Buf {
let buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: None,
contents: bytemuck::cast_slice(data),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
});
Buf { buffer, len: data.len() }
}
/// A zeroed buffer of `len` elements.
pub fn zeros(&self, len: usize) -> Result<Buf> {
let bytes = (len.max(1) as u64).checked_mul(4).ok_or_else(|| Error::Budget("GPU buffer size overflows".into()))?;
if bytes > self.device.limits().max_buffer_size || bytes > self.device.limits().max_storage_buffer_binding_size {
return Err(Error::Budget(format!("a {} MB GPU buffer is over this GPU's limit", bytes >> 20)));
}
let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
label: None,
size: bytes,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
Ok(Buf { buffer, len })
}
fn pipeline(&self, entry: &'static str) -> Arc<wgpu::ComputePipeline> {
let mut map = self.pipelines.lock().unwrap_or_else(PoisonError::into_inner);
map.entry(entry)
.or_insert_with(|| {
Arc::new(self.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(entry),
layout: None,
module: &self.module,
entry_point: Some(entry),
compilation_options: Default::default(),
cache: None,
}))
})
.clone()
}
/// Queue one dispatch of `entry`. `bindings` are the kernel's buffers by binding index.
pub(crate) fn run(&self, entry: &'static str, bindings: &[(u32, &Buf)], params: &[u32], groups: (u32, u32, u32)) {
let pipeline = self.pipeline(entry);
let uniform = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: None,
contents: bytemuck::cast_slice(&words(params)),
usage: wgpu::BufferUsages::UNIFORM,
});
let mut entries: Vec<wgpu::BindGroupEntry<'_>> =
bindings.iter().map(|(i, b)| wgpu::BindGroupEntry { binding: *i, resource: b.buffer.as_entire_binding() }).collect();
entries.push(wgpu::BindGroupEntry { binding: 5, resource: uniform.as_entire_binding() });
let group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { label: None, layout: &pipeline.get_bind_group_layout(0), entries: &entries });
let mut slot = self.encoder.lock().unwrap_or_else(PoisonError::into_inner);
let enc = slot.get_or_insert_with(|| self.device.create_command_encoder(&Default::default()));
let mut pass = enc.begin_compute_pass(&Default::default());
pass.set_pipeline(&pipeline);
pass.set_bind_group(0, &group, &[]);
pass.dispatch_workgroups(groups.0, groups.1, groups.2);
}
/// Submit what is queued and wait for it to finish.
pub fn flush(&self) -> Result<()> {
let enc = self.encoder.lock().unwrap_or_else(PoisonError::into_inner).take();
if let Some(enc) = enc {
self.queue.submit([enc.finish()]);
}
let _ = self.device.poll(wgpu::PollType::wait_indefinitely());
self.take_error()
}
/// Copy a buffer back to the CPU (after submitting everything queued).
pub fn download(&self, buf: &Buf) -> Result<Vec<f32>> {
let bytes = buf.len as u64 * 4;
if bytes == 0 {
return Ok(vec![]);
}
let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
label: None,
size: bytes,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let mut enc =
self.encoder.lock().unwrap_or_else(PoisonError::into_inner).take().unwrap_or_else(|| self.device.create_command_encoder(&Default::default()));
enc.copy_buffer_to_buffer(&buf.buffer, 0, &staging, 0, bytes);
self.queue.submit([enc.finish()]);
let (tx, rx) = std::sync::mpsc::channel();
staging.map_async(wgpu::MapMode::Read, .., move |r| {
let _ = tx.send(r);
});
let _ = self.device.poll(wgpu::PollType::wait_indefinitely());
self.take_error()?;
rx.recv().map_err(|_| Error::Compute("GPU readback never completed".into()))?.map_err(|e| Error::Compute(format!("GPU readback: {e}")))?;
let out = {
let view = staging.get_mapped_range(..).map_err(|e| Error::Compute(format!("GPU readback: {e}")))?;
bytemuck::cast_slice::<u8, f32>(&view).to_vec()
};
staging.unmap();
Ok(out)
}
/// Queue `C[b] = act(alpha * A[b] . B[b] + bias) (+ residual)`.
pub fn matmul(&self, op: &Matmul<'_>) -> Result<()> {
let (a, sa) = op.a;
let (b, sb) = op.b;
let (c, sc) = op.c;
if op.m == 0 || op.n == 0 || op.batch == 0 {
return Ok(());
}
// bounds: the last element each operand touches must lie inside its buffer
let last = |s: Strides, rows: usize, cols: usize| s.offset + (op.batch - 1) * s.batch + (rows - 1) * s.row + (cols.max(1) - 1) * s.col;
for (name, buf, s, rows, cols) in [("A", a, sa, op.m, op.k), ("B", b, sb, op.k, op.n), ("C", c, sc, op.m, op.n)] {
if op.k > 0 && last(s, rows, cols) >= buf.len {
return Err(Error::Input(format!("matmul operand {name} reaches element {} of a {}-element buffer", last(s, rows, cols), buf.len)));
}
}
if op.bias.is_some_and(|bias| bias.len < op.n) || op.residual.is_some_and(|r| r.len < c.len) {
return Err(Error::Input("matmul bias or residual is too small".into()));
}
let flags = u32::from(op.bias.is_some())
| (match op.act {
Act::None => 0,
Act::Gelu => 1,
Act::Relu => 2,
} << 1)
| (u32::from(op.residual.is_some()) << 3);
let params = [
u32_of(op.m, "m")?,
u32_of(op.n, "n")?,
u32_of(op.k, "k")?,
u32_of(op.batch, "batch")?,
u32_of(sa.offset, "a offset")?,
u32_of(sa.batch, "a batch")?,
u32_of(sa.row, "a row")?,
u32_of(sa.col, "a col")?,
u32_of(sb.offset, "b offset")?,
u32_of(sb.batch, "b batch")?,
u32_of(sb.row, "b row")?,
u32_of(sb.col, "b col")?,
u32_of(sc.offset, "c offset")?,
u32_of(sc.batch, "c batch")?,
u32_of(sc.row, "c row")?,
flags,
op.alpha.to_bits(),
];
let dummy;
let bias = match op.bias {
Some(b) => b,
None => {
dummy = self.upload(&[0.0]);
&dummy
}
};
// a buffer is never bound twice in one dispatch (read and read-write), so no residual means a dummy
let dummy_res;
let res = match op.residual {
Some(r) => r,
None => {
dummy_res = self.upload(&[0.0]);
&dummy_res
}
};
self.run(
"matmul",
&[(0, a), (1, b), (2, bias), (3, c), (4, res)],
&params,
(u32_of(op.n.div_ceil(64), "groups")?, u32_of(op.m.div_ceil(64), "groups")?, u32_of(op.batch, "batch")?),
);
Ok(())
}
}
pub mod encoder;
mod ops;
/// A linear layer's weight `[out, in]` and bias `[out]` on the GPU.
pub struct GLinear {
pub(crate) w: Buf,
pub(crate) b: Buf,
pub(crate) inp: usize,
pub(crate) out: usize,
}
/// A layer norm's scale and shift on the GPU.
pub struct GNorm {
pub(crate) w: Buf,
pub(crate) b: Buf,
}
fn flat(t: candle_core::Tensor) -> Result<Vec<f32>> {
Ok(t.flatten_all()?.to_vec1::<f32>()?)
}
impl GLinear {
/// From `{name}.weight` (`[out, in]`) and `{name}.bias`.
pub fn load(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result<GLinear> {
Ok(GLinear {
w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?),
b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?),
inp,
out,
})
}
/// A 1 x 1 convolution (`[out, in, 1, 1]`) as a linear layer over channels-last tokens.
pub fn load_conv1x1(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result<GLinear> {
Ok(GLinear {
w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp, 1, 1])?)?),
b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?),
inp,
out,
})
}
}
impl GNorm {
pub fn load(g: &Gpu, w: &crate::weights::Weights, name: &str, dim: usize) -> Result<GNorm> {
Ok(GNorm { w: g.upload(&flat(w.get(&format!("{name}.weight"), &[dim])?)?), b: g.upload(&flat(w.get(&format!("{name}.bias"), &[dim])?)?) })
}
}
impl Gpu {
/// `out = act(x . W^T + bias) (+ residual)` for `rows` rows of `l.inp` values.
pub fn linear(&self, x: &Buf, rows: usize, l: &GLinear, out: &Buf, act: Act, residual: Option<&Buf>) -> Result<()> {
// the weight is stored [out, in]: B[k][j] = w[j * in + k]
self.matmul(&Matmul {
m: rows,
n: l.out,
k: l.inp,
batch: 1,
a: (x, Strides::row_major(l.inp)),
b: (&l.w, Strides { offset: 0, batch: 0, row: 1, col: l.inp }),
bias: Some(&l.b),
c: (out, Strides::row_major(l.out)),
residual,
alpha: 1.0,
act,
})
}
}
#[cfg(test)]
mod tests;
crates/sam2/src/gpu/ops.rs
+167
View File
@@ -0,0 +1,167 @@
//! The non-matmul kernels, as methods of [`Gpu`] (each queues one dispatch).
use super::{Buf, Gpu};
use crate::{Error, Result};
fn n32(v: usize, what: &str) -> Result<u32> {
u32::try_from(v).map_err(|_| Error::Compute(format!("{what} {v} does not fit a GPU index")))
}
/// A 2-D grid of groups covering `groups` (a dimension holds 65535).
fn grid(groups: usize) -> Result<(u32, u32, u32)> {
let gx = groups.clamp(1, 65535);
Ok((n32(gx, "groups")?, n32(groups.div_ceil(gx).max(1), "groups")?, 1))
}
/// One 256-thread group per 256 items.
fn items(n: usize) -> Result<(u32, u32, u32)> {
grid(n.div_ceil(256))
}
fn need(buf: &Buf, len: usize, what: &str) -> Result<()> {
if buf.len < len {
return Err(Error::Input(format!("{what}: a buffer of {} elements is smaller than the {len} needed", buf.len)));
}
Ok(())
}
impl Gpu {
/// `y = layer_norm(x) * weight + bias` over the last dimension of a `[rows, cols]` buffer.
pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32) -> Result<()> {
need(x, rows * cols, "layernorm x")?;
need(y, rows * cols, "layernorm y")?;
need(weight, cols, "layernorm weight")?;
need(bias, cols, "layernorm bias")?;
self.run("layernorm", &[(0, x), (1, weight), (2, bias), (3, y)], &[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits()], grid(rows)?);
Ok(())
}
/// Softmax over each of `rows` rows of `cols`, in place.
pub fn softmax(&self, x: &Buf, rows: usize, cols: usize) -> Result<()> {
need(x, rows * cols, "softmax")?;
self.run("softmax", &[(3, x)], &[n32(rows, "rows")?, n32(cols, "cols")?], grid(rows)?);
Ok(())
}
/// `out = a + b` over `n` elements.
pub fn add(&self, a: &Buf, b: &Buf, out: &Buf, n: usize) -> Result<()> {
need(a, n, "add a")?;
need(b, n, "add b")?;
need(out, n, "add out")?;
self.run("add", &[(0, a), (1, b), (3, out)], &[n32(n, "n")?], items(n)?);
Ok(())
}
/// `[h, w, c]` to `[windows, ws, ws, c]`, zero padded at the bottom and right.
pub fn window_partition(&self, x: &Buf, out: &Buf, h: usize, w: usize, c: usize, ws: usize) -> Result<()> {
if ws == 0 {
return Err(Error::Input("window size 0".into()));
}
let (nwx, nwy) = (w.div_ceil(ws), h.div_ceil(ws));
let n = nwx * nwy * ws * ws * c;
need(x, h * w * c, "window_partition input")?;
need(out, n, "window_partition output")?;
self.run(
"window_partition",
&[(0, x), (3, out)],
&[n32(n, "n")?, n32(h, "h")?, n32(w, "w")?, n32(c, "c")?, n32(ws, "ws")?, n32(nwx, "windows")?],
items(n)?,
);
Ok(())
}
/// The inverse of [`Gpu::window_partition`] (the padding is dropped).
pub fn window_unpartition(&self, x: &Buf, out: &Buf, h: usize, w: usize, c: usize, ws: usize) -> Result<()> {
if ws == 0 {
return Err(Error::Input("window size 0".into()));
}
let (nwx, nwy) = (w.div_ceil(ws), h.div_ceil(ws));
need(x, nwx * nwy * ws * ws * c, "window_unpartition input")?;
need(out, h * w * c, "window_unpartition output")?;
self.run(
"window_unpartition",
&[(0, x), (3, out)],
&[n32(h * w * c, "n")?, n32(h, "h")?, n32(w, "w")?, n32(c, "c")?, n32(ws, "ws")?, n32(nwx, "windows")?],
items(h * w * c)?,
);
Ok(())
}
/// 2 x 2 max pool of `nb` images of `[h, w, c]` tokens (`h` and `w` even) read with the given strides
/// (image, row, token) so it can pool one third of a qkv buffer in place; output contiguous.
#[allow(clippy::too_many_arguments)]
pub fn maxpool2(&self, x: &Buf, out: &Buf, nb: usize, h: usize, w: usize, c: usize, strides: (usize, usize, usize)) -> Result<()> {
if !h.is_multiple_of(2) || !w.is_multiple_of(2) {
return Err(Error::Input(format!("maxpool2 needs even sizes, got {w}x{h}")));
}
let n = nb * (h / 2) * (w / 2) * c;
let last = (nb.saturating_sub(1)) * strides.0 + (h - 1) * strides.1 + (w - 1) * strides.2 + c;
need(x, last, "maxpool2 input")?;
need(out, n, "maxpool2 output")?;
self.run(
"maxpool2",
&[(0, x), (3, out)],
&[
n32(n, "n")?,
n32(nb, "nb")?,
n32(h, "h")?,
n32(w, "w")?,
n32(c, "c")?,
n32(strides.0, "stride")?,
n32(strides.1, "stride")?,
n32(strides.2, "stride")?,
],
items(n)?,
);
Ok(())
}
/// The patch embedding: `[3, h, w]` pixels, weight `[co, 3, 7, 7]`, bias `[co]`; output `[h/4, w/4, co]`.
pub fn patch_embed(&self, pixels: &Buf, weight: &Buf, bias: &Buf, out: &Buf, h: usize, w: usize, co: usize) -> Result<()> {
let (ho, wo) = ((h + 6 - 7) / 4 + 1, (w + 6 - 7) / 4 + 1);
need(pixels, 3 * h * w, "patch_embed pixels")?;
need(weight, co * 3 * 49, "patch_embed weight")?;
need(bias, co, "patch_embed bias")?;
need(out, ho * wo * co, "patch_embed output")?;
let n = ho * wo * co;
self.run(
"patch_embed",
&[(0, pixels), (1, weight), (2, bias), (3, out)],
&[n32(n, "n")?, n32(h, "h")?, n32(w, "w")?, n32(co, "co")?, n32(ho, "ho")?, n32(wo, "wo")?],
items(n)?,
);
Ok(())
}
/// Rotary encoding of `[rows, cols]`: rows below `rotated` are rotated with the `[table_rows, cols]`
/// tables (row `r` uses table row `r % table_rows`), the rest copied.
#[allow(clippy::too_many_arguments)]
pub fn rope(&self, x: &Buf, cos: &Buf, sin: &Buf, out: &Buf, rows: usize, cols: usize, table_rows: usize, rotated: usize) -> Result<()> {
if !cols.is_multiple_of(2) || table_rows == 0 {
return Err(Error::Input("rope needs an even width and a table".into()));
}
need(x, rows * cols, "rope input")?;
need(out, rows * cols, "rope output")?;
need(cos, table_rows * cols, "rope cos")?;
need(sin, table_rows * cols, "rope sin")?;
self.run(
"rope",
&[(0, x), (1, cos), (2, sin), (3, out)],
&[n32(rows * cols, "n")?, n32(cols, "cols")?, n32(table_rows, "table")?, n32(rotated.min(rows), "rotated")?],
items(rows * cols)?,
);
Ok(())
}
/// `out = lateral + nearest_upsample_2x(prev)` for `[h, w, c]` (`h` and `w` even).
pub fn upsample_add(&self, lateral: &Buf, prev: &Buf, out: &Buf, h: usize, w: usize, c: usize) -> Result<()> {
if !h.is_multiple_of(2) || !w.is_multiple_of(2) {
return Err(Error::Input(format!("upsample_add needs even sizes, got {w}x{h}")));
}
need(lateral, h * w * c, "upsample_add lateral")?;
need(prev, h / 2 * (w / 2) * c, "upsample_add prev")?;
need(out, h * w * c, "upsample_add out")?;
self.run("upsample_add", &[(0, lateral), (1, prev), (3, out)], &[n32(h * w * c, "n")?, n32(h, "h")?, n32(w, "w")?, n32(c, "c")?], items(h * w * c)?);
Ok(())
}
}
crates/sam2/src/gpu/tests.rs
+301
View File
@@ -0,0 +1,301 @@
//! Each kernel against a plain CPU computation. They skip when there is no GPU.
use std::sync::{Mutex, PoisonError};
use super::*;
/// The one device the tests share (several devices at once upset the Vulkan loader), held for the whole
/// test so the tests take turns.
pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, &'static Gpu);
impl std::ops::Deref for Shared {
type Target = Gpu;
fn deref(&self) -> &Gpu {
self.1
}
}
pub(crate) fn gpu() -> Option<Shared> {
static TURN: Mutex<()> = Mutex::new(());
static DEVICE: std::sync::OnceLock<Option<Gpu>> = std::sync::OnceLock::new();
let guard = TURN.lock().unwrap_or_else(PoisonError::into_inner);
let device = DEVICE.get_or_init(|| match Gpu::new() {
Ok(g) => Some(g),
Err(e) => {
eprintln!("skipped: {e}");
None
}
});
device.as_ref().map(|g| Shared(guard, g))
}
/// Deterministic pseudo-random values in -1..1.
pub(crate) fn noise(n: usize, seed: u32) -> Vec<f32> {
let mut s = seed.wrapping_mul(2654435761).wrapping_add(1);
(0..n)
.map(|_| {
s ^= s << 13;
s ^= s >> 17;
s ^= s << 5;
(s % 20001) as f32 / 10000.0 - 1.0
})
.collect()
}
pub(crate) fn assert_close(got: &[f32], want: &[f32], tol: f32, what: &str) {
assert_eq!(got.len(), want.len(), "{what}: length");
let worst = got.iter().zip(want).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
assert!(worst <= tol, "{what}: max abs error {worst} (tolerance {tol})");
}
fn gelu(x: f32) -> f32 {
use candle_core::{Device, Tensor};
let t = Tensor::new(&[x], &Device::Cpu).and_then(|t| t.gelu_erf()).and_then(|t| t.to_vec1::<f32>());
t.map(|v| v[0]).unwrap_or(f32::NAN)
}
/// `C[b][i][j] = sum_k A[b][i][k] * B[b][k][j]` by the definition, with the strides given.
#[allow(clippy::too_many_arguments)]
fn naive(a: &[f32], sa: Strides, b: &[f32], sb: Strides, m: usize, n: usize, k: usize, batch: usize) -> Vec<f32> {
let mut c = vec![0f32; batch * m * n];
for z in 0..batch {
for i in 0..m {
for j in 0..n {
let mut acc = 0f32;
for l in 0..k {
acc += a[sa.offset + z * sa.batch + i * sa.row + l * sa.col] * b[sb.offset + z * sb.batch + l * sb.row + j * sb.col];
}
c[z * m * n + i * n + j] = acc;
}
}
}
c
}
#[test]
fn matmul_matches_the_definition_for_awkward_sizes_and_strides() {
let Some(g) = gpu() else { return };
// (m, n, k, batch): not multiples of the 64 x 64 x 16 tile, and degenerate ones
for (m, n, k, batch) in [(100, 50, 37, 1), (1, 1, 1, 1), (65, 129, 17, 2), (64, 64, 16, 1), (3, 200, 96, 3), (130, 7, 300, 1)] {
let (av, bv) = (noise(batch * m * k, 1), noise(batch * k * n, 2));
let (a, b) = (g.upload(&av), g.upload(&bv));
let c = g.zeros(batch * m * n).unwrap();
let (sa, sb) = (Strides { offset: 0, batch: m * k, row: k, col: 1 }, Strides { offset: 0, batch: k * n, row: n, col: 1 });
let sc = Strides { offset: 0, batch: m * n, row: n, col: 1 };
g.matmul(&Matmul { m, n, k, batch, a: (&a, sa), b: (&b, sb), bias: None, c: (&c, sc), residual: None, alpha: 1.0, act: Act::None }).unwrap();
g.flush().unwrap();
assert_close(&g.download(&c).unwrap(), &naive(&av, sa, &bv, sb, m, n, k, batch), 2e-4 * (k as f32).sqrt(), &format!("{m}x{n}x{k} x{batch}"));
}
}
#[test]
fn matmul_with_transposed_operands_bias_activation_and_residual() {
let Some(g) = gpu() else { return };
let (m, n, k) = (70, 90, 41);
let (av, wv, bias, res) = (noise(m * k, 3), noise(n * k, 4), noise(n, 5), noise(m * n, 6));
let (a, w, bi, r) = (g.upload(&av), g.upload(&wv), g.upload(&bias), g.upload(&res));
let c = g.zeros(m * n).unwrap();
// B is a weight stored [n, k] (a linear layer): B[k][j] = w[j * k + l]
let sb = Strides { offset: 0, batch: 0, row: 1, col: k };
let sa = Strides::row_major(k);
g.matmul(&Matmul {
m,
n,
k,
batch: 1,
a: (&a, sa),
b: (&w, sb),
bias: Some(&bi),
c: (&c, Strides::row_major(n)),
residual: Some(&r),
alpha: 0.5,
act: Act::Gelu,
})
.unwrap();
g.flush().unwrap();
let base = naive(&av, sa, &wv, sb, m, n, k, 1);
let want: Vec<f32> = base.iter().enumerate().map(|(i, v)| gelu(v * 0.5 + bias[i % n]) + res[i]).collect();
assert_close(&g.download(&c).unwrap(), &want, 1e-3, "linear + gelu + residual");
// relu, and a transposed A (m contiguous)
let at: Vec<f32> = (0..k * m).map(|i| av[(i % m) * k + i / m]).collect();
let (a2, c2) = (g.upload(&at), g.zeros(m * n).unwrap());
g.matmul(&Matmul {
m,
n,
k,
batch: 1,
a: (&a2, Strides { offset: 0, batch: 0, row: 1, col: m }),
b: (&w, sb),
bias: None,
c: (&c2, Strides::row_major(n)),
residual: None,
alpha: 1.0,
act: Act::Relu,
})
.unwrap();
g.flush().unwrap();
let want: Vec<f32> = base.iter().map(|v| v.max(0.0)).collect();
assert_close(&g.download(&c2).unwrap(), &want, 1e-3, "transposed A, relu");
}
#[test]
fn matmul_refuses_operands_that_overrun_their_buffers() {
let Some(g) = gpu() else { return };
let (a, b, c) = (g.upload(&[1.0; 6]), g.upload(&[1.0; 6]), g.zeros(4).unwrap());
let ok = Matmul {
m: 2,
n: 2,
k: 3,
batch: 1,
a: (&a, Strides::row_major(3)),
b: (&b, Strides::row_major(2)),
bias: None,
c: (&c, Strides::row_major(2)),
residual: None,
alpha: 1.0,
act: Act::None,
};
assert!(g.matmul(&ok).is_ok());
let too_big = Matmul { m: 2, n: 2, k: 4, ..ok };
assert!(matches!(g.matmul(&too_big), Err(Error::Input(_))));
assert!(g.zeros(usize::MAX / 2).is_err());
g.flush().unwrap();
}
// ---------------------------------------------------------------- the other kernels
fn run_and_read(g: &Gpu, out: &Buf) -> Vec<f32> {
g.flush().unwrap();
g.download(out).unwrap()
}
#[test]
fn layernorm_and_softmax_match_the_definition() {
let Some(g) = gpu() else { return };
for (rows, cols) in [(7, 96), (1, 1), (300, 768), (70_000, 8)] {
let x = noise(rows * cols, 7);
let (w, b) = (noise(cols, 8), noise(cols, 9));
let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&w), g.upload(&b), g.zeros(rows * cols).unwrap());
g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6).unwrap();
let want: Vec<f32> = x
.chunks(cols)
.flat_map(|r| {
let mean = r.iter().sum::<f32>() / cols as f32;
let var = r.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / cols as f32;
r.iter().enumerate().map(|(i, v)| (v - mean) / (var + 1e-6).sqrt() * w[i] + b[i]).collect::<Vec<_>>()
})
.collect();
assert_close(&run_and_read(&g, &gy), &want, 2e-4, &format!("layernorm {rows}x{cols}"));
}
for (rows, cols) in [(5, 64), (3, 1), (100, 4096), (2, 28736), (70_000, 64)] {
let x = noise(rows * cols, 10).iter().map(|v| v * 6.0).collect::<Vec<_>>();
let gx = g.upload(&x);
g.softmax(&gx, rows, cols).unwrap();
let want: Vec<f32> = x
.chunks(cols)
.flat_map(|r| {
let m = r.iter().cloned().fold(f32::MIN, f32::max);
let e: Vec<f32> = r.iter().map(|v| (v - m).exp()).collect();
let s: f32 = e.iter().sum();
e.into_iter().map(|v| v / s).collect::<Vec<_>>()
})
.collect();
assert_close(&run_and_read(&g, &gx), &want, 1e-5, &format!("softmax {rows}x{cols}"));
}
}
#[test]
fn windows_pooling_and_add_match_naive_loops() {
let Some(g) = gpu() else { return };
let (h, w, c, ws) = (13, 22, 5, 8);
let x = noise(h * w * c, 11);
let (nwx, nwy) = (w.div_ceil(ws), h.div_ceil(ws));
let (gx, parts, back) = (g.upload(&x), g.zeros(nwx * nwy * ws * ws * c).unwrap(), g.zeros(h * w * c).unwrap());
g.window_partition(&gx, &parts, h, w, c, ws).unwrap();
g.window_unpartition(&parts, &back, h, w, c, ws).unwrap();
let mut want = vec![0f32; nwx * nwy * ws * ws * c];
for win in 0..nwx * nwy {
for y in 0..ws {
for xx in 0..ws {
let (gy, gxx) = ((win / nwx) * ws + y, (win % nwx) * ws + xx);
for ch in 0..c {
if gy < h && gxx < w {
want[((win * ws + y) * ws + xx) * c + ch] = x[(gy * w + gxx) * c + ch];
}
}
}
}
}
g.flush().unwrap();
assert_close(&g.download(&parts).unwrap(), &want, 0.0, "window_partition");
assert_close(&g.download(&back).unwrap(), &x, 0.0, "window round trip");
assert!(g.window_partition(&gx, &parts, h, w, c, 0).is_err());
// pooling a q third of a [nb, 8, 6, 3 * c] buffer in place
let (nb, ph, pw, pc) = (3, 8, 6, 4);
let qkv = noise(nb * ph * pw * 3 * pc, 12);
let (gq, out) = (g.upload(&qkv), g.zeros(nb * (ph / 2) * (pw / 2) * pc).unwrap());
g.maxpool2(&gq, &out, nb, ph, pw, pc, (ph * pw * 3 * pc, pw * 3 * pc, 3 * pc)).unwrap();
let mut want = vec![];
for img in 0..nb {
for oy in 0..ph / 2 {
for ox in 0..pw / 2 {
for ch in 0..pc {
let at = |y: usize, xx: usize| qkv[((img * ph + y) * pw + xx) * 3 * pc + ch];
want.push(at(2 * oy, 2 * ox).max(at(2 * oy, 2 * ox + 1)).max(at(2 * oy + 1, 2 * ox)).max(at(2 * oy + 1, 2 * ox + 1)));
}
}
}
}
assert_close(&run_and_read(&g, &out), &want, 0.0, "maxpool2");
assert!(g.maxpool2(&gq, &out, nb, 7, pw, pc, (1, 1, 1)).is_err());
let (a, b) = (noise(100_001, 13), noise(100_001, 14));
let (ga, gb, go) = (g.upload(&a), g.upload(&b), g.zeros(a.len()).unwrap());
g.add(&ga, &gb, &go, a.len()).unwrap();
let want: Vec<f32> = a.iter().zip(&b).map(|(x, y)| x + y).collect();
assert_close(&run_and_read(&g, &go), &want, 0.0, "add");
}
#[test]
fn patch_embed_rope_and_upsample_match_the_cpu_definitions() {
use candle_core::{Device, Tensor};
let Some(g) = gpu() else { return };
// the 7 x 7 stride 4 convolution against candle's conv2d
let (h, w, co) = (64usize, 96usize, 8usize);
let (px, wt, bias) = (noise(3 * h * w, 15), noise(co * 147, 16), noise(co, 17));
let (gp, gw, gb) = (g.upload(&px), g.upload(&wt), g.upload(&bias));
let (ho, wo) = (h / 4, w / 4);
let out = g.zeros(ho * wo * co).unwrap();
g.patch_embed(&gp, &gw, &gb, &out, h, w, co).unwrap();
let t = Tensor::from_vec(px, (1, 3, h, w), &Device::Cpu).unwrap();
let k = Tensor::from_vec(wt, (co, 3, 7, 7), &Device::Cpu).unwrap();
let want =
t.conv2d(&k, 3, 4, 1, 1).unwrap().broadcast_add(&Tensor::from_vec(bias, (1, co, 1, 1), &Device::Cpu).unwrap()).unwrap().permute((0, 2, 3, 1)).unwrap();
assert_close(&run_and_read(&g, &out), &want.flatten_all().unwrap().to_vec1::<f32>().unwrap(), 2e-4, "patch_embed");
// rotary encoding: pairs rotated with row % table rows, later rows copied
let (rows, cols, table, rotated) = (37usize, 16usize, 10usize, 30usize);
let (x, cs, sn) = (noise(rows * cols, 18), noise(table * cols, 19), noise(table * cols, 20));
let (gx, gc, gs, go) = (g.upload(&x), g.upload(&cs), g.upload(&sn), g.zeros(rows * cols).unwrap());
g.rope(&gx, &gc, &gs, &go, rows, cols, table, rotated).unwrap();
let mut want = x.clone();
for r in 0..rotated {
for j in 0..cols / 2 {
let (x0, x1) = (x[r * cols + 2 * j], x[r * cols + 2 * j + 1]);
let (t0, t1) = ((r % table) * cols + 2 * j, (r % table) * cols + 2 * j + 1);
want[r * cols + 2 * j] = x0 * cs[t0] - x1 * sn[t0];
want[r * cols + 2 * j + 1] = x1 * cs[t1] + x0 * sn[t1];
}
}
assert_close(&run_and_read(&g, &go), &want, 1e-6, "rope");
assert!(g.rope(&gx, &gc, &gs, &go, rows, 15, table, rotated).is_err());
// neck: lateral + nearest upsample
let (h, w, c) = (6usize, 10usize, 7usize);
let (lat, prev) = (noise(h * w * c, 21), noise(h / 2 * (w / 2) * c, 22));
let (gl, gp2, go) = (g.upload(&lat), g.upload(&prev), g.zeros(h * w * c).unwrap());
g.upsample_add(&gl, &gp2, &go, h, w, c).unwrap();
let want: Vec<f32> = (0..h * w * c).map(|i| lat[i] + prev[((i / c / w / 2) * (w / 2) + (i / c % w) / 2) * c + i % c]).collect();
assert_close(&run_and_read(&g, &go), &want, 0.0, "upsample_add");
}
crates/sam2/src/hiera.rs
+43 -11
View File
@@ -116,6 +116,39 @@ impl Attention {
}
}
/// What a block is made of: its widths, heads, window (0: global attention) and query pooling.
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct BlockPlan {
pub dim: usize,
pub dim_out: usize,
pub heads: usize,
pub window: usize,
pub query_stride: Option<usize>,
}
impl BlockPlan {
pub(crate) fn new(cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result<BlockPlan> {
let first = stage > 0 && block_in_stage == 0;
let dim = at(&cfg.embed_dims, if first { stage - 1 } else { stage }, "embed dim")?;
let dim_out = at(&cfg.embed_dims, stage, "embed dim")?;
let window = if first { at(&cfg.windows, stage - 1, "window")? } else { at(&cfg.windows, stage, "window")? };
let window = if cfg.global_blocks.contains(&total) { 0 } else { window };
let query_stride = (stage > 0 && stage <= cfg.query_pool_stages && block_in_stage == 0).then_some(cfg.query_stride);
Ok(BlockPlan { dim, dim_out, heads: at(&cfg.heads, stage, "heads")?, window, query_stride })
}
/// Every block of the backbone, in order.
pub(crate) fn all(cfg: &HieraConfig) -> Result<Vec<BlockPlan>> {
let mut out = vec![];
for (stage, n) in cfg.blocks.iter().enumerate() {
for b in 0..*n {
out.push(BlockPlan::new(cfg, stage, b, out.len())?);
}
}
Ok(out)
}
}
/// One Hiera block: pre-norm attention (window or global, optionally pooling the queries) and MLP.
pub struct Block {
norm1: LayerNorm,
@@ -130,13 +163,8 @@ pub struct Block {
impl Block {
fn load(w: &Weights, cfg: &HieraConfig, stage: usize, block_in_stage: usize, total: usize) -> Result<Block> {
let prefix = format!("vision_encoder.backbone.blocks.{total}");
let first = stage > 0 && block_in_stage == 0;
let dim = at(&cfg.embed_dims, if first { stage - 1 } else { stage }, "embed dim")?;
let dim_out = at(&cfg.embed_dims, stage, "embed dim")?;
let window = if first { at(&cfg.windows, stage - 1, "window")? } else { at(&cfg.windows, stage, "window")? };
let window = if cfg.global_blocks.contains(&total) { 0 } else { window };
let query_stride = (stage > 0 && stage <= cfg.query_pool_stages && block_in_stage == 0).then_some(cfg.query_stride);
let heads = at(&cfg.heads, stage, "heads")?;
let plan = BlockPlan::new(cfg, stage, block_in_stage, total)?;
let (dim, dim_out, heads) = (plan.dim, plan.dim_out, plan.heads);
Ok(Block {
norm1: layer_norm(w, &format!("{prefix}.layer_norm1"), dim, cfg.eps)?,
attn: Attention {
@@ -144,13 +172,13 @@ impl Block {
proj: linear(w, &format!("{prefix}.attn.proj"), dim_out, dim_out)?,
heads,
scale: ((dim_out / heads.max(1)) as f64).powf(-0.5),
query_stride,
query_stride: plan.query_stride,
},
norm2: layer_norm(w, &format!("{prefix}.layer_norm2"), dim_out, cfg.eps)?,
mlp: Mlp::load(w, &format!("{prefix}.mlp"), dim_out, dim_out * cfg.mlp_ratio)?,
proj: (dim != dim_out).then(|| linear(w, &format!("{prefix}.proj"), dim, dim_out)).transpose()?,
window,
query_stride,
window: plan.window,
query_stride: plan.query_stride,
})
}
@@ -269,6 +297,10 @@ impl Hiera {
}
/// Patch embedding of `[B, 3, H, W]` pixels: `[B, H/4, W/4, C]` (before the positional embedding).
pub(crate) fn config(&self) -> &HieraConfig {
&self.cfg
}
pub fn patch_embed(&self, pixels: &Tensor) -> Result<Tensor> {
let x = pixels.conv2d(&self.patch_w, 3, 4, 1, 1)?;
let x = x.broadcast_add(&self.patch_b.reshape((1, self.hidden, 1, 1))?)?;
@@ -276,7 +308,7 @@ impl Hiera {
}
/// Background embedding resized to the grid plus the tiled window embedding, `[1, H, W, C]`.
fn pos(&self, h: usize, w: usize) -> Result<Tensor> {
pub(crate) fn pos(&self, h: usize, w: usize) -> Result<Tensor> {
let bg = self.cfg.pos_background;
let big = bicubic_resize(&self.pos_embed, self.hidden, (bg, bg), (h, w));
let ws = self.window0.max(1);
crates/sam2/src/lib.rs
+1
View File
@@ -12,6 +12,7 @@
use std::fmt;
pub mod decoder;
pub mod gpu;
pub mod hiera;
pub mod mask;
pub mod memattn;
crates/sam2/src/neck.rs
+2 -2
View File
@@ -129,7 +129,7 @@ impl Neck {
/// Hiera plus neck: `Sam2VisionModel`.
pub struct VisionEncoder {
pub hiera: Hiera,
pub hiera: std::sync::Arc<Hiera>,
pub neck: Neck,
levels: usize,
}
@@ -144,7 +144,7 @@ pub struct ImageFeatures {
impl VisionEncoder {
pub fn load_tiny(w: &Weights) -> Result<VisionEncoder> {
let cfg = HieraConfig::tiny();
Ok(VisionEncoder { hiera: Hiera::load(w, &cfg)?, neck: Neck::load(w, &[768, 384, 192, 96], 256, &[2, 3])?, levels: 3 })
Ok(VisionEncoder { hiera: std::sync::Arc::new(Hiera::load(w, &cfg)?), neck: Neck::load(w, &[768, 384, 192, 96], 256, &[2, 3])?, levels: 3 })
}
pub fn forward(&self, pixels: &Tensor) -> Result<ImageFeatures> {
crates/sam2/tests/reference.rs
+40
View File
@@ -592,3 +592,43 @@ fn the_tracker_stops_at_its_budget() {
let r = t.propagate(0, Direction::Forward, None, &Cancel::new(), &mut |_| {});
assert!(matches!(r, Err(Error::Budget(_))), "{r:?}");
}
/// The GPU encoder (wgpu kernels, any vendor) against the CPU path and PyTorch on the same image. The two
/// float implementations differ by accumulation order and the GPU's erf (1.5e-7), compounded through 12
/// blocks. Measured (RX 7900 GRE, 2026-10-15): 6.8e-6 against PyTorch, 1.2e-5 against the CPU path.
const GPU_TOL: f32 = 1e-4;
#[test]
fn gpu_encoder_matches_the_cpu_path_and_pytorch() {
use effectcraft_sam2::gpu::{Gpu, encoder::GpuEncoder};
use effectcraft_sam2::model::ImageModel;
let Some((w, r)) = load() else { return };
let gpu = match Gpu::new() {
Ok(g) => std::sync::Arc::new(g),
Err(e) => {
eprintln!("skipped: {e}");
return;
}
};
eprintln!("GPU: {}", gpu.adapter_name());
let model = ImageModel::load(&w).unwrap();
let pixels = r.get("input.pixel_values").unwrap();
let cpu = model.encode(pixels).unwrap();
let enc = GpuEncoder::load(gpu, &w, model.vision.hiera.clone()).unwrap();
let t = std::time::Instant::now();
let got = enc.encode(pixels).unwrap();
eprintln!("GPU encode (first run, includes shader compile): {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
let t = std::time::Instant::now();
let got = (enc.encode(pixels).unwrap(), got).0;
eprintln!("GPU encode: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
let low_cpu = cpu.features.features.get(2).unwrap();
let e = max_abs(&got.low, low_cpu);
eprintln!("low level vs CPU: max abs error {e:e}");
assert!(e < GPU_TOL, "low level error {e}");
let e = max_abs(&got.low, r.get("vision_encoder.neck.1").unwrap());
eprintln!("low level vs PyTorch: max abs error {e:e}");
assert!(e < GPU_TOL, "low level vs PyTorch error {e}");
let (e0, e1) = (max_abs(&got.high.0, &cpu.high.0), max_abs(&got.high.1, &cpu.high.1));
eprintln!("skip features vs CPU: {e0:e} / {e1:e}");
assert!(e0 < GPU_TOL && e1 < GPU_TOL, "skip feature errors {e0} / {e1}");
}