nir,zink,asahi: support passing through gl_PrimitiveID

When this pass is used with Zink, gl_PrimitiveID needs to be passed
through, however this is unnecessary for other divers.

Analogous to previous commit

Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Fixes: d0342e28b3 ("nir: Add helper to create passthrough GS shader")
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32397>

src/compiler/nir/nir.h

@@ -6126,7 +6126,8 @@ nir_shader *nir_create_passthrough_gs(const nir_shader_compiler_options *options
                                       enum mesa_prim primitive_type,
                                       enum mesa_prim output_primitive_type,
                                       bool emulate_edgeflags,
-                                      bool force_line_strip_out);
+                                      bool force_line_strip_out,
+                                      bool passthrough_prim_id);
 
 bool nir_lower_fragcolor(nir_shader *shader, unsigned max_cbufs);
 bool nir_lower_fragcoord_wtrans(nir_shader *shader);

src/compiler/nir/nir_passthrough_gs.c

@@ -123,7 +123,8 @@ nir_create_passthrough_gs(const nir_shader_compiler_options *options,
                           enum mesa_prim primitive_type,
                           enum mesa_prim output_primitive_type,
                           bool emulate_edgeflags,
-                          bool force_line_strip_out)
+                          bool force_line_strip_out,
+                          bool passthrough_prim_id)
 {
    unsigned int vertices_out = vertices_for_prim(primitive_type);
    emulate_edgeflags = emulate_edgeflags && (prev_stage->info.outputs_written & VARYING_BIT_EDGE);
@@ -158,6 +159,8 @@ nir_create_passthrough_gs(const nir_shader_compiler_options *options,
    /* Create input/output variables. */
    nir_foreach_shader_out_variable(var, prev_stage) {
       assert(!var->data.patch);
+      assert(var->data.location != VARYING_SLOT_PRIMITIVE_ID &&
+            "not a VS output");
 
       /* input vars can't be created for those */
       if (var->data.location == VARYING_SLOT_LAYER ||
@@ -200,6 +203,24 @@ nir_create_passthrough_gs(const nir_shader_compiler_options *options,
       out_vars[num_outputs++] = out;
    }
 
+   if (passthrough_prim_id) {
+      /* When a geometry shader is not used, a fragment shader may read primitive
+       * ID and get an implicit value without the vertex shader writing an ID. This
+       * case needs to work even when we inject a GS internally.
+       *
+       * However, if a geometry shader precedes a fragment shader that reads
+       * primitive ID, Vulkan requires that the geometry shader write primitive ID.
+       * To handle this case correctly, we must write primitive ID, copying the
+       * fixed-function gl_PrimitiveIDIn input which matches what the fragment
+       * shader will expect.
+       */
+      in_vars[num_inputs++] = nir_create_variable_with_location(
+            nir, nir_var_shader_in, VARYING_SLOT_PRIMITIVE_ID, glsl_int_type());
+
+      out_vars[num_outputs++] = nir_create_variable_with_location(
+            nir, nir_var_shader_out, VARYING_SLOT_PRIMITIVE_ID, glsl_int_type());
+   }
+
    unsigned int start_vert = 0;
    unsigned int end_vert = vertices_out;
    unsigned int vert_step = 1;
@@ -240,7 +261,12 @@ nir_create_passthrough_gs(const nir_shader_compiler_options *options,
             uint64_t mask = BITFIELD64_BIT(in_vars[j]->data.location);
             index = nir_bcsel(&b, nir_ieq_imm(&b, nir_iand_imm(&b, flat_interp_mask_def, mask), 0), nir_imm_int(&b, idx), pv_vert_index);
          }
-         nir_deref_instr *value = nir_build_deref_array(&b, nir_build_deref_var(&b, in_vars[j]), index);
+
+         /* gl_PrimitiveIDIn is not arrayed, all other inputs are */
+         nir_deref_instr *value = nir_build_deref_var(&b, in_vars[j]);
+         if (in_vars[j]->data.location != VARYING_SLOT_PRIMITIVE_ID)
+            value = nir_build_deref_array(&b, value, index);
+
          copy_vars(&b, nir_build_deref_var(&b, out_vars[oj]), value);
          ++oj;
       }

src/gallium/drivers/asahi/agx_state.c

@@ -4341,7 +4341,7 @@ agx_get_passthrough_gs(struct agx_context *ctx,
 
    nir_shader *gs = nir_create_passthrough_gs(
       &agx_nir_options, prev, mode, rast_prim(mode, poly_mode), edgeflags,
-      false /* force line strip out */);
+      false /* force line strip out */, false);
 
    ralloc_free(prev);
 

src/gallium/drivers/zink/zink_program.c

@@ -2584,7 +2584,8 @@ zink_set_primitive_emulation_keys(struct zink_context *ctx)
                   prim,
                   ctx->gfx_pipeline_state.rast_prim,
                   lower_edge_flags,
-                  lower_line_stipple || lower_quad_prim);
+                  lower_line_stipple || lower_quad_prim,
+                  true);
             }
             zink_lower_system_values_to_inlined_uniforms(nir);