radv: export implicit primitive ID in NIR for legacy VS or TES

It's implicit for VS or TES, while it's required for GS or MS.

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

src/amd/compiler/aco_instruction_selection.cpp

@@ -8872,6 +8872,9 @@ visit_intrinsic(isel_context* ctx, nir_intrinsic_instr* instr)
              * even if there is no SW GS. */
             bld.copy(Definition(dst), get_arg(ctx, ctx->args->ac.gs_prim_id));
             break;
+         } else if (ctx->shader->info.stage == MESA_SHADER_VERTEX) {
+            bld.copy(Definition(dst), get_arg(ctx, ctx->args->ac.vs_prim_id));
+            break;
          }
          unreachable("Unimplemented shader stage for nir_intrinsic_load_primitive_id");
       }
@@ -10632,16 +10635,6 @@ create_vs_exports(isel_context* ctx)
    assert(outinfo);
    ctx->block->kind |= block_kind_export_end;
 
-   if (outinfo->export_prim_id && ctx->stage.hw != HWStage::NGG) {
-      ctx->outputs.mask[VARYING_SLOT_PRIMITIVE_ID] |= 0x1;
-      if (ctx->stage.has(SWStage::TES))
-         ctx->outputs.temps[VARYING_SLOT_PRIMITIVE_ID * 4u] =
-            get_arg(ctx, ctx->args->ac.tes_patch_id);
-      else
-         ctx->outputs.temps[VARYING_SLOT_PRIMITIVE_ID * 4u] =
-            get_arg(ctx, ctx->args->ac.vs_prim_id);
-   }
-
    /* Hardware requires position data to always be exported, even if the
     * application did not write gl_Position.
     */

src/amd/llvm/ac_nir_to_llvm.c

@@ -3684,6 +3684,8 @@ static void visit_intrinsic(struct ac_nir_context *ctx, nir_intrinsic_instr *ins
       } else if (ctx->stage == MESA_SHADER_TESS_EVAL) {
          result = ctx->tes_patch_id_replaced ? ctx->tes_patch_id_replaced
                                              : ac_get_arg(&ctx->ac, ctx->args->tes_patch_id);
+      } else if (ctx->stage == MESA_SHADER_VERTEX) {
+         result = ac_get_arg(&ctx->ac, ctx->args->vs_prim_id);
       } else
          fprintf(stderr, "Unknown primitive id intrinsic: %d", ctx->stage);
       break;

src/amd/vulkan/radv_nir_to_llvm.c

@@ -1046,7 +1046,7 @@ handle_vs_outputs_post(struct radv_shader_context *ctx, bool export_prim_id, boo
    }
 
    /* Allocate a temporary array for the output values. */
-   unsigned num_outputs = util_bitcount64(ctx->output_mask) + export_prim_id;
+   unsigned num_outputs = util_bitcount64(ctx->output_mask);
    outputs = malloc(num_outputs * sizeof(outputs[0]));
 
    for (unsigned i = 0; i < AC_LLVM_MAX_OUTPUTS; ++i) {
@@ -1072,20 +1072,6 @@ handle_vs_outputs_post(struct radv_shader_context *ctx, bool export_prim_id, boo
       noutput++;
    }
 
-   /* Export PrimitiveID. */
-   if (export_prim_id) {
-      outputs[noutput].slot_name = VARYING_SLOT_PRIMITIVE_ID;
-      outputs[noutput].slot_index = 0;
-      outputs[noutput].usage_mask = 0x1;
-      if (ctx->stage == MESA_SHADER_TESS_EVAL)
-         outputs[noutput].values[0] = ac_get_arg(&ctx->ac, ctx->args->ac.tes_patch_id);
-      else
-         outputs[noutput].values[0] = ac_get_arg(&ctx->ac, ctx->args->ac.vs_prim_id);
-      for (unsigned j = 1; j < 4; j++)
-         outputs[noutput].values[j] = ctx->ac.f32_0;
-      noutput++;
-   }
-
    radv_llvm_export_vs(ctx, outputs, noutput, outinfo, export_clip_dists);
 
    free(outputs);

src/amd/vulkan/radv_pipeline.c

@@ -2769,6 +2769,29 @@ radv_lower_multiview(nir_shader *nir)
    return progress;
 }
 
+static bool
+radv_export_implicit_primitive_id(nir_shader *nir)
+{
+   nir_function_impl *impl = nir_shader_get_entrypoint(nir);
+   nir_builder b;
+   nir_builder_init(&b, impl);
+
+   b.cursor = nir_after_cf_list(&impl->body);
+
+   nir_variable *var = nir_variable_create(nir, nir_var_shader_out, glsl_int_type(), NULL);
+   var->data.location = VARYING_SLOT_PRIMITIVE_ID;
+   var->data.interpolation = INTERP_MODE_NONE;
+
+   nir_store_var(&b, var, nir_load_primitive_id(&b), 1);
+
+   /* Update outputs_written to reflect that the pass added a new output. */
+   nir->info.outputs_written |= BITFIELD64_BIT(VARYING_SLOT_PRIMITIVE_ID);
+
+   nir_metadata_preserve(impl, nir_metadata_block_index | nir_metadata_dominance);
+
+   return true;
+}
+
 static void
 radv_link_shaders(struct radv_pipeline *pipeline,
                   const struct radv_pipeline_key *pipeline_key,
@@ -2876,6 +2899,19 @@ radv_link_shaders(struct radv_pipeline *pipeline,
       }
    }
 
+   /* Export the primitive ID when VS or TES don't export it because it's implicit, while it's
+    * required for GS or MS. The primitive ID is added during lowering for NGG.
+    */
+   if (stages[MESA_SHADER_FRAGMENT].nir &&
+       (stages[MESA_SHADER_FRAGMENT].nir->info.inputs_read & VARYING_BIT_PRIMITIVE_ID) &&
+       !(stages[pipeline->graphics.last_vgt_api_stage].nir->info.outputs_written & VARYING_BIT_PRIMITIVE_ID) &&
+       ((pipeline->graphics.last_vgt_api_stage == MESA_SHADER_VERTEX &&
+         !stages[MESA_SHADER_VERTEX].info.is_ngg) ||
+        (pipeline->graphics.last_vgt_api_stage == MESA_SHADER_TESS_EVAL &&
+         !stages[MESA_SHADER_TESS_EVAL].info.is_ngg))) {
+      radv_export_implicit_primitive_id(stages[pipeline->graphics.last_vgt_api_stage].nir);
+   }
+
    if (!optimize_conservatively) {
       bool uses_xfb = pipeline->graphics.last_vgt_api_stage != -1 &&
                       radv_nir_stage_uses_xfb(stages[pipeline->graphics.last_vgt_api_stage].nir);