rhdl_ubft/card/

fabric.rs

//! The on-chip fabric of the one-card cluster: what RoCE would carry
//! between replicas, as an RHDL experiment (D2).
//!
//! # What it is, and what it is not
//!
//! Each replica sends RDMA writes as a meta beat (opcode, QPN, addresses,
//! length) and data beats, exactly as it would to the RoCE stack. The fabric
//! turns each write into AXI writes into the destination replica's HBM
//! region. It reproduces the parts of RoCE the source audit established
//! (Atlas, `docs/rhdl_ubft-roce.md`) and nothing more:
//!
//! * **The host-programmed connection table.** Replica `i`'s local QPN
//!   `(N−1)·i + r` (`r < N − 1`) reaches peer `j = r` if `r < i`, else
//!   `r + 1`. Any other QPN has no entry on replica `i`: it is not delivered
//!   and counts as `unmapped`. A QPN that has an entry is delivered to that
//!   entry's peer even when the sender meant another. Upstream's kernel
//!   indexes QPNs as `id·N + j`; with `upstream_qpns` set (the replay of
//!   upstream's sends) the fabric counts every write whose actual peer
//!   differs from that meaning as `misrouted`.
//! * **No loopback.** No QPN reaches the sender itself.
//! * **Packet boundaries.** A write takes data beats up to the first
//!   `tlast`. Beats past `len` bytes before `tlast` are not written and
//!   count as `oversize`; a `tlast` before `len` bytes ends the write short
//!   (`short_packet`); `len = 0` writes nothing (`zero_length`); a write past
//!   the destination region, or whose HBM address would wrap 64 bits,
//!   counts as `out_of_bounds` and is not written. A runtime `n` outside
//!   `2..=16` or an `id ≥ n` makes every descriptor `bad_config` (nothing
//!   is delivered).
//!   The final beat's strobe covers only the bytes left of `len`.
//! * **Physical lanes.** AXI write lanes belong to the 64-byte-aligned
//!   address, so the payload stream is re-aligned onto lines: a write at
//!   `raddr` starts `raddr mod 64` bytes into its first line, each line takes
//!   the previous beat's tail and the current beat's head ([`realign`]), and
//!   one extra line carries the last beat's tail when the shift pushes it
//!   past. Shifts are in 4-byte steps (rhdl_ubft addresses words), so a
//!   `raddr` that is not a multiple of 4 is refused (`misaligned`), and a
//!   region base must be 64-byte aligned (else `bad_config`).
//!
//! # Receive authorization
//!
//! A mapped write is then checked as a whole, before any AXI write, by
//! Atlas's [`authorize_write`]: the trusted ingress (this link's `id`), the
//! destination, the requested interval and `W` trusted [`WriteWindow`]s.
//! A rejected descriptor is drained without writing and counted once, by
//! the guard's reason: membership (`bad_config`), empty (`zero_length`),
//! outside the region or wrapping (`out_of_bounds`), or no window
//! (`unauthorized`). The windows are configuration and change only while
//! the fabric is drained.
//!
//! What the audit left unknown is a labelled assumption of this model, not
//! a claim about RoCE: the source is the fabric port the write came in on
//! (never a payload field); one write at a time per source, and the next
//! starts only when every byte of the previous one is acknowledged by HBM
//! (B responses), which keeps each source's writes in order and makes a
//! later flag write land after its payload. That B response is this
//! model's commit point, not a statement about when real RDMA data is
//! visible. A fixed `latency` (cycles before a write starts) stands in for
//! the network.

use rhdl::prelude::*;
use rhdl_fpga::core::dff::DFF;

use super::authorization::{WriteCheck, WriteWindow, authorize_write};
use crate::ports::{AxiWrite, AxiWriteResp, Axis256, Axis512};

/// Error and progress counters of one source link.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
pub struct LinkCounters {
    /// Writes delivered (all their bytes acknowledged).
    pub delivered: b32,
    /// Delivered to a different peer than the sender meant (`qpn − id·N`).
    pub misrouted: b32,
    /// QPNs with no connection entry (not delivered).
    pub unmapped: b32,
    /// `tlast` before `len` bytes.
    pub short_packet: b32,
    /// Data beats after `len` bytes and before `tlast` (not written).
    pub oversize: b32,
    /// Writes with `len = 0`.
    pub zero_length: b32,
    /// Writes reaching past the destination region (not written).
    pub out_of_bounds: b32,
    /// Opcodes other than write (not delivered).
    pub not_write: b32,
    /// Descriptors dropped because `n` or `id` is out of range.
    pub bad_config: b32,
    /// Mapped writes no trusted window allows (not delivered).
    pub unauthorized: b32,
    /// Writes whose `raddr` is not a multiple of 4 (not delivered; the
    /// re-aligner shifts in 4-byte steps, the granularity rhdl_ubft uses).
    pub misaligned: b32,
}

