mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-10 02:11:41 -06:00
#1539: match reference box padding and use a flat GPU module
This commit is contained in:
Kristian Myram
@@ -0,0 +1,471 @@
|
||||
// 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!("gpu/kernels.wgsl"), include_str!("gpu/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,
|
||||
/// Two half-precision values per 32-bit word (a weight matrix only; `len` still counts values).
|
||||
pub(crate) f16: bool,
|
||||
}
|
||||
|
||||
/// 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>>,
|
||||
turn: Mutex<()>,
|
||||
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")))
|
||||
}
|
||||
|
||||
/// Whether weight matrices are kept on the GPU in half precision (`EFFECTCRAFT_SAM2_F16=1`).
|
||||
pub fn half_weights() -> bool {
|
||||
std::env::var("EFFECTCRAFT_SAM2_F16").is_ok_and(|v| matches!(v.trim(), "1" | "on" | "true"))
|
||||
}
|
||||
|
||||
/// The process's one GPU device, once something has asked for it.
|
||||
static DEVICE: std::sync::OnceLock<std::result::Result<Arc<Gpu>, String>> = std::sync::OnceLock::new();
|
||||
|
||||
impl Gpu {
|
||||
/// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders
|
||||
/// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go,
|
||||
/// the device stays. An error (no GPU) is remembered too.
|
||||
pub fn shared() -> Result<Arc<Gpu>> {
|
||||
DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(|e| e.to_string())).clone().map_err(Error::Compute)
|
||||
}
|
||||
|
||||
/// Whether the shared device has been looked for yet, and found: `None` before anything has asked
|
||||
/// for it (this never creates it, so a caller that only wants to describe the GPU does not stall
|
||||
/// on the adapter and shader compilation), then `Some(true)` or `Some(false)` (no GPU).
|
||||
pub fn probe() -> Option<bool> {
|
||||
DEVICE.get().map(|d| d.is_ok())
|
||||
}
|
||||
|
||||
/// A device on the best real GPU (a software adapter is refused: the CPU path is faster).
|
||||
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), turn: Mutex::new(()), error, adapter: info.name })
|
||||
}
|
||||
|
||||
/// One GPU pass at a time: a caller holds this while it queues, flushes and reads back a whole job, so
|
||||
/// two threads never interleave their dispatches in one command stream.
|
||||
pub fn exclusive(&self) -> std::sync::MutexGuard<'_, ()> {
|
||||
self.turn.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
|
||||
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 {
|
||||
// `write_buffer` instead of a mapped-at-creation buffer: on a lost or out-of-memory device the latter panics inside wgpu,
|
||||
// while this records the error for `take_error`.
|
||||
let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: None,
|
||||
size: (data.len().max(1) as u64) * 4,
|
||||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
if !data.is_empty() {
|
||||
self.queue.write_buffer(&buffer, 0, bytemuck::cast_slice(data));
|
||||
}
|
||||
Buf { buffer, len: data.len(), f16: false }
|
||||
}
|
||||
|
||||
/// A weight matrix: stored as half precision (half the memory) when `EFFECTCRAFT_SAM2_F16` asks for it and every
|
||||
/// value fits, as f32 otherwise.
|
||||
pub fn upload_weights(&self, data: &[f32]) -> Buf {
|
||||
if !half_weights() || data.is_empty() || data.iter().any(|v| v.abs() > 60_000.0) {
|
||||
return self.upload(data);
|
||||
}
|
||||
self.upload_f16(data)
|
||||
}
|
||||
|
||||
/// `data` as half precision, two values per word.
|
||||
pub(crate) fn upload_f16(&self, data: &[f32]) -> Buf {
|
||||
let words: Vec<u32> = data
|
||||
.chunks(2)
|
||||
.map(|p| {
|
||||
let lo = half::f16::from_f32(p.first().copied().unwrap_or(0.0)).to_bits();
|
||||
let hi = half::f16::from_f32(p.get(1).copied().unwrap_or(0.0)).to_bits();
|
||||
u32::from(lo) | (u32::from(hi) << 16)
|
||||
})
|
||||
.collect();
|
||||
let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: None,
|
||||
size: (words.len() as u64) * 4,
|
||||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
self.queue.write_buffer(&buffer, 0, bytemuck::cast_slice(&words));
|
||||
Buf { buffer, len: data.len(), f16: true }
|
||||
}
|
||||
|
||||
/// 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, f16: false })
|
||||
}
|
||||
|
||||
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(());
|
||||
}
|
||||
if sc.col != 1 {
|
||||
return Err(Error::Input(format!("matmul output needs a column stride of 1, got {}", sc.col)));
|
||||
}
|
||||
// bounds: the last element each operand touches must lie inside its buffer
|
||||
let last = |s: Strides, rows: usize, cols: usize| {
|
||||
s.offset.checked_add((op.batch - 1).checked_mul(s.batch)?)?.checked_add((rows - 1).checked_mul(s.row)?)?.checked_add((cols - 1).checked_mul(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 name != "C" && op.k == 0 {
|
||||
continue;
|
||||
}
|
||||
let last = last(s, rows, cols).ok_or_else(|| Error::Input(format!("matmul operand {name} layout overflows its element index")))?;
|
||||
if last >= buf.len {
|
||||
return Err(Error::Input(format!("matmul operand {name} reaches element {last} of a {}-element buffer", 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(
|
||||
if b.f16 { "matmul_w16" } else { "matmul" },
|
||||
&[(0, a), (1, b), (2, bias), (3, c), (4, res)],
|
||||
¶ms,
|
||||
(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;
|
||||
pub mod memattn;
|
||||
pub mod memenc;
|
||||
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_weights(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?),
|
||||
b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?),
|
||||
inp,
|
||||
out,
|
||||
})
|
||||
}
|
||||
|
||||
/// As [`GLinear::load`], with output channel `j` scaled by `scale[j]` (the memory fuser's layer scale).
|
||||
pub fn load_scaled(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize, scale: &[f32]) -> Result<GLinear> {
|
||||
if scale.len() != out {
|
||||
return Err(Error::Weights(format!("`{name}`: {} scales for {out} outputs", scale.len())));
|
||||
}
|
||||
let mut wt = flat(w.get(&format!("{name}.weight"), &[out, inp])?)?;
|
||||
let mut bias = flat(w.get(&format!("{name}.bias"), &[out])?)?;
|
||||
for (j, s) in scale.iter().enumerate() {
|
||||
if let Some(row) = wt.get_mut(j * inp..(j + 1) * inp) {
|
||||
row.iter_mut().for_each(|v| *v *= s);
|
||||
}
|
||||
if let Some(b) = bias.get_mut(j) {
|
||||
*b *= s;
|
||||
}
|
||||
}
|
||||
Ok(GLinear { w: g.upload_weights(&wt), b: g.upload(&bias), 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_weights(&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;
|
||||
Reference in New Issue
Block a user