ferric-tensor 0.0.1

//! Ferric L2 — a general N-dimensional tensor runtime on the GPU fabric.
//!
//! This is the substrate the whole ecosystem is meant to stand on: not fixed-shape, hand-fused
//! kernels for one architecture, but a real tensor with **arbitrary rank, strided views, and
//! broadcasting**, plus **general** elementwise ops, **general** reductions over any axes, and
//! **batched** matmul. The transformer kernels in `ferric-core` become fused fast-paths of this.
//!
//! Design: eager execution, tensors are `Arc`-shared f32 buffers described by (shape, strides,
//! offset). Views (`reshape`/`permute`/`transpose`/`broadcast_to`) are zero-copy stride tricks;
//! `contiguous()` materializes. One general strided kernel powers elementwise + broadcasting; a
//! segmented kernel powers reductions; a batched kernel powers matmul. Validated against a strided
//! CPU reference on general shapes (broadcasting, non-contiguous inputs, arbitrary reduction axes).
//!
//! Next fabric layers (in progress): dtypes (f16/bf16/int), autograd tape for training, op fusion,
//! and the heterogeneous scheduler.

use ferric_core::Context;
use std::sync::Arc;
use wgpu::util::DeviceExt;

pub mod autograd; // reverse-mode autodiff (training)
pub mod cpu; // strided CPU reference (validation source of truth)
pub use autograd::Var;

/// A general N-D f32 tensor: an Arc-shared device buffer viewed through (shape, strides, offset).
#[derive(Clone)]
pub struct Tensor {
    ctx: Arc<Context>,
    buf: Arc<wgpu::Buffer>,
    pub shape: Vec<usize>,
    pub strides: Vec<usize>, // element strides (row-major by default)
    offset: usize,
}

fn contig_strides(shape: &[usize]) -> Vec<usize> {
    let mut s = vec![1usize; shape.len()];
    for i in (0..shape.len().saturating_sub(1)).rev() {
        s[i] = s[i + 1] * shape[i + 1];
    }
    s
}
fn numel(shape: &[usize]) -> usize { shape.iter().product() }

/// Broadcast two shapes NumPy-style (right-aligned). Returns the result shape.
fn broadcast_shapes(a: &[usize], b: &[usize]) -> Vec<usize> {
    let r = a.len().max(b.len());
    let mut out = vec![0usize; r];
    for i in 0..r {
        let da = if i + a.len() >= r { a[i + a.len() - r] } else { 1 };
        let db = if i + b.len() >= r { b[i + b.len() - r] } else { 1 };
        assert!(da == db || da == 1 || db == 1, "shapes {a:?} and {b:?} not broadcastable at dim {i}");
        out[i] = da.max(db);
    }
    out
}

impl Tensor {
    pub fn numel(&self) -> usize { numel(&self.shape) }
    pub fn rank(&self) -> usize { self.shape.len() }
    pub fn is_contiguous(&self) -> bool { self.strides == contig_strides(&self.shape) && self.offset == 0 }

