marfrit/rosenblatt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rosenblatt U6: GEMM tiler — block a large matmul into rocket-op-sized tiles

Browse Source 
rkt_gemm_plan() partitions an MxN (full-K) output into <=tile_m x tile_n
sub-matmuls; rkt_gemm_op_fits() confirms a shape fits one rocket op via the
verified builder. First slice of the CPU+NPU backend: lets a gemma-scale GEMM
be expressed as many NPU ops. K-tiling (DPU partial-sum accumulation) deferred.

selftest: exact coverage of 64x200 by 8 tiles, every tile fits one op (K=64),
impossible-K reported unfit. Golden match unaffected. ALL PASS.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
claude-noether
2026-07-13 21:20:46 +00:00
parent d783dd1b28
commit 51ca5a34ba
4 changed files with 125 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
userspace/npu-probe/Makefile
+1 -1
View File
@@ -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
userspace/npu-probe/rkt_gemm.c
+36
View File
@@ -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;
}
userspace/npu-probe/rkt_gemm.h
+36
View File
@@ -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 <stdint.h>
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 */
userspace/npu-probe/selftest.c
+52 -17
View File
@@ -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 <stdint.h>
#include <stdio.h>
@@ -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