From 13ecef56d06ded1b872c62fc6251a54bad88e36e Mon Sep 17 00:00:00 2001 From: Alyssa Rosenzweig Date: Wed, 28 Feb 2024 22:23:46 -0400 Subject: [PATCH] libagx: accelerate prim restart unroll across wg Signed-off-by: Alyssa Rosenzweig Part-of: --- src/asahi/lib/agx_nir_lower_gs.c | 4 ++- src/asahi/lib/shaders/geometry.cl | 52 ++++++++++++++++----------- src/gallium/drivers/asahi/agx_state.c | 2 +- 3 files changed, 36 insertions(+), 22 deletions(-) diff --git a/src/asahi/lib/agx_nir_lower_gs.c b/src/asahi/lib/agx_nir_lower_gs.c index ede3620fa9d..14a20f6ab53 100644 --- a/src/asahi/lib/agx_nir_lower_gs.c +++ b/src/asahi/lib/agx_nir_lower_gs.c @@ -1394,10 +1394,12 @@ agx_nir_gs_setup_indirect(nir_builder *b, const void *data) void agx_nir_unroll_restart(nir_builder *b, const void *data) { + b->shader->info.workgroup_size[0] = 1024; + const struct agx_unroll_restart_key *key = data; nir_def *ia = nir_load_input_assembly_buffer_agx(b); nir_def *draw = nir_channel(b, nir_load_workgroup_id(b), 0); - nir_def *lane = nir_load_subgroup_invocation(b); + nir_def *lane = nir_channel(b, nir_load_local_invocation_id(b), 0); nir_def *mode = nir_imm_int(b, key->prim); if (key->index_size_B == 1) diff --git a/src/asahi/lib/shaders/geometry.cl b/src/asahi/lib/shaders/geometry.cl index b637effb120..760aa9a884c 100644 --- a/src/asahi/lib/shaders/geometry.cl +++ b/src/asahi/lib/shaders/geometry.cl @@ -234,13 +234,20 @@ libagx_vertex_id_for_topology(enum mesa_prim mode, bool flatshade_first, } } -/* Return the ID of the first thread where cond is true. cond must be true for - * some thread. +/* + * Return the ID of the first thread in the workgroup where cond is true, or + * 1024 if cond is false across the workgroup. */ static uint -first_true_thread(bool cond) +first_true_thread_in_workgroup(bool cond, local uint *scratch) { - return ctz(ballot(cond)); + barrier(CLK_LOCAL_MEM_FENCE); + scratch[get_sub_group_id()] = ballot(cond); + barrier(CLK_LOCAL_MEM_FENCE); + + uint first_group = ctz(ballot(scratch[get_sub_group_local_id()])); + uint off = ctz(first_group < 32 ? scratch[first_group] : 0); + return (first_group * 32) + off; } /* @@ -281,9 +288,9 @@ setup_unroll_for_draw(global struct agx_ia_state *ia, constant uint *in_draw, } #define UNROLL(INDEX, suffix) \ - void libagx_unroll_restart_##suffix(global struct agx_ia_state *ia, \ - enum mesa_prim mode, uint draw, \ - uint tid) \ + kernel void libagx_unroll_restart_##suffix(global struct agx_ia_state *ia, \ + enum mesa_prim mode, uint draw, \ + uint tid) \ { \ /* For an indirect multidraw, we are dispatched maxDraws times and \ * terminate trailing invocations. \ @@ -298,10 +305,16 @@ setup_unroll_for_draw(global struct agx_ia_state *ia, constant uint *in_draw, constant INDEX *in = (constant INDEX *)ia->index_buffer; \ in += in_draw[2]; \ \ - global INDEX *out; \ - if (tid == 0) \ - out = setup_unroll_for_draw(ia, in_draw, draw, mode, sizeof(INDEX)); \ - out = (global INDEX *)sub_group_broadcast((uintptr_t)out, 0); \ + local uintptr_t out_ptr; \ + if (tid == 0) { \ + out_ptr = (uintptr_t)setup_unroll_for_draw(ia, in_draw, draw, mode, \ + sizeof(INDEX)); \ + } \ + \ + barrier(CLK_LOCAL_MEM_FENCE); \ + global INDEX *out = (global INDEX *)out_ptr; \ + \ + local uint scratch[32]; \ \ uint out_prims = 0; \ INDEX restart_idx = ia->restart_index; \ @@ -314,24 +327,22 @@ setup_unroll_for_draw(global struct agx_ia_state *ia, constant uint *in_draw, /* Search for next restart or the end. Lanes load in parallel. */ \ uint next_restart = needle; \ for (;;) { \ - uint idx = next_restart + tid; \ /* Relies on shortcircuiting */ \ + uint idx = next_restart + tid; \ bool restart = idx >= count || in[idx] == restart_idx; \ - if (sub_group_any(restart)) { \ - uint first_restart_tid = first_true_thread(restart); \ - next_restart += first_restart_tid; \ + \ + uint next_offs = first_true_thread_in_workgroup(restart, scratch); \ + \ + next_restart += next_offs; \ + if (next_offs < 1024) \ break; \ - } else { \ - /* No restart in this group, try the next */ \ - next_restart += get_sub_group_size(); \ - } \ } \ \ /* Emit up to the next restart. Lanes output in parallel */ \ uint subcount = next_restart - needle; \ uint subprims = u_decomposed_prims_for_vertices(mode, subcount); \ uint out_prims_base = out_prims; \ - for (uint i = tid; i < subprims; i += get_sub_group_size()) { \ + for (uint i = tid; i < subprims; i += 1024) { \ for (uint vtx = 0; vtx < per_prim; ++vtx) { \ uint id = libagx_vertex_id_for_topology(mode, flatshade_first, \ i, vtx, subprims); \ @@ -341,6 +352,7 @@ setup_unroll_for_draw(global struct agx_ia_state *ia, constant uint *in_draw, offset < in_size_el ? in[offset] : 0; \ } \ } \ + \ out_prims += subprims; \ needle = next_restart + 1; \ } \ diff --git a/src/gallium/drivers/asahi/agx_state.c b/src/gallium/drivers/asahi/agx_state.c index ae7a1d9aaa0..5acdb666e63 100644 --- a/src/gallium/drivers/asahi/agx_state.c +++ b/src/gallium/drivers/asahi/agx_state.c @@ -4356,7 +4356,7 @@ agx_draw_without_restart(struct agx_batch *batch, /* Unroll the index buffer for each draw */ const struct pipe_grid_info grid_setup = { - .block = {32, 1, 1}, + .block = {1024, 1, 1}, .grid = {indirect->draw_count, 1, 1}, };