    // ---- construction / io ----
    pub fn from_vec(ctx: &Arc<Context>, data: &[f32], shape: &[usize]) -> Tensor {
        assert_eq!(data.len(), numel(shape), "data len != shape product");
        let buf = ctx.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("tensor"),
            contents: bytemuck::cast_slice(data),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
        });
        Tensor { ctx: ctx.clone(), buf: Arc::new(buf), shape: shape.to_vec(), strides: contig_strides(shape), offset: 0 }
    }
    pub fn zeros(ctx: &Arc<Context>, shape: &[usize]) -> Tensor { Self::from_vec(ctx, &vec![0.0; numel(shape)], shape) }

    /// Materialize (contiguous) and read back to host in logical row-major order.
    pub async fn to_vec(&self) -> Vec<f32> {
        let c = self.contiguous();
        readback(&c.ctx, &c.buf, c.numel()).await
    }

    // ---- zero-copy views ----
    pub fn reshape(&self, shape: &[usize]) -> Tensor {
        assert_eq!(numel(shape), self.numel(), "reshape changes numel");
        let c = self.contiguous();
        Tensor { ctx: c.ctx, buf: c.buf, strides: contig_strides(shape), shape: shape.to_vec(), offset: 0 }
    }
    pub fn permute(&self, perm: &[usize]) -> Tensor {
        assert_eq!(perm.len(), self.rank(), "permute rank mismatch");
        Tensor {
            ctx: self.ctx.clone(), buf: self.buf.clone(), offset: self.offset,
            shape: perm.iter().map(|&p| self.shape[p]).collect(),
            strides: perm.iter().map(|&p| self.strides[p]).collect(),
        }
    }
    pub fn transpose(&self, a: usize, b: usize) -> Tensor {
        let mut p: Vec<usize> = (0..self.rank()).collect();
        p.swap(a, b);
        self.permute(&p)
    }
    /// Broadcast to a larger shape (right-aligned); broadcast dims get stride 0.
    pub fn broadcast_to(&self, shape: &[usize]) -> Tensor {
        let r = shape.len();
        assert!(r >= self.rank(), "cannot broadcast to fewer dims");
        let mut strides = vec![0usize; r];
        for i in 0..self.rank() {
            let (si, di) = (self.rank() - 1 - i, r - 1 - i);
            if self.shape[si] == shape[di] {
                strides[di] = self.strides[si];
            } else {
                assert_eq!(self.shape[si], 1, "cannot broadcast dim {} of {:?} to {:?}", si, self.shape, shape);
                strides[di] = 0;
            }
        }
        Tensor { ctx: self.ctx.clone(), buf: self.buf.clone(), shape: shape.to_vec(), strides, offset: self.offset }
    }

    /// Materialize a (possibly strided/broadcast) view into a fresh contiguous buffer.
    pub fn contiguous(&self) -> Tensor {
        if self.is_contiguous() {
            return self.clone();
        }
        let n = self.numel();
        let out = empty(&self.ctx, n);
        // info: [rank, n, offset, shape..., strides...]
        let mut info = vec![self.rank() as u32, n as u32, self.offset as u32];
        info.extend(self.shape.iter().map(|&x| x as u32));
        info.extend(self.strides.iter().map(|&x| x as u32));
        run(&self.ctx, GATHER_WGSL, "gather", &[&self.buf, &out, &u32buf(&self.ctx, &info)], groups(n));
        Tensor { ctx: self.ctx.clone(), buf: Arc::new(out), shape: self.shape.clone(), strides: contig_strides(&self.shape), offset: 0 }
    }

    // ---- general broadcasting elementwise ----
    fn binary(&self, other: &Tensor, op: u32) -> Tensor {
        let shape = broadcast_shapes(&self.shape, &other.shape);
        let a = self.broadcast_to(&shape);
        let b = other.broadcast_to(&shape);
        let n = numel(&shape);
        let out = empty(&self.ctx, n);
        // info: [rank, op, n, offA, offB, shape..., aStr..., bStr...]
        let mut info = vec![shape.len() as u32, op, n as u32, a.offset as u32, b.offset as u32];
        info.extend(shape.iter().map(|&x| x as u32));
        info.extend(a.strides.iter().map(|&x| x as u32));
        info.extend(b.strides.iter().map(|&x| x as u32));
        run(&self.ctx, BINARY_WGSL, "binary", &[&a.buf, &b.buf, &out, &u32buf(&self.ctx, &info)], groups(n));
        Tensor { ctx: self.ctx.clone(), buf: Arc::new(out), shape: shape.clone(), strides: contig_strides(&shape), offset: 0 }
    }
    pub fn add(&self, o: &Tensor) -> Tensor { self.binary(o, 0) }
    pub fn sub(&self, o: &Tensor) -> Tensor { self.binary(o, 1) }
    pub fn mul(&self, o: &Tensor) -> Tensor { self.binary(o, 2) }
    pub fn div(&self, o: &Tensor) -> Tensor { self.binary(o, 3) }
    pub fn maximum(&self, o: &Tensor) -> Tensor { self.binary(o, 4) }

    fn unary(&self, op: u32) -> Tensor {
        let c = self.contiguous();
        let n = c.numel();
        let out = empty(&self.ctx, n);
        run(&self.ctx, UNARY_WGSL, "unary", &[&c.buf, &out, &u32buf(&self.ctx, &[op, n as u32])], groups(n));
        Tensor { ctx: self.ctx.clone(), buf: Arc::new(out), shape: c.shape, strides: c.strides, offset: 0 }
    }
    pub fn exp(&self) -> Tensor { self.unary(0) }
    pub fn neg(&self) -> Tensor { self.unary(1) }
    pub fn relu(&self) -> Tensor { self.unary(2) }
    pub fn sqrt(&self) -> Tensor { self.unary(3) }
    pub fn relu_mask(&self) -> Tensor { self.unary(4) } // 1 where x>0 else 0 (relu' )
    pub(crate) fn ctx_arc(&self) -> Arc<Context> { self.ctx.clone() }

    // ---- general reduction over arbitrary axes ----
    fn reduce(&self, axes: &[usize], op: u32, keepdim: bool) -> Tensor {
        let mut ax: Vec<usize> = axes.to_vec();
        ax.sort_unstable();
        ax.dedup();
        let keep: Vec<usize> = (0..self.rank()).filter(|d| !ax.contains(d)).collect();
        // permute reduced axes to the end, materialize → [outer, red]
        let perm: Vec<usize> = keep.iter().chain(ax.iter()).copied().collect();
        let moved = self.permute(&perm).contiguous();
        let red: usize = ax.iter().map(|&d| self.shape[d]).product();
        let outer: usize = moved.numel() / red.max(1);
        let out = empty(&self.ctx, outer);
        run(&self.ctx, REDUCE_WGSL, "reduce", &[&moved.buf, &out, &u32buf(&self.ctx, &[outer as u32, red as u32, op])], groups(outer));
        let mut oshape: Vec<usize> = keep.iter().map(|&d| self.shape[d]).collect();
        if keepdim {
            oshape = (0..self.rank()).map(|d| if ax.contains(&d) { 1 } else { self.shape[d] }).collect();
        }
        if oshape.is_empty() { oshape.push(1); }
        Tensor { ctx: self.ctx.clone(), buf: Arc::new(out), strides: contig_strides(&oshape), shape: oshape, offset: 0 }
    }
    pub fn sum(&self, axes: &[usize], keepdim: bool) -> Tensor { self.reduce(axes, 0, keepdim) }
    pub fn max(&self, axes: &[usize], keepdim: bool) -> Tensor { self.reduce(axes, 1, keepdim) }
    pub fn mean(&self, axes: &[usize], keepdim: bool) -> Tensor {
        let n: usize = axes.iter().map(|&d| self.shape[d]).product();
        let s = self.sum(axes, keepdim);
        let inv = Tensor::from_vec(&self.ctx, &[1.0 / n as f32], &[1]);
        s.mul(&inv)
    }

    // ---- batched matmul: [..., m, k] x [..., k, n] -> [..., m, n], batch dims broadcast ----
    pub fn matmul(&self, other: &Tensor) -> Tensor {
        let (ra, rb) = (self.rank(), other.rank());
        assert!(ra >= 2 && rb >= 2, "matmul needs rank >= 2");
        let (m, ka) = (self.shape[ra - 2], self.shape[ra - 1]);
        let (kb, n) = (other.shape[rb - 2], other.shape[rb - 1]);
        assert_eq!(ka, kb, "matmul inner dims {ka} != {kb}");
        let batch_a = &self.shape[..ra - 2];
        let batch_b = &other.shape[..rb - 2];
        let batch = broadcast_shapes(batch_a, batch_b);
        let bn: usize = numel(&batch);
        let a_full: Vec<usize> = batch.iter().chain([m, ka].iter()).copied().collect();
        let b_full: Vec<usize> = batch.iter().chain([kb, n].iter()).copied().collect();
        let a = self.broadcast_to(&a_full).contiguous();
        let b = other.broadcast_to(&b_full).contiguous();
        let out = empty(&self.ctx, bn * m * n);
        run(&self.ctx, MATMUL_WGSL, "bmm", &[&a.buf, &b.buf, &out, &u32buf(&self.ctx, &[bn as u32, m as u32, ka as u32, n as u32])], groups(bn * m * n));
        let oshape: Vec<usize> = batch.iter().chain([m, n].iter()).copied().collect();
        Tensor { ctx: self.ctx.clone(), buf: Arc::new(out), strides: contig_strides(&oshape), shape: oshape, offset: 0 }
    }
}