/// One source link's inputs.
#[derive(Clone, Copy, Debug, PartialEq, Digital)]
pub struct LinkIn<const W: usize> {
    /// Replicas.
    pub n: b5,
    /// This source.
    pub id: b5,
    /// Every replica's region base (byte address in HBM).
    pub bases: [b64; 16],
    /// Region size in bytes (each replica's).
    pub region: b64,
    /// Cycles before a write starts (the network stand-in).
    pub latency: b16,
    /// The sender indexes QPNs as upstream does (`id·N + j`): count
    /// writes that reach a different peer than that meant.
    pub upstream_qpns: bool,
    /// Trusted write windows (see the module docs).
    pub windows: [WriteWindow; W],
    /// The source's `m_axis_tx_meta`.
    pub meta: Axis256,
    /// The source's `m_axis_tx_data`.
    pub data: Axis512,
    /// HBM write responses.
    pub mem: AxiWriteResp,
}

/// One source link's outputs.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
pub struct LinkOut {
    /// `tx_meta` `tready`.
    pub meta_ready: bool,
    /// `tx_data` `tready`.
    pub data_ready: bool,
    /// HBM writes.
    pub mem: AxiWrite,
    /// Counters.
    pub counters: LinkCounters,
}

/// Link phases.
const IDLE: b3 = b3(0);
/// Waiting out the latency.
const DELAY: b3 = b3(1);
/// Moving data beats.
const DATA: b3 = b3(2);
/// Dropping data beats up to `tlast` (unmapped, not a write, out of bounds).
const DROP: b3 = b3(3);
/// Waiting for the last B responses.
const DRAIN: b3 = b3(4);
/// Emitting the last line, which holds only the tail of the last beat.
const FLUSH: b3 = b3(5);

/// A link's registers.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
#[doc(hidden)]
pub struct LinkState {
    phase: b3,
    /// Destination address of the next beat.
    addr: b64,
    /// Bytes still to write.
    left: b32,
    /// Cycles of latency left.
    wait: b16,
    /// The write's `len` was 0.
    zero: bool,
    /// AW of the current beat accepted.
    aw_done: bool,
    /// W of the current beat accepted.
    w_done: bool,
    /// Beat registered for writing.
    beat_valid: bool,
    /// Its data.
    beat_data: [b64; 8],
    /// Its strobe.
    beat_strb: b64,
    /// Its address.
    beat_addr: b64,
    /// The beat with `tlast` has been taken.
    last_seen: bool,
    /// The write's offset within a 64-byte line, in 4-byte words.
    shift: b4,
    /// The previous payload beat (its tail starts the next line).
    prev_data: [b64; 8],
    /// Its byte enables (masked to the declared length).
    prev_keep: b64,
    /// Bytes of the line in hand that are still inside the write: line `j`
    /// covers payload bytes `64j − offset ..`, so this is `offset + len −
    /// 64j` (saturating).
    rem: b32,
    /// Write responses outstanding.
    pending: b8,
    counters: LinkCounters,
}

/// One source link (see the module docs).
#[derive(Clone, Debug, Synchronous, SynchronousDQ)]
pub struct FabricLink<const W: usize> {
    state: DFF<LinkState>,
}

impl<const W: usize> Default for FabricLink<W> {
    fn default() -> Self {
        Self { state: DFF::new(LinkState::default()) }
    }
}

impl<const W: usize> SynchronousIO for FabricLink<W> {
    type I = LinkIn<W>;
    type O = LinkOut;
    type Kernel = fabric_link_kernel<W>;
}

/// The strobe of a beat with `left` bytes still to write.
#[kernel]
pub fn strobe(left: b32) -> b64 {
    if left >= b32(64) { !b64(0) } else { (b64(1) << left.resize::<64>()) - b64(1) }
}

