libagx: accelerate prim restart unroll across wg

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/28483>
This commit is contained in:
Alyssa Rosenzweig
2024-02-28 22:23:46 -04:00
committed by Marge Bot
parent 890a96e2a7
commit 13ecef56d0
3 changed files with 36 additions and 22 deletions
+3 -1
View File
@@ -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)
+32 -20
View File
@@ -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; \
} \
+1 -1
View File
@@ -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},
};