// ---------- device plumbing (uses ferric-core Context's public device/queue) ----------
fn empty(ctx: &Context, n: usize) -> wgpu::Buffer {
    ctx.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("t"), size: (n.max(1) * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}
fn u32buf(ctx: &Context, data: &[u32]) -> wgpu::Buffer {
    ctx.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("info"), contents: bytemuck::cast_slice(data),
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
    })
}
fn groups(n: usize) -> (u32, u32, u32) { (((n as u32) + 63) / 64, 1, 1) }
fn run(ctx: &Context, wgsl: &str, label: &str, binds: &[&wgpu::Buffer], g: (u32, u32, u32)) {
    let module = ctx.device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some(label), source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(wgsl)),
    });
    let pipe = ctx.device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: Some(label), layout: None, module: &module, entry_point: Some("main"),
        compilation_options: Default::default(), cache: None,
    });
    let entries: Vec<wgpu::BindGroupEntry> = binds.iter().enumerate()
        .map(|(i, b)| wgpu::BindGroupEntry { binding: i as u32, resource: b.as_entire_binding() }).collect();
    let bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some(label), layout: &pipe.get_bind_group_layout(0), entries: &entries,
    });
    let mut enc = ctx.device.create_command_encoder(&Default::default());
    {
        let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some(label), timestamp_writes: None });
        pass.set_pipeline(&pipe);
        pass.set_bind_group(0, &bg, &[]);
        pass.dispatch_workgroups(g.0, g.1, g.2);
    }
    ctx.queue.submit([enc.finish()]);
}
async fn readback(ctx: &Context, buf: &wgpu::Buffer, n: usize) -> Vec<f32> {
    let bytes = (n * 4) as u64;
    let staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("staging"), size: bytes, usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, mapped_at_creation: false,
    });
    let mut enc = ctx.device.create_command_encoder(&Default::default());
    enc.copy_buffer_to_buffer(buf, 0, &staging, 0, bytes);
    ctx.queue.submit([enc.finish()]);
    let (tx, rx) = flume::bounded(1);
    staging.slice(..).map_async(wgpu::MapMode::Read, move |r| { let _ = tx.send(r); });
    let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
    rx.recv_async().await.unwrap().unwrap();
    let data = staging.slice(..).get_mapped_range().unwrap();
    let out = bytemuck::cast_slice(&data).to_vec();
    drop(data);
    staging.unmap();
    out
}