/// One 64-byte line of a re-aligned write: the payload stream shifted up by
/// `sw` 4-byte words, from the previous beat's tail and the current beat's
/// head (`sw = 0`: the current beat as it is).
#[allow(clippy::needless_range_loop)]
#[kernel]
pub fn realign(prev: [b64; 8], cur: [b64; 8], sw: b4) -> [b64; 8] {
    // Thirty-two 32-bit words: previous beat, then current beat.
    let mut win = [b32(0); 32];
    for l in 0..8 {
        win[2 * l] = prev[l].resize();
        win[2 * l + 1] = (prev[l] >> 32).resize();
        win[16 + 2 * l] = cur[l].resize();
        win[16 + 2 * l + 1] = (cur[l] >> 32).resize();
    }
    let mut out = [b32(0); 16];
    for w in 0..16 {
        for k in 0..16 {
            if sw == b4(k as u128) {
                out[w] = win[16 - k + w];
            }
        }
    }
    let mut lanes = [b64(0); 8];
    for l in 0..8 {
        lanes[l] = out[2 * l].resize::<64>() | (out[2 * l + 1].resize::<64>() << 32);
    }
    lanes
}

/// The byte enables of a re-aligned line, shifted like [`realign`].
#[kernel]
pub fn realign_keep(prev: b64, cur: b64, sw: b4) -> b64 {
    let window: b128 = (cur.resize::<128>() << 64) | prev.resize::<128>();
    let by: b128 = b128(64) - (sw.resize::<128>() << 2);
    (window >> by).resize()
}

