ferrotensor_backend/engine/
matmul.rs

1//! Matrix multiplication kernels.
2
3use super::{uniform, ComputeEngine};
4use crate::{UnifiedTensor, WgpuContext};
5
6/// Operand layout understood by the `matmul_blocked` kernel.
7#[derive(Clone, Copy)]
8#[repr(u32)]
9enum Layout {
10    /// `a @ b` with `a: [M, K]`, `b: [K, N]`.
11    Plain = 0,
12    /// `a @ b^T` with `a: [M, K]`, `b: [N, K]`.
13    TransposedB = 1,
14    /// `a^T @ b` with `a: [K, M]`, `b: [K, N]`.
15    TransposedA = 2,
16}
17
18impl ComputeEngine {
19    /// `out[M, N] = a[M, K] @ b[K, N]`.
20    pub fn matmul(
21        &self,
22        ctx: &WgpuContext,
23        a: &UnifiedTensor,
24        b: &UnifiedTensor,
25        out: &UnifiedTensor,
26    ) {
27        let dims = (a.shape[0] as u32, a.shape[1] as u32, b.shape[1] as u32);
28        self.matmul_dispatch(ctx, Layout::Plain, dims, a, b, out);
29    }
30
31    /// `out[M, N] = a[M, K] @ b[N, K]^T` (e.g. the input gradient of a
32    /// linear layer).
33    pub fn matmul_bt(
34        &self,
35        ctx: &WgpuContext,
36        a: &UnifiedTensor,
37        b: &UnifiedTensor,
38        out: &UnifiedTensor,
39    ) {
40        let dims = (a.shape[0] as u32, a.shape[1] as u32, b.shape[0] as u32);
41        self.matmul_dispatch(ctx, Layout::TransposedB, dims, a, b, out);
42    }
43
44    /// `out[M, N] = a[K, M]^T @ b[K, N]` (e.g. the weight gradient of a
45    /// linear layer).
46    pub fn matmul_at(
47        &self,
48        ctx: &WgpuContext,
49        a: &UnifiedTensor,
50        b: &UnifiedTensor,
51        out: &UnifiedTensor,
52    ) {
53        let dims = (a.shape[1] as u32, a.shape[0] as u32, b.shape[1] as u32);
54        self.matmul_dispatch(ctx, Layout::TransposedA, dims, a, b, out);
55    }
56
57    fn matmul_dispatch(
58        &self,
59        ctx: &WgpuContext,
60        layout: Layout,
61        (m, k, n): (u32, u32, u32),
62        a: &UnifiedTensor,
63        b: &UnifiedTensor,
64        out: &UnifiedTensor,
65    ) {
66        let params = uniform(ctx, bytemuck::cast_slice(&[m, k, n, layout as u32]));
67        // One 16x16 workgroup per 64x64 output tile.
68        self.run(
69            ctx,
70            &self.matmul,
71            &[
72                (0, a.buffer()),
73                (1, b.buffer()),
74                (2, out.buffer()),
75                (3, &params),
76            ],
77            (n.div_ceil(64), m.div_ceil(64), 1),
78        );
79    }
80}
ferrotensor_backend/engine/matmul.rs

ferrotensor_backend/engine/
matmul.rs