// ---------- general kernels ----------
// row-major decode of a linear output index into per-input strided offsets.
const BINARY_WGSL: &str = r#"
@group(0) @binding(0) var<storage,read>        a: array<f32>;
@group(0) @binding(1) var<storage,read>        b: array<f32>;
@group(0) @binding(2) var<storage,read_write>  out: array<f32>;
@group(0) @binding(3) var<storage,read>        info: array<u32>; // rank,op,n,offA,offB,shape[r],aStr[r],bStr[r]
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x; let rank = info[0]; let op = info[1]; let n = info[2];
    if (i >= n) { return; }
    var ia = info[3]; var ib = info[4]; var rem = i;
    for (var dd: u32 = 0u; dd < rank; dd = dd + 1u) {
        let d = rank - 1u - dd;
        let sz = info[5u + d];
        let idx = rem % sz; rem = rem / sz;
        ia = ia + idx * info[5u + rank + d];
        ib = ib + idx * info[5u + 2u * rank + d];
    }
    let x = a[ia]; let y = b[ib];
    var r: f32 = 0.0;
    switch (op) {
        case 0u: { r = x + y; }
        case 1u: { r = x - y; }
        case 2u: { r = x * y; }
        case 3u: { r = x / y; }
        case 4u: { r = max(x, y); }
        default: { r = x + y; }
    }
    out[i] = r;
}
"#;

