ac/nir/ngg: Add ability to store primitive ID as per-primitive.

This configuration will be enabled in RADV in a subsequent commit.

On GFX10.3:
Do this together with the primitive export, to avoid adding extra
CF, and to ensure optimal access of the export space.

On GFX11:
It's not an export but a memory store instruction, so always do
it earlier and ensure the optimal attribute ring access pattern.

Signed-off-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Georg Lehmann <dadschoorse@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32270>

src/amd/common/ac_nir.h

@@ -142,6 +142,7 @@ typedef struct {
    bool passthrough;
    bool use_edgeflags;
    bool export_primitive_id;
+   bool export_primitive_id_per_prim;
    uint32_t instance_rate_inputs;
    uint32_t user_clip_plane_enable_mask;
 

src/amd/common/ac_nir_lower_ngg.c

@@ -632,6 +632,18 @@ emit_ngg_nogs_prim_export(nir_builder *b, lower_ngg_nogs_state *s, nir_def *arg)
       }
 
       ac_nir_export_primitive(b, arg, NULL);
+
+      /* Store implicit primitive ID when configured as a per-primitive output on GFX10.3.
+       * Because this uses the export space, do it together with the primitive export.
+       */
+      if (s->options->gfx_level == GFX10_3 && s->options->export_primitive_id_per_prim) {
+         const uint8_t offset = s->options->vs_output_param_offset[VARYING_SLOT_PRIMITIVE_ID];
+         nir_def *prim_id = nir_load_primitive_id(b);
+         nir_def *undef = nir_undef(b, 1, 32);
+         nir_def *param_components[4] = { prim_id, undef, undef, undef };
+
+         ac_nir_export_parameters(b, &offset, 1, 0, &param_components, NULL, NULL);
+      }
    }
    nir_pop_if(b, if_gs_thread);
 }
@@ -665,6 +677,46 @@ emit_ngg_nogs_prim_id_store_shared(nir_builder *b, lower_ngg_nogs_state *s)
    nir_pop_if(b, if_gs_thread);
 }
 
+/* Store implicit primitive ID when configured as a per-primitive output on GFX11+.
+ * This is done separately from the primitive export on GFX11 in order to
+ * optimize attribute ring access.
+ */
+static void
+emit_ngg_nogs_prim_id_store_per_prim_gfx11(nir_builder *b, lower_ngg_nogs_state *s)
+{
+   assert( s->options->gfx_level >= GFX11);
+
+   const uint8_t offset = s->options->vs_output_param_offset[VARYING_SLOT_PRIMITIVE_ID];
+
+   nir_def *is_gs_thread = nir_load_var(b, s->gs_exported_var);
+   nir_def *prim_id = nir_load_primitive_id(b);
+   nir_def *undef = nir_undef(b, 1, 32);
+   nir_def *param_components[4] = { prim_id, undef, undef, undef };
+
+   /* Emit this outside of CF to allow CSE to pick it up. */
+   nir_def *vindex = nir_load_local_invocation_index(b);
+
+   /* Align the number of active invocations to 8 for optimal attribute ring access pattern,
+    * see export_vertex_params_gfx11 for full explanation.
+    */
+   nir_def *highest_gs_thread = nir_ufind_msb(b, nir_ballot(b, 1, s->options->wave_size, is_gs_thread));
+   nir_def *max_num_gs_threads = nir_iadd_imm_nuw(b, highest_gs_thread, 1);
+   nir_def *num_gs_attr_ring_threads = nir_iand_imm(b, nir_iadd_imm(b, max_num_gs_threads, 7), ~7);
+
+   nir_if *if_store_gs_attr_ring = nir_push_if(b, nir_is_subgroup_invocation_lt_amd(b, num_gs_attr_ring_threads));
+   {
+      nir_def *attr_rsrc = nir_load_ring_attr_amd(b);
+      nir_def *attr_offset = nir_load_ring_attr_offset_amd(b);
+      nir_def *voffset = nir_imm_int(b, 0);
+
+      nir_store_buffer_amd(b, nir_vec(b, param_components, 4), attr_rsrc, voffset, attr_offset, vindex,
+                           .base = offset * 16,
+                           .memory_modes = nir_var_shader_out,
+                           .access = ACCESS_COHERENT | ACCESS_IS_SWIZZLED_AMD);
+   }
+   nir_pop_if(b, if_store_gs_attr_ring);
+}
+
 static void
 emit_store_ngg_nogs_es_primitive_id(nir_builder *b, lower_ngg_nogs_state *s)
 {
@@ -2682,9 +2734,12 @@ ac_nir_lower_ngg_nogs(nir_shader *shader, const ac_nir_lower_ngg_options *option
       .gs_exported_var = gs_exported_var,
       .max_num_waves = DIV_ROUND_UP(options->max_workgroup_size, options->wave_size),
       .has_user_edgeflags = has_user_edgeflags,
-      .skip_primitive_id = streamout_enabled && options->export_primitive_id,
+      .skip_primitive_id = streamout_enabled && (options->export_primitive_id || options->export_primitive_id_per_prim),
    };
 
+   /* Can't export the primitive ID both as per-vertex and per-primitive. */
+   assert(!options->export_primitive_id || !options->export_primitive_id_per_prim);
+
    const bool need_prim_id_store_shared =
       options->export_primitive_id && shader->info.stage == MESA_SHADER_VERTEX;
 
@@ -2692,6 +2747,12 @@ ac_nir_lower_ngg_nogs(nir_shader *shader, const ac_nir_lower_ngg_options *option
       shader->info.outputs_written |= VARYING_BIT_PRIMITIVE_ID;
    }
 
+   if (options->export_primitive_id_per_prim) {
+      /* The HW preloads the primitive ID to VGPRs of GS threads for VS, but not for TES. */
+      assert(shader->info.stage == MESA_SHADER_VERTEX);
+      assert(options->gfx_level >= GFX10_3);
+   }
+
    nir_builder builder = nir_builder_create(impl);
    nir_builder *b = &builder; /* This is to avoid the & */
 
@@ -2765,6 +2826,8 @@ ac_nir_lower_ngg_nogs(nir_shader *shader, const ac_nir_lower_ngg_options *option
       /* Wait for GS threads to store primitive ID in LDS. */
       nir_barrier(b, .execution_scope = SCOPE_WORKGROUP, .memory_scope = SCOPE_WORKGROUP,
                             .memory_semantics = NIR_MEMORY_ACQ_REL, .memory_modes = nir_var_mem_shared);
+   } else if (options->export_primitive_id_per_prim && options->gfx_level >= GFX11) {
+      emit_ngg_nogs_prim_id_store_per_prim_gfx11(b, &state);
    }
 
    nir_def *es_thread =