""" block_scale.py — the smallest *block* (load–compute–store) accelerator. This is the documentation anchor for the **block** custom-hook pattern (``docs/guide/custom_hooks/block.md``): an elementwise affine map ``y[i] = A * x[i] + B`` over a resident array block in shared memory. The teaching point is the division of labour between the kernel body and the hook: * ``run_proc`` owns the **data movement** — it reads the operand block off the ``m_axi`` master and writes the result block back. Both ``read_array`` / ``write_array`` calls live in the extractable kernel body, so codegen **auto-generates** them (``int32_array_utils::read_array_slice`` / ``write_array_slice`` bursts over the resident ``[0, n)`` range) — no memory movement is hand-written. * ``compute`` is a ``@synthesizable`` hook that is **pure compute** over the materialized C++ array: ``out[i] = A*x[i] + B``. It never touches the port. It is the array generalization of the scalar regmap example (``examples/regmap/simp_fun.py``): there the scalar I/O is auto-generated by the register path and ``compute(x, a, b)`` is a pure scalar hook; here the array I/O is auto-generated by the ``m_axi`` path and ``compute(x, n)`` is a pure block hook. A small command carries the operand/result addresses and the block length ``n`` (a real accelerator has to be told where its buffers are and how long they are); the buffer's compile-time bound is the ``max_n`` HwParam. The Python body of ``compute`` is the bit-exact golden; the hand-written C++ hook ``block_scale_compute_impl.cpp`` must match it. Kept deliberately tiny — the docs pull from it, in the spirit of ``examples/basic_vec/kernels.py``. """ from __future__ import annotations from dataclasses import dataclass, field from typing import ClassVar import numpy as np import numpy.typing as npt from waveflow.hw.clock import Clock from waveflow.hw.dataschema import DataArray, DataList, IntField, MemAddr from waveflow.hw.hw_component import HwComponent, HwParam from waveflow.hw.hw_testbench import HwTestbench from waveflow.hw.interface import StreamIF, StreamIFMaster, StreamIFSlave from waveflow.hw.memif import DirectMMIF, MMIFMaster from waveflow.hw.memory import MemComponent from waveflow.hw.synth import synthesizable from waveflow.simulation.simobj import ProcessGen, SimObj from waveflow.simulation.simulation import Simulation # --------------------------------------------------------------------------- # Data model + parameters # --------------------------------------------------------------------------- INCLUDE_DIR = "include" Int32 = IntField.specialize(bitwidth=32, signed=True, include_dir=INCLUDE_DIR) NField = IntField.specialize(bitwidth=32, signed=False) STREAM_BW = 32 # command stream width (bits) MEM_BW = 32 # m_axi data width (bits) MEM_AWIDTH = 64 # m_axi address width (bits) MAX_N = 256 # compile-time bound on the block length (buffer size) A = 3 # affine scale (compile-time) B = -4 # affine bias (compile-time) DEFAULT_N = 16 # block length used by the reference run MAX_MEM_WORDS = 2 * MAX_N # operand block + result block AddrField = MemAddr.specialize(bitwidth=MEM_AWIDTH) def block_affine(x: np.ndarray, a: int = A, b: int = B) -> np.ndarray: """The golden: full-precision ``y = a*x + b`` per element, as int32. Test data is kept small enough that ``a*x + b`` never overflows int32, so this is bit-identical to the C++ hook's ``ap_int<32>`` arithmetic.""" xv = np.asarray(x, dtype=np.int64) return (a * xv + b).astype(np.int32) class BlockCmd(DataList): """Command: where the operand block is, where the result goes, how long.""" elements = { "n": {"schema": NField, "description": "Block length (elements)"}, "x_addr": {"schema": AddrField, "description": "Operand block base address"}, "y_addr": {"schema": AddrField, "description": "Result block base address"}, } class BlockBuf(DataArray): """Local-buffer type used only to type the hook's array param/return in the generated C++ (``ap_int<32> x[MAX_N]`` / ``out[MAX_N]``). Carries only (element type, compile-time bound) for codegen — the SimPy model works on numpy.""" cpp_typing_only = True element_type = Int32 static = True max_shape = (MAX_N,) cpp_storage = "raw" SCHEMA_CLASSES = [BlockCmd] # --------------------------------------------------------------------------- # Accelerator (SimPy model + codegen source) # --------------------------------------------------------------------------- @dataclass class BlockScaleComponent(HwComponent): """Synthesizable block-scale kernel — the codegen source for ``gen/block_scale.cpp``. ``run_proc`` is the kernel body (stream-controlled, so the codegen root is ``run_proc``): read one :class:`BlockCmd`, read the ``n``-element operand block off ``m_mem``, run the pure ``compute`` hook, write the result block back. The two array ops auto-generate; only ``compute`` is hand-written.""" cpp_kernel_name: ClassVar[str | None] = "block_scale" cpp_namespace: ClassVar[str | None] = "block_scale_impl" in_bw: HwParam[int] = STREAM_BW mem_bw: HwParam[int] = MEM_BW mem_awidth: HwParam[int] = MEM_AWIDTH # m_axi address width (TB-side typing) max_n: HwParam[int] = MAX_N # compile-time buffer bound clk: Clock = field(default_factory=lambda: Clock(freq=1e9)) def __post_init__(self) -> None: super().__post_init__() self.s_in = StreamIFSlave(name=f"{self.name}_s_in", sim=self.sim, bitwidth=self.in_bw) self.m_mem = MMIFMaster( name=f"{self.name}_m_mem", sim=self.sim, bitwidth=self.mem_bw) for ep in (self.s_in, self.m_mem): self.add_endpoint(ep) def run_proc(self) -> ProcessGen[None]: """Kernel body (single ap_ctrl_hs invocation). The ``read_array`` / ``write_array`` calls are auto-generated by codegen into ``int32_array_utils::read_array_slice`` / ``write_array_slice`` bursts over the resident ``[0, n)`` range; ``max_count`` gives each buffer its compile-time bound. Only ``compute`` is hand-written.""" cmd = yield from self.s_in.get(BlockCmd) x = yield from self.m_mem.read_array(Int32, cmd.n, cmd.x_addr, max_count=self.max_n) y = yield from self.compute(x, cmd.n) yield from self.m_mem.write_array(y, Int32, cmd.y_addr, cmd.n, max_count=self.max_n) @synthesizable def compute(self, x: BlockBuf, n: int) -> ProcessGen[BlockBuf]: """The pure-compute hook (hand-written as ``block_scale_compute_impl.cpp``). Pure over the materialized array — no port access. The Python body is the bit-exact golden; the C++ contract it must match is the hand-written impl file.""" return block_affine(np.asarray(x)[:int(n)]) yield # unreachable — makes this a generator (ProcessGen) # --------------------------------------------------------------------------- # SimPy controller (testbench driver) # --------------------------------------------------------------------------- @dataclass(kw_only=True) class BlockController(SimObj): """Allocates the operand + result regions, writes the operand, issues the command; the result is read back after the run completes.""" mem: MemComponent x: npt.NDArray[np.int32] word_bw: int = MEM_BW def __post_init__(self) -> None: super().__post_init__() self.m_cmd = StreamIFMaster(name=f"{self.name}_m_cmd", sim=self.sim, bitwidth=STREAM_BW) self.y_addr: int | None = None def run_proc(self) -> ProcessGen[None]: from waveflow.hw.arrayutils import get_nwords n = len(self.x) x_addr = self.mem.alloc(get_nwords(Int32, word_bw=self.mem.word_size, shape=n)) self.y_addr = self.mem.alloc(get_nwords(Int32, word_bw=self.mem.word_size, shape=n)) # Write the operand, then issue the command — the accelerator waits on # the command, so the operand is resident before it reads. yield from self.mem.m_mm.write_array( np.asarray(self.x, dtype=np.int32), Int32, x_addr, word_bw=self.word_bw) yield from self.m_cmd.write(BlockCmd(n=n, x_addr=x_addr, y_addr=self.y_addr)) def connect(sim: Simulation, ctrl: BlockController, accel: BlockScaleComponent, mem: MemComponent, clk: Clock) -> None: in_stream = StreamIF(sim=sim, clk=clk) mem_link = DirectMMIF(sim=sim, clk=clk, byte_addressable=True) in_stream.bind("master", ctrl.m_cmd) in_stream.bind("slave", accel.s_in) mem_link.bind( "master", accel.m_mem) mem_link.bind( "slave", mem.s_mm) def run_sim(x: np.ndarray, *, clk_freq: float = 1e9) -> np.ndarray: """Run the SimPy model and return the kernel-produced ``y``.""" sim = Simulation() clk = Clock(freq=clk_freq) mem = MemComponent(name="mem", sim=sim, inline=False, clk=clk) accel = BlockScaleComponent(name="block_scale", sim=sim, clk=clk) ctrl = BlockController(name="ctrl", sim=sim, mem=mem, x=np.asarray(x, dtype=np.int32)) connect(sim, ctrl, accel, mem, clk) sim.run_sim() # Read the result back after the run completes (no response stream to wait on). assert ctrl.y_addr is not None return np.asarray(mem.read_array(ctrl.y_addr, Int32, count=len(x)), dtype=np.int32) # --------------------------------------------------------------------------- # Codegen-source testbench (lowered to gen/block_scale_tb.cpp) # --------------------------------------------------------------------------- @dataclass class BlockScaleTBHls(HwTestbench): """Sequential codegen-source testbench for the block-scale kernel. ``main()`` lowers to ``gen/block_scale_tb.cpp``: read the command + operand block from disk, allocate the operand + result regions (patching their addresses into the command), push the command, run the DUT against the ``mem`` pointer, and write the kernel-produced result back for the functional verify step.""" cpp_kernel_name: ClassVar[str | None] = "block_scale" def main(self) -> None: dut = BlockScaleComponent() mem = MemComponent(name="mem", sim=None, inline=False, nwords_tot=MAX_MEM_WORDS) cmd = BlockCmd() cmd.read_uint32_file(self.data_dir + "/cmd.bin") x = BlockBuf() x.read_uint32_file_array(self.data_dir + "/x_array.bin", count=cmd.n) y = BlockBuf() cmd.x_addr = mem.alloc_array(x, Int32, count=cmd.n) cmd.y_addr = mem.alloc_array(y, Int32, count=cmd.n) dut.s_in.push(cmd) dut.run(mem=mem) out = mem.read_array(cmd.y_addr, Int32, count=cmd.n) out.write_uint32_file_array(self.data_dir + "/y_array.bin", count=cmd.n) def main() -> None: rng = np.random.default_rng(7) x = rng.integers(-100, 100, size=DEFAULT_N).astype(np.int32) y = run_sim(x) gold = block_affine(x) ok = np.array_equal(y, gold) print(f"block_scale sim: n={DEFAULT_N}, A={A}, B={B}, passed={ok}") if not ok: print(f" expected={gold.tolist()}") print(f" got ={y.tolist()}") if __name__ == "__main__": main()