diff --git a/userspace/npu-probe/Makefile b/userspace/npu-probe/Makefile index 30c34e3..0512838 100644 --- a/userspace/npu-probe/Makefile +++ b/userspace/npu-probe/Makefile @@ -3,7 +3,7 @@ CC ?= gcc CFLAGS ?= -Wall -Wextra -O2 -EMIT_OBJS = rkt_matmul_cna.o rkt_matmul_coredpu.o rkt_matmul.o +EMIT_OBJS = rkt_matmul_cna.o rkt_matmul_coredpu.o rkt_matmul.o rkt_gemm.o .PHONY: all test lib clean all: selftest diff --git a/userspace/npu-probe/rkt_gemm.c b/userspace/npu-probe/rkt_gemm.c new file mode 100644 index 0000000..6da9db4 --- /dev/null +++ b/userspace/npu-probe/rkt_gemm.c @@ -0,0 +1,36 @@ +/* SPDX-License-Identifier: MIT + * + * rkt_gemm.c — U6 GEMM tiler (see rkt_gemm.h). + */ +#include "rkt_gemm.h" +#include "rkt_matmul.h" + +int rkt_gemm_plan(uint32_t M, uint32_t N, uint32_t tile_m, uint32_t tile_n, + struct rkt_gemm_tile *tiles, int max_tiles) +{ + if (M == 0 || N == 0 || tile_m == 0 || tile_n == 0) + return -1; + + int t = 0; + for (uint32_t r = 0; r < M; r += tile_m) { + for (uint32_t c = 0; c < N; c += tile_n) { + if (t >= max_tiles) + return -1; + tiles[t].row = r; + tiles[t].col = c; + tiles[t].m = (r + tile_m <= M) ? tile_m : (M - r); + tiles[t].n = (c + tile_n <= N) ? tile_n : (N - c); + t++; + } + } + return t; +} + +int rkt_gemm_op_fits(uint32_t m, uint32_t n, uint32_t k) +{ + uint64_t scratch[512]; + /* Addresses/zero-points are placeholders — fit is shape-only. */ + return rkt_build_matmul_regcmd(scratch, 512, m, n, k, + 0x1000, 0x2000, 0x3000, + 128, 128, 128) >= 0; +} diff --git a/userspace/npu-probe/rkt_gemm.h b/userspace/npu-probe/rkt_gemm.h new file mode 100644 index 0000000..dd4bccc --- /dev/null +++ b/userspace/npu-probe/rkt_gemm.h @@ -0,0 +1,36 @@ +/* SPDX-License-Identifier: MIT + * + * rkt_gemm.h — U6 backend: tile a large INT8 GEMM into rocket-op-sized + * sub-matmuls, each buildable by rkt_build_matmul_regcmd. + * + * A single rocket op has a fixed CBUF budget (12 banks), so a gemma-scale + * GEMM must be blocked. This layer tiles the M (rows/tokens) and N (output + * channels) dimensions. K (contraction / input channels) is NOT tiled here: + * the full K must fit one op — K-splitting needs partial-sum accumulation via + * the DPU add path and is deferred (see rkt_matmul_coredpu add_tensor path). + */ +#ifndef RKT_GEMM_H +#define RKT_GEMM_H + +#include + +struct rkt_gemm_tile { + uint32_t row; /* first row in M this tile covers */ + uint32_t col; /* first col in N this tile covers */ + uint32_t m; /* number of rows */ + uint32_t n; /* number of cols */ +}; + +/* + * Partition an M x N output (full K) into tiles of at most tile_m x tile_n, + * row-major. Writes tiles[] and returns the tile count, or -1 on bad args or + * if more than max_tiles would be produced. Pure geometry — call + * rkt_gemm_op_fits() to confirm a tile actually fits one rocket op. + */ +int rkt_gemm_plan(uint32_t M, uint32_t N, uint32_t tile_m, uint32_t tile_n, + struct rkt_gemm_tile *tiles, int max_tiles); + +/* True (1) if an m x n, full-k matmul fits a single rocket op. */ +int rkt_gemm_op_fits(uint32_t m, uint32_t n, uint32_t k); + +#endif /* RKT_GEMM_H */ diff --git a/userspace/npu-probe/selftest.c b/userspace/npu-probe/selftest.c index 5b49da9..a2afabf 100644 --- a/userspace/npu-probe/selftest.c +++ b/userspace/npu-probe/selftest.c @@ -1,14 +1,13 @@ /* SPDX-License-Identifier: MIT * - * selftest.c — host-side structural gate for the rocket regcmd builder. + * selftest.c — host-side structural gate for the rocket regcmd builder + the + * U6 GEMM tiler. * - * NOT a golden-byte check (that needs a Mesa build or the NPU). It verifies - * that the unified matmul builder (rkt_build_matmul_regcmd) and the raw - * stages assemble into one buffer without overflow, are deterministic (the - * builders are pure functions of their inputs), report capacity exhaustion - * instead of overrunning, and emit a sane per-block target histogram. This is - * the compile+link+run smoke gate — run after every change to the emit/tiling - * code. + * NOT a golden-byte check (that is verify/, which diffs against Mesa). This + * verifies the builders assemble without overflow, are deterministic, report + * capacity exhaustion, emit a sane per-block target histogram, and that the + * GEMM tiler covers an output exactly with op-fitting tiles. Run after every + * change to the emit / tiling / planner code. */ #include #include @@ -17,12 +16,12 @@ #include "rkt_matmul.h" #include "rkt_matmul_cna.h" #include "rkt_matmul_coredpu.h" +#include "rkt_gemm.h" #define CAP 512 -/* Print the per-block target histogram and flag any stage that emitted - * nothing. target = block_enum+1; RAW64 PC operation-descriptor words fall - * through to raw/other and are expected. Returns fail count. */ +/* Per-block target histogram; flags any stage that emitted nothing. RAW64 PC + * operation-descriptor words fall through to raw/other and are expected. */ static int histogram(const char *tag, const uint64_t *buf, int n) { int pc = 0, cna = 0, core = 0, dpu = 0, rdma = 0, raw = 0; @@ -46,8 +45,7 @@ static int histogram(const char *tag, const uint64_t *buf, int n) return 0; } -/* Raw-stage path: drive rkt_emit_cna + rkt_emit_coredpu with representative - * (not tiling-accurate) params — exercises the emit/packing layer directly. */ +/* Raw-stage path: representative (not tiling-accurate) params. */ static int build_stages(uint64_t *buf, int cap) { struct cna_params c; @@ -85,7 +83,6 @@ static int build_stages(uint64_t *buf, int cap) return n1 + n2; } -/* Unified builder: derives every field from (M,N,K) via the Mesa tiling math. */ static int build_matmul(uint64_t *buf, int cap) { return rkt_build_matmul_regcmd(buf, cap, 4, 8, 16, @@ -121,9 +118,7 @@ int main(void) printf(" capacity exhaustion reported (-1): OK\n"); } - /* A shape that needs CBUF task-splitting must be refused, not mis-emitted. */ - uint64_t big[CAP]; - if (rkt_build_matmul_regcmd(big, CAP, 100000, 8, 16, + if (rkt_build_matmul_regcmd(a, CAP, 100000, 8, 16, 0x1000, 0x2000, 0x3000, 128, 128, 128) != -1) { printf("FAIL: oversized M not refused (split path unimplemented)\n"); fails++; @@ -140,6 +135,46 @@ int main(void) fails += histogram("stages", a, sn); } + printf("== U6 GEMM tiler: rkt_gemm_plan / rkt_gemm_op_fits ==\n"); + { + struct rkt_gemm_tile tiles[512]; + const uint32_t M = 64, N = 200, K = 64; + const uint32_t tm = 32, tn = 64; + int nt = rkt_gemm_plan(M, N, tm, tn, tiles, 512); + if (nt <= 0) { + printf("FAIL: rkt_gemm_plan returned %d\n", nt); + fails++; + } else { + uint64_t covered = 0; + int all_fit = 1; + for (int i = 0; i < nt; i++) { + covered += (uint64_t)tiles[i].m * tiles[i].n; + if (!rkt_gemm_op_fits(tiles[i].m, tiles[i].n, K)) + all_fit = 0; + } + if (covered != (uint64_t)M * N) { + printf("FAIL: tiles cover %llu of %u output elems\n", + (unsigned long long)covered, M * N); + fails++; + } else { + printf(" exact coverage of %ux%u by %d tiles: OK\n", + M, N, nt); + } + if (!all_fit) { + printf("FAIL: a tile does not fit one rocket op\n"); + fails++; + } else { + printf(" every tile fits one op (K=%u): OK\n", K); + } + } + if (rkt_gemm_op_fits(8, 8, 1000000)) { + printf("FAIL: impossible-K op reported as fitting\n"); + fails++; + } else { + printf(" oversized-K op correctly reported unfit: OK\n"); + } + } + if (fails) printf("\n%d CHECK(S) FAILED\n", fails); else