nir, spirv: Don't mark NV_mesh_shader primitive indices as per-primitive.

They are not per-primitive in NV_mesh_shader, but a flat array.

Signed-off-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Caio Oliveira <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/13466>

src/compiler/nir/nir_gather_info.c

@@ -170,7 +170,10 @@ mark_whole_variable(nir_shader *shader, nir_variable *var,
 {
    const struct glsl_type *type = var->type;
 
-   if (nir_is_arrayed_io(var, shader->info.stage)) {
+   if (nir_is_arrayed_io(var, shader->info.stage) ||
+       /* For NV_mesh_shader. */
+       (shader->info.stage == MESA_SHADER_MESH &&
+        var->data.location == VARYING_SLOT_PRIMITIVE_INDICES)) {
       assert(glsl_type_is_array(type));
       type = glsl_get_array_element(type);
    }
@@ -194,7 +197,8 @@ mark_whole_variable(nir_shader *shader, nir_variable *var,
 }
 
 static unsigned
-get_io_offset(nir_deref_instr *deref, nir_variable *var, bool is_arrayed)
+get_io_offset(nir_deref_instr *deref, nir_variable *var, bool is_arrayed,
+              bool skip_non_arrayed)
 {
    if (var->data.compact) {
       assert(deref->deref_type == nir_deref_type_array);
@@ -210,6 +214,9 @@ get_io_offset(nir_deref_instr *deref, nir_variable *var, bool is_arrayed)
          if (is_arrayed && nir_deref_instr_parent(d)->deref_type == nir_deref_type_var)
             break;
 
+         if (!is_arrayed && skip_non_arrayed)
+            break;
+
          if (!nir_src_is_const(d->arr.index))
             return -1;
 
@@ -240,6 +247,7 @@ try_mask_partial_io(nir_shader *shader, nir_variable *var,
 {
    const struct glsl_type *type = var->type;
    bool is_arrayed = nir_is_arrayed_io(var, shader->info.stage);
+   bool skip_non_arrayed = shader->info.stage == MESA_SHADER_MESH;
 
    if (is_arrayed) {
       assert(glsl_type_is_array(type));
@@ -250,7 +258,7 @@ try_mask_partial_io(nir_shader *shader, nir_variable *var,
    if (var->data.per_view)
       return false;
 
-   unsigned offset = get_io_offset(deref, var, is_arrayed);
+   unsigned offset = get_io_offset(deref, var, is_arrayed, skip_non_arrayed);
    if (offset == -1)
       return false;
 

src/compiler/nir/nir_lower_io.c

@@ -152,6 +152,12 @@ nir_is_arrayed_io(const nir_variable *var, gl_shader_stage stage)
    if (var->data.patch || !glsl_type_is_array(var->type))
       return false;
 
+   if (stage == MESA_SHADER_MESH) {
+      /* NV_mesh_shader: this is flat array for the whole workgroup. */
+      if (var->data.location == VARYING_SLOT_PRIMITIVE_INDICES)
+         return false;
+   }
+
    if (var->data.mode == nir_var_shader_in)
       return stage == MESA_SHADER_GEOMETRY ||
              stage == MESA_SHADER_TESS_CTRL ||
@@ -438,6 +444,13 @@ emit_store(struct lower_io_state *state, nir_ssa_def *data,
       var->data.precision == GLSL_PRECISION_MEDIUM ||
       var->data.precision == GLSL_PRECISION_LOW;
    semantics.per_view = var->data.per_view;
+
+   /* NV_mesh_shader: prevent assigning several slots to primitive indices. */
+   if (b->shader->info.stage == MESA_SHADER_MESH &&
+       var->data.location == VARYING_SLOT_PRIMITIVE_INDICES &&
+       !nir_is_arrayed_io(var, b->shader->info.stage))
+      semantics.num_slots = 1;
+
    nir_intrinsic_set_io_semantics(store, semantics);
 
    nir_builder_instr_insert(b, &store->instr);
@@ -2684,6 +2697,15 @@ add_const_offset_to_base_block(nir_block *block, nir_builder *b,
 
       if (((modes & nir_var_shader_in) && is_input(intrin)) ||
           ((modes & nir_var_shader_out) && is_output(intrin))) {
+         nir_io_semantics sem = nir_intrinsic_io_semantics(intrin);
+
+         /* NV_mesh_shader: ignore MS primitive indices. */
+         if (b->shader->info.stage == MESA_SHADER_MESH &&
+             sem.location == VARYING_SLOT_PRIMITIVE_INDICES &&
+             !(b->shader->info.per_primitive_outputs &
+               BITFIELD64_BIT(VARYING_SLOT_PRIMITIVE_INDICES)))
+            continue;
+
          nir_src *offset = nir_get_io_offset_src(intrin);
 
          /* TODO: Better handling of per-view variables here */
@@ -2693,7 +2715,6 @@ add_const_offset_to_base_block(nir_block *block, nir_builder *b,
 
             nir_intrinsic_set_base(intrin, nir_intrinsic_base(intrin) + off);
 
-            nir_io_semantics sem = nir_intrinsic_io_semantics(intrin);
             sem.location += off;
             /* non-indirect indexing should reduce num_slots */
             sem.num_slots = is_dual_slot(intrin) ? 2 : 1;

src/compiler/spirv/vtn_variables.c

@@ -1329,19 +1329,6 @@ gather_var_kind_cb(struct vtn_builder *b, struct vtn_value *val, int member,
    case SpvDecorationPerPrimitiveNV:
       vtn_var->var->data.per_primitive = true;
       break;
-   case SpvDecorationBuiltIn:
-      if (b->shader->info.stage == MESA_SHADER_MESH) {
-         SpvBuiltIn builtin = dec->operands[0];
-         switch (builtin) {
-         case SpvBuiltInPrimitiveIndicesNV:
-            vtn_var->var->data.per_primitive = true;
-            break;
-         default:
-            /* Nothing to do. */
-            break;
-         }
-      }
-      break;
    default:
       /* Nothing to do. */
       break;