const UNARY_WGSL: &str = r#"
@group(0) @binding(0) var<storage,read>        x: array<f32>;
@group(0) @binding(1) var<storage,read_write>  out: array<f32>;
@group(0) @binding(2) var<storage,read>        info: array<u32>; // op, n
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x; if (i >= info[1]) { return; }
    let v = x[i]; var r: f32 = v;
    switch (info[0]) {
        case 0u: { r = exp(v); }
        case 1u: { r = -v; }
        case 2u: { r = max(v, 0.0); }
        case 3u: { r = sqrt(v); }
        case 4u: { if (v > 0.0) { r = 1.0; } else { r = 0.0; } }
        default: { r = v; }
    }
    out[i] = r;
}
"#;

const GATHER_WGSL: &str = r#"
@group(0) @binding(0) var<storage,read>        x: array<f32>;
@group(0) @binding(1) var<storage,read_write>  out: array<f32>;
@group(0) @binding(2) var<storage,read>        info: array<u32>; // rank,n,offset,shape[r],strides[r]
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x; let rank = info[0]; let n = info[1];
    if (i >= n) { return; }
    var src = info[2]; var rem = i;
    for (var dd: u32 = 0u; dd < rank; dd = dd + 1u) {
        let d = rank - 1u - dd;
        let sz = info[3u + d];
        let idx = rem % sz; rem = rem / sz;
        src = src + idx * info[3u + rank + d];
    }
    out[i] = x[src];
}
"#;

const REDUCE_WGSL: &str = r#"
@group(0) @binding(0) var<storage,read>        x: array<f32>;   // [outer, red] contiguous
@group(0) @binding(1) var<storage,read_write>  out: array<f32>; // [outer]
@group(0) @binding(2) var<storage,read>        info: array<u32>; // outer, red, op(0=sum,1=max)
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x; let outer = info[0]; let red = info[1]; let op = info[2];
    if (i >= outer) { return; }
    let base = i * red;
    if (op == 1u) {
        var acc = x[base];
        for (var j: u32 = 1u; j < red; j = j + 1u) { acc = max(acc, x[base + j]); }
        out[i] = acc;
    } else {
        var acc = 0.0;
        for (var j: u32 = 0u; j < red; j = j + 1u) { acc = acc + x[base + j]; }
        out[i] = acc;
    }
}
"#;

const MATMUL_WGSL: &str = r#"
@group(0) @binding(0) var<storage,read>        a: array<f32>;   // [batch, m, k]
@group(0) @binding(1) var<storage,read>        b: array<f32>;   // [batch, k, n]
@group(0) @binding(2) var<storage,read_write>  out: array<f32>; // [batch, m, n]
@group(0) @binding(3) var<storage,read>        info: array<u32>; // batch, m, k, n
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x; let batch = info[0]; let m = info[1]; let k = info[2]; let n = info[3];
    if (idx >= batch * m * n) { return; }
    let j = idx % n; let i = (idx / n) % m; let bt = idx / (m * n);
    let ao = bt * m * k + i * k; let bo = bt * k * n;
    var acc = 0.0;
    for (var l: u32 = 0u; l < k; l = l + 1u) { acc = acc + a[ao + l] * b[bo + l * n + j]; }
    out[idx] = acc;
}
"#;