Mesa (main): ac/nir/ngg: count the number of generated primitives for VS and TES

[GitLab Mirror]

Module: Mesa
Branch: main
Commit: fe57fe1fd8ad31677f977b44b38b6b971770511c
URL:    http://cgit.freedesktop.org/mesa/mesa/commit/?id=fe57fe1fd8ad31677f977b44b38b6b971770511c

Author: Samuel Pitoiset <samuel.pitoiset at gmail.com>
Date:   Fri May 27 13:51:12 2022 +0200

ac/nir/ngg: count the number of generated primitives for VS and TES

Signed-off-by: Samuel Pitoiset <samuel.pitoiset at gmail.com>
Reviewed-by: Timur Kristóf <timur.kristof at gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/15639>

---

 src/amd/common/ac_nir.h           |  1 +
 src/amd/common/ac_nir_lower_ngg.c | 21 +++++++++++++++++++++
 src/amd/vulkan/radv_shader.c      |  3 ++-
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/amd/common/ac_nir.h b/src/amd/common/ac_nir.h
index f5c59acbcf9..a532db6e289 100644
--- a/src/amd/common/ac_nir.h
+++ b/src/amd/common/ac_nir.h
@@ -125,6 +125,7 @@ ac_nir_lower_ngg_nogs(nir_shader *shader,
                       bool export_prim_id,
                       bool provoking_vtx_last,
                       bool use_edgeflags,
+                      bool has_prim_query,
                       uint32_t instance_rate_inputs);
 
 void
diff --git a/src/amd/common/ac_nir_lower_ngg.c b/src/amd/common/ac_nir_lower_ngg.c
index 0631964fb06..69caffe203d 100644
--- a/src/amd/common/ac_nir_lower_ngg.c
+++ b/src/amd/common/ac_nir_lower_ngg.c
@@ -53,6 +53,7 @@ typedef struct
    bool export_prim_id;
    bool early_prim_export;
    bool use_edgeflags;
+   bool has_prim_query;
    unsigned wave_size;
    unsigned max_num_waves;
    unsigned num_vertices_per_primitives;
@@ -445,6 +446,24 @@ emit_ngg_nogs_prim_export(nir_builder *b, lower_ngg_nogs_state *st, nir_ssa_def
          nir_pop_if(b, if_prim_valid);
       }
 
+      if (st->has_prim_query) {
+         nir_if *if_shader_query = nir_push_if(b, nir_load_shader_query_enabled_amd(b));
+         {
+            /* Number of active GS threads. Each has 1 output primitive. */
+            nir_ssa_def *num_gs_threads = nir_bit_count(b, nir_ballot(b, 1, st->wave_size, nir_imm_bool(b, true)));
+            /* Activate only 1 lane and add the number of primitives to GDS. */
+            nir_if *if_elected = nir_push_if(b, nir_elect(b, 1));
+            {
+               /* Use a different GDS offset than NGG GS to ensure that pipeline statistics
+                * queries won't return the number of primitives generated by VS/TES.
+                */
+               nir_gds_atomic_add_amd(b, 32, num_gs_threads, nir_imm_int(b, 4), nir_imm_int(b, 0x100));
+            }
+            nir_pop_if(b, if_elected);
+         }
+         nir_pop_if(b, if_shader_query);
+      }
+
       nir_export_primitive_amd(b, arg);
    }
    nir_pop_if(b, if_gs_thread);
@@ -1350,6 +1369,7 @@ ac_nir_lower_ngg_nogs(nir_shader *shader,
                       bool export_prim_id,
                       bool provoking_vtx_last,
                       bool use_edgeflags,
+                      bool has_prim_query,
                       uint32_t instance_rate_inputs)
 {
    nir_function_impl *impl = nir_shader_get_entrypoint(shader);
@@ -1367,6 +1387,7 @@ ac_nir_lower_ngg_nogs(nir_shader *shader,
       .export_prim_id = export_prim_id,
       .early_prim_export = early_prim_export,
       .use_edgeflags = use_edgeflags,
+      .has_prim_query = has_prim_query,
       .num_vertices_per_primitives = num_vertices_per_primitives,
       .provoking_vtx_idx = provoking_vtx_last ? (num_vertices_per_primitives - 1) : 0,
       .position_value_var = position_value_var,
diff --git a/src/amd/vulkan/radv_shader.c b/src/amd/vulkan/radv_shader.c
index 4dc180f8616..178ca2bf3e3 100644
--- a/src/amd/vulkan/radv_shader.c
+++ b/src/amd/vulkan/radv_shader.c
@@ -1235,7 +1235,8 @@ void radv_lower_ngg(struct radv_device *device, struct radv_pipeline_stage *ngg_
       NIR_PASS_V(nir, ac_nir_lower_ngg_nogs, max_vtx_in, num_vertices_per_prim,
                  info->workgroup_size, info->wave_size, info->has_ngg_culling,
                  info->has_ngg_early_prim_export, info->is_ngg_passthrough, export_prim_id,
-                 pl_key->vs.provoking_vtx_last, false, pl_key->vs.instance_rate_inputs);
+                 pl_key->vs.provoking_vtx_last, false, pl_key->primitives_generated_query,
+                 pl_key->vs.instance_rate_inputs);
    } else if (nir->info.stage == MESA_SHADER_GEOMETRY) {
       assert(info->is_ngg);
       NIR_PASS_V(nir, ac_nir_lower_ngg_gs, info->wave_size, info->workgroup_size,