/// The link's logic.
#[allow(clippy::too_many_lines)]
#[kernel]
pub fn fabric_link_kernel<const W: usize>(
    _cr: ClockReset,
    i: LinkIn<W>,
    q: FabricLinkQ<W>,
) -> (LinkOut, FabricLinkD<W>) {
    let r = q.state;
    let mut s = r;
    let mut c = r.counters;
    let mut meta_ready = false;
    let mut data_ready = false;

    // ------------------------------------------------ write channel
    // One registered beat at a time: AW and W independently, B counted.
    let mem = AxiWrite {
        aw_valid: r.beat_valid && !r.aw_done,
        aw_addr: r.beat_addr,
        w_valid: r.beat_valid && !r.w_done,
        w_data: r.beat_data,
        w_strb: r.beat_strb,
        b_ready: true,
    };
    let aw_now = mem.aw_valid && i.mem.aw_ready;
    let w_now = mem.w_valid && i.mem.w_ready;
    let aw_all = r.aw_done || aw_now;
    let w_all = r.w_done || w_now;
    let sent = r.beat_valid && aw_all && w_all;
    s.aw_done = aw_all && !sent;
    s.w_done = w_all && !sent;
    if sent {
        s.beat_valid = false;
    }
    let pending_after = r.pending + (if sent { b8(1) } else { b8(0) }) - (if i.mem.b_valid { b8(1) } else { b8(0) });
    s.pending = pending_after;
    let _ = mem;

    // ------------------------------------------------ descriptors
    if r.phase == IDLE {
        meta_ready = true;
        if i.meta.valid {
            let lane0 = i.meta.data[0];
            let lane1 = i.meta.data[1];
            let lane2 = i.meta.data[2];
            let op: b3 = (lane0 & b64(7)).resize();
            let qpn: b32 = ((lane0 >> 3) & b64(0xff_ffff)).resize();
            let raddr: b64 = (lane1 >> 11) & b64(0xffff_ffff_ffff);
            let len: b32 = ((lane1 >> 59) | ((lane2 & b64(0x7ff_ffff)) << 5)).resize();
            let n: b32 = i.n.resize();
            let id: b32 = i.id.resize();
            // Host-programmed table: (N−1)·id + r reaches peer r or r + 1.
            let first = (n - b32(1)) * id;
            let mapped = qpn >= first && qpn < first + n - b32(1);
            let rank = qpn - first;
            let dest: b32 = if rank < id { rank } else { rank + b32(1) };
            // The peer upstream's id·N + j indexing meant.
            let meant = qpn - n * id;
            let mut base = b64(0);
            for k in 0..16 {
                if dest == b32(k as u128) {
                    base = i.bases[k];
                }
            }
            // The whole interval, checked before any AXI write.
            let verdict = authorize_write::<W>(
                WriteCheck {
                    members: i.n,
                    source: i.id,
                    destination: dest.resize::<5>(),
                    offset: raddr,
                    length: len,
                    base,
                    region: i.region,
                },
                i.windows,
            );
            // Region bases are 64-byte aligned (a line is one AXI beat).
            let config_ok = n >= b32(2) && n <= b32(16) && id < n && (base & b64(63)) == b64(0);
            if !config_ok {
                c.bad_config = r.counters.bad_config + b32(1);
                s.phase = DROP;
            } else if op != b3(1) {
                c.not_write = r.counters.not_write + b32(1);
                s.phase = DROP;
            } else if !mapped {
                c.unmapped = r.counters.unmapped + b32(1);
                s.phase = DROP;
            } else if (raddr & b64(3)) != b64(0) {
                c.misaligned = r.counters.misaligned + b32(1);
                s.phase = DROP;
            } else if !verdict.allowed {
                if verdict.reason == b3(1) {
                    c.bad_config = r.counters.bad_config + b32(1);
                } else if verdict.reason == b3(2) {
                    c.zero_length = r.counters.zero_length + b32(1);
                } else if verdict.reason == b3(5) {
                    c.unauthorized = r.counters.unauthorized + b32(1);
                } else {
                    c.out_of_bounds = r.counters.out_of_bounds + b32(1);
                }
                s.phase = DROP;
            } else {
                if i.upstream_qpns && meant != dest {
                    c.misrouted = r.counters.misrouted + b32(1);
                }
                // Lines are 64-byte aligned; the payload starts `offset`
                // bytes into the first one.
                let at = base + raddr;
                let offset: b32 = (at & b64(63)).resize();
                s.addr = at & !b64(63);
                s.shift = (offset >> 2).resize();
                s.prev_data = [b64(0); 8];
                s.prev_keep = b64(0);
                s.rem = offset + len;
                s.left = len;
                s.zero = false;
                s.wait = i.latency;
                s.last_seen = false;
                s.phase = DELAY;
            }
        }
    } else if r.phase == DELAY {
        if r.wait == b16(0) {
            s.phase = DATA;
        } else {
            s.wait = r.wait - b16(1);
        }
    } else if r.phase == DATA {
        // Take a beat when the beat register is free.
        // At most 254 beats await their B response, so `pending` never wraps.
        data_ready = (!r.beat_valid || sent) && r.pending < b8(254);
        if i.data.valid && data_ready {
            // Bytes past the declared length never enter a line.
            let keep_in = i.data.keep & strobe(r.left);
            let rem_after = if r.rem > b32(64) { r.rem - b32(64) } else { b32(0) };
            if r.left > b32(0) {
                s.beat_valid = true;
                s.beat_data = realign(r.prev_data, i.data.data, r.shift);
                s.beat_strb = realign_keep(r.prev_keep, keep_in, r.shift) & strobe(r.rem);
                s.beat_addr = r.addr;
                s.addr = r.addr + b64(64);
                s.left = if r.left > b32(64) { r.left - b32(64) } else { b32(0) };
                s.prev_data = i.data.data;
                s.prev_keep = keep_in;
                s.rem = rem_after;
            } else {
                c.oversize = r.counters.oversize + b32(1);
            }
            if i.data.last {
                if r.left > b32(64) {
                    c.short_packet = r.counters.short_packet + b32(1);
                }
                // A shifted write spills its last in-length beat's tail into
                // one more line. The bytes still owed to memory decide it,
                // not the input still expected: when `tlast` comes on an
                // oversize beat (`left` already 0), the tail of the beat
                // before is still in `prev_data` and `rem` counts it.
                let owed = if r.left > b32(0) { rem_after } else { r.rem };
                let spill = r.shift != b4(0) && owed > b32(0);
                s.phase = if spill { FLUSH } else { DRAIN };
            }
        }
    } else if r.phase == FLUSH {
        // The last line: the previous beat's tail only.
        if !r.beat_valid || sent {
            s.beat_valid = true;
            s.beat_data = realign(r.prev_data, [b64(0); 8], r.shift);
            s.beat_strb = realign_keep(r.prev_keep, b64(0), r.shift) & strobe(r.rem);
            s.beat_addr = r.addr;
            s.addr = r.addr + b64(64);
            s.rem = b32(0);
            s.phase = DRAIN;
        }
    } else if r.phase == DROP {
        data_ready = true;
        if i.data.valid && i.data.last {
            s.phase = IDLE;
        }
    } else {
        // DRAIN: every written byte acknowledged before the next write.
        if !s.beat_valid && pending_after == b8(0) {
            c.delivered = r.counters.delivered + b32(1);
            s.phase = IDLE;
        }
    }
    s.counters = c;
    let d = FabricLinkD::<W> { state: s };
    (
        LinkOut {
            meta_ready,
            data_ready,
            mem: AxiWrite {
                aw_valid: r.beat_valid && !r.aw_done,
                aw_addr: r.beat_addr,
                w_valid: r.beat_valid && !r.w_done,
                w_data: r.beat_data,
                w_strb: r.beat_strb,
                b_ready: true,
            },
            counters: r.counters,
        },
        d,
    )
}