Skip to content

GPU bangs chain their rounds on the device, three command buffers in flight #1447

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nicolas-abril merged 5 commits into from
Oct 19, 2026
Merged

GPU bangs chain their rounds on the device, three command buffers in flight #1447

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c684b0b
A bang chains its rounds on the device: the pack takes the next front…
nicolas-abril Oct 9, 2026
e686860
Drop the added comments
nicolas-abril Oct 9, 2026
1e81d7a
A GPU bang sizes its submissions by time, and the CPU path no longer …
nicolas-abril Oct 13, 2026
69cf38e
A GPU bang runs one round per command buffer, two in flight, chained …
nicolas-abril Oct 14, 2026
1ac3299
GPU bang: Metal keeps three rounds in flight until the frontier empti…
nicolas-abril Oct 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 94 additions & 40 deletions bend2/comp.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3553,6 +3553,8 @@ enum { ERR_RING = 1, ERR_TAGS, ERR_HEAP, ERR_FIDS, ERR_NATS, ERR_RFCS,
#define H_BUMP 0
#define H_CAP 1
#define H_CURSOR LINE
#define H_FRONT (H_CURSOR + 1)
#define H_GROUPS (H_CURSOR + 2)
#define H_ROOT_DONE (2 * LINE)
#define H_ERROR_CODE (3 * LINE)
#define H_ROOT_WORD (4 * LINE)
Expand Down Expand Up @@ -4273,6 +4275,25 @@ INLINE u32 ring_flip(u32 i) {
return (i % CUBE_T << CUBE_LOG) + i / CUBE_T;
}

// A turn's five dispatches (x, y, z) sit at H_GROUPS, before H_ROOT_DONE.
typedef char turn_fits[H_GROUPS + 8 <= H_ROOT_DONE ? 1 : -1];

#define TURN_ROWS(at) ((at) % 4 == 0 ? 1 : CUBE_G)

INLINE u32 turn_groups(u64 f, u32 at) {
bool run = f != 0 && (at == 2 ? f >= CUBE_T : at == 1 ? f < LANES
: at == 4 || f < CUBE_T);
return run ? TURN_ROWS(at) : 0;
}

INLINE void turn_set(DEV u64* H, u64 f) {
DEV u32* g = (DEV u32*)(H + H_GROUPS);
H[H_FRONT] = f;
for (u32 i = 0; i < 15; i += 1) {
g[i] = i % 3 != 0 ? 1 : turn_groups(f, i / 3);
}
}

#define ring_pick(b, s, c) ((b) + (s) * (a32_add(c, 1) & (CUBE_T - 1)))

// Task
Expand Down Expand Up @@ -4580,14 +4601,23 @@ kernel void bend_dev(DEV u64* H [[buffer(0)]], constant u32& pass [[buffer(1)]],
u32 row [[threadgroup_position_in_grid]],
u32 lane [[thread_position_in_threadgroup]]) {
#else
extern "C" __global__ void bend_dev(DEV u64* H, u32 pass) {
extern "C" __global__ void bend_dev(DEV u64* H, u32 pass, u32 at) {
__shared__ u32 vote[3];
u32 grids = gridDim.x;
u32 row = blockIdx.x;
u32 lane = threadIdx.x;
#endif
#ifndef __METAL_VERSION__
if (!turn_groups(H[H_FRONT], at)) {
return;
}
#endif
if (pass == 2) {
bank_pack(H, lane);
if (lane == 0) {
turn_set(H, root_done(H) || err_seen(H) ? 0
: a32_exch(a32_at(H, H_CURSOR), 0));
}
return;
}
u32 stride = grids == 1 ? CUBE_G : 1;
Expand Down Expand Up @@ -4897,17 +4927,13 @@ static void gpu_note(const char* path) {

#if BEND_METAL || BEND_CUDA

static void gpu_kernel(u32 pass, u32 groups);
static void gpu_kernel(u32 pass, u32 at);

static void gpu_run(u32 f) {
if (f < CUBE_T) {
gpu_kernel(0, 1);
}
if (f < LANES) {
gpu_kernel(0, CUBE_G);
// Dispatch at runs pass {0, 0, 1, 3, 2}[at] on TURN_ROWS(at) groups.
static void gpu_run(void) {
for (u32 i = 0; i < 5; i += 1) {
gpu_kernel((u32[]){ 0, 0, 1, 3, 2 }[i], i);
}
gpu_kernel(f < CUBE_T ? 3 : 1, CUBE_G);
gpu_kernel(2, 1);
}

#endif
Expand Down Expand Up @@ -4996,26 +5022,46 @@ static void gpu_load(u64 bytes) {
}
}

static void gpu_kernel(u32 pass, u32 groups) {
[gpu_enc setComputePipelineState:gpu_pso];
[gpu_enc setBuffer:gpu_buf offset:0 atIndex:0];
static void gpu_kernel(u32 pass, u32 at) {
[gpu_enc setBytes:&pass length:sizeof pass atIndex:1];
[gpu_enc setThreadgroupMemoryLength:TG_HOLD * 8 atIndex:0];
[gpu_enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
[gpu_enc dispatchThreadgroupsWithIndirectBuffer:gpu_buf
indirectBufferOffset:H_GROUPS * 8 + at * 12
threadsPerThreadgroup:MTLSizeMake(CUBE_T, 1, 1)];
[gpu_enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
}

static void gpu_pass(u32 f) {
@autoreleasepool {
id<MTLCommandBuffer> cb = [gpu_que commandBuffer];
gpu_enc = [cb computeCommandEncoder];
gpu_run(f);
[gpu_enc endEncoding];
[cb commit];
[cb waitUntilCompleted];
if ([cb error]) {
gpu_fail([cb error]);
// A bang gives each round its own command buffer, GPU_DEEP in flight, each
// waiting on the last through one event, until the pack empties the
// frontier. A round queued past that dispatches no groups.

#define GPU_DEEP 3

static id<MTLEvent> gpu_ev;
static u64 gpu_ev_at;

static void gpu_pass(u64* H) {
id<MTLCommandBuffer> cb[GPU_DEEP] = { nil };
gpu_ev = gpu_ev ? gpu_ev : [gpu_dev newEvent];
for (u32 i = 0, left = GPU_DEEP; left != 0; i += 1) {
@autoreleasepool {
id<MTLCommandBuffer> c = cb[i % GPU_DEEP];
[c waitUntilCompleted];
if ([c error]) {
gpu_fail([c error]);
}
left -= left < GPU_DEEP || ((volatile u64*)H)[H_FRONT] == 0;
cb[i % GPU_DEEP] = c = left == GPU_DEEP ? [gpu_que commandBuffer] : nil;
if (c) {
[c encodeWaitForEvent:gpu_ev value:gpu_ev_at];
gpu_enc = [c computeCommandEncoder];
[gpu_enc setComputePipelineState:gpu_pso];
[gpu_enc setBuffer:gpu_buf offset:0 atIndex:0];
[gpu_enc setThreadgroupMemoryLength:TG_HOLD * 8 atIndex:0];
gpu_run();
[gpu_enc endEncoding];
[c encodeSignalEvent:gpu_ev value:++gpu_ev_at];
[c commit];
}
}
}
}
Expand Down Expand Up @@ -5154,18 +5200,24 @@ static void gpu_load(u64 bytes) {
}
}

static void gpu_kernel(u32 pass, u32 groups) {
void* args[] = { &CORPUS, &pass };
if (cuLaunchKernel(gpu_pso, groups, 1, 1, CUBE_T, 1, 1, 0, NULL,
static void gpu_kernel(u32 pass, u32 at) {
void* args[] = { &CORPUS, &pass, &at };
if (cuLaunchKernel(gpu_pso, TURN_ROWS(at), 1, 1, CUBE_T, 1, 1, 0, NULL,
args, NULL) != CUDA_SUCCESS) {
err_fail("device launch failed");
}
}

static void gpu_pass(u32 f) {
gpu_run(f);
if (cuCtxSynchronize() != CUDA_SUCCESS) {
err_fail("device fault");
// A bang launches 1, 2, 4 .. 64 rounds between syncs, until the pack
// empties the frontier. A pass launched past that returns at once.
static void gpu_pass(u64* H) {
for (u32 n = 1; H[H_FRONT] != 0; n = n < 32 ? 2 * n : 64) {
for (u32 i = 0; i < n; i += 1) {
gpu_run();
}
if (cuCtxSynchronize() != CUDA_SUCCESS) {
err_fail("device fault");
}
}
}

Expand All @@ -5175,7 +5227,7 @@ static void gpu_pass(u32 f) {
#define gpu_make(p) true
#define gpu_span() 0
#define gpu_load(b)
#define gpu_pass(f)
#define gpu_pass(H)

#endif

Expand Down Expand Up @@ -5213,16 +5265,17 @@ static void ring_rewind(u64* H, u32 rows) {
}

static void cube_run(u64* H, bool gpu) {
for (;;) {
u32 f = a32_exch(a32_at(H, H_CURSOR), 0);
if (root_done(H)) {
return;
}
u32 f = a32_exch(a32_at(H, H_CURSOR), 0);
if (gpu) {
turn_set(H, f);
}
while (!root_done(H)) {
if (f == 0) {
err_fail("frontier drained without a result");
}
if (gpu) {
gpu_pass(f);
gpu_pass(H);
f = (u32)H[H_FRONT];
} else {
u32 rows = f;
u32 want = (pool_size - 1) / (CUBE_T / LINE) + 1;
Expand All @@ -5236,14 +5289,15 @@ static void cube_run(u64* H, bool gpu) {
pool_turn(true, rows);
}
pool_turn(false, rows);
f = a32_load(a32_at(H, H_CURSOR));
f = a32_exch(a32_at(H, H_CURSOR), 0);
ring_rewind(H, f > rows ? f : rows);
}
u32 ec = a32_load(a32_at(H, H_ERROR_CODE));
if (ec != 0) {
err_post(H, ec);
}
}
a32_store(a32_at(H, H_CURSOR), 0);
}

// Corpus
Expand Down