spirv: Use task_payload mode for generic task outputs and mesh inputs.

This new mode will be only used for the actual payload variables and
not the number of launched mesh shader workgroups, which will still
be treated as an output.

Signed-off-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Jason Ekstrand <jason.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/14930>

src/compiler/spirv/nir_spirv.h

@@ -83,6 +83,7 @@ struct spirv_to_nir_options {
    nir_address_format phys_ssbo_addr_format;
    nir_address_format push_const_addr_format;
    nir_address_format shared_addr_format;
+   nir_address_format task_payload_addr_format;
    nir_address_format global_addr_format;
    nir_address_format temp_addr_format;
    nir_address_format constant_addr_format;

src/compiler/spirv/spirv_to_nir.c

@@ -2456,6 +2456,9 @@ vtn_mem_semantics_to_nir_var_modes(struct vtn_builder *b,
       modes |= nir_var_mem_global;
    if (semantics & SpvMemorySemanticsOutputMemoryMask) {
       modes |= nir_var_shader_out;
+
+      if (b->shader->info.stage == MESA_SHADER_TASK)
+         modes |= nir_var_mem_task_payload;
    }
 
    return modes;

src/compiler/spirv/tests/helpers.h

@@ -53,6 +53,7 @@ protected:
       spirv_options.phys_ssbo_addr_format = nir_address_format_64bit_global;
       spirv_options.push_const_addr_format = nir_address_format_32bit_offset;
       spirv_options.shared_addr_format = nir_address_format_32bit_offset;
+      spirv_options.task_payload_addr_format = nir_address_format_32bit_offset;
 
       nir_shader_compiler_options nir_options;
       memset(&nir_options, 0, sizeof(nir_options));

src/compiler/spirv/vtn_private.h

@@ -499,6 +499,7 @@ enum vtn_variable_mode {
    vtn_variable_mode_push_constant,
    vtn_variable_mode_workgroup,
    vtn_variable_mode_cross_workgroup,
+   vtn_variable_mode_task_payload,
    vtn_variable_mode_generic,
    vtn_variable_mode_constant,
    vtn_variable_mode_input,

src/compiler/spirv/vtn_variables.c

@@ -777,7 +777,9 @@ vtn_variable_copy(struct vtn_builder *b, struct vtn_pointer *dest,
 static void
 set_mode_system_value(struct vtn_builder *b, nir_variable_mode *mode)
 {
-   vtn_assert(*mode == nir_var_system_value || *mode == nir_var_shader_in);
+   vtn_assert(*mode == nir_var_system_value || *mode == nir_var_shader_in ||
+              /* Hack for NV_mesh_shader due to lack of dedicated storage class. */
+              *mode == nir_var_mem_task_payload);
    *mode = nir_var_system_value;
 }
 
@@ -1140,7 +1142,9 @@ vtn_get_builtin_location(struct vtn_builder *b,
       *location = VARYING_SLOT_PRIMITIVE_INDICES;
       break;
    case SpvBuiltInTaskCountNV:
+      /* NV_mesh_shader only. */
       *location = VARYING_SLOT_TASK_COUNT;
+      *mode = nir_var_shader_out;
       break;
    case SpvBuiltInMeshViewCountNV:
       *location = SYSTEM_VALUE_MESH_VIEW_COUNT;
@@ -1315,12 +1319,10 @@ apply_var_decoration(struct vtn_builder *b,
 
    case SpvDecorationPerTaskNV:
       vtn_fail_if(
-         !(b->shader->info.stage == MESA_SHADER_TASK && var_data->mode == nir_var_shader_out) &&
-         !(b->shader->info.stage == MESA_SHADER_MESH && var_data->mode == nir_var_shader_in),
-         "PerTaskNV decoration only allowed for Task shader outputs or Mesh shader inputs");
-      /* Don't set anything, because this decoration is implied by being a
-       * non-builtin Task Output or Mesh Input.
-       */
+         (b->shader->info.stage != MESA_SHADER_MESH &&
+          b->shader->info.stage != MESA_SHADER_TASK) ||
+         var_data->mode != nir_var_mem_task_payload,
+         "PerTaskNV decoration only allowed on Task/Mesh payload variables.");
       break;
 
    case SpvDecorationPerViewNV:
@@ -1544,10 +1546,22 @@ vtn_storage_class_to_mode(struct vtn_builder *b,
    case SpvStorageClassInput:
       mode = vtn_variable_mode_input;
       nir_mode = nir_var_shader_in;
+
+      /* NV_mesh_shader: fixup due to lack of dedicated storage class */
+      if (b->shader->info.stage == MESA_SHADER_MESH) {
+         mode = vtn_variable_mode_task_payload;
+         nir_mode = nir_var_mem_task_payload;
+      }
       break;
    case SpvStorageClassOutput:
       mode = vtn_variable_mode_output;
       nir_mode = nir_var_shader_out;
+
+      /* NV_mesh_shader: fixup due to lack of dedicated storage class */
+      if (b->shader->info.stage == MESA_SHADER_TASK) {
+         mode = vtn_variable_mode_task_payload;
+         nir_mode = nir_var_mem_task_payload;
+      }
       break;
    case SpvStorageClassPrivate:
       mode = vtn_variable_mode_private;
@@ -1643,6 +1657,9 @@ vtn_mode_to_address_format(struct vtn_builder *b, enum vtn_variable_mode mode)
    case vtn_variable_mode_accel_struct:
       return nir_address_format_64bit_global;
 
+   case vtn_variable_mode_task_payload:
+      return b->options->task_payload_addr_format;
+
    case vtn_variable_mode_function:
       if (b->physical_ptrs)
          return b->options->temp_addr_format;
@@ -1936,6 +1953,7 @@ vtn_create_variable(struct vtn_builder *b, struct vtn_value *val,
 
    case vtn_variable_mode_workgroup:
    case vtn_variable_mode_cross_workgroup:
+   case vtn_variable_mode_task_payload:
       /* Create the variable normally */
       var->var = rzalloc(b->shader, nir_variable);
       var->var->name = ralloc_strdup(var->var, val->name);
@@ -2024,16 +2042,6 @@ vtn_create_variable(struct vtn_builder *b, struct vtn_value *val,
                                           vtn_value_type_type),
                              var_decoration_cb, var);
 
-      /* PerTask I/O is always a single block without any Location, so
-       * initialize the base_location of the block and let
-       * assign_missing_member_locations() do the rest.
-       */
-      if ((b->shader->info.stage == MESA_SHADER_TASK && var->mode == vtn_variable_mode_output) ||
-          (b->shader->info.stage == MESA_SHADER_MESH && var->mode == vtn_variable_mode_input)) {
-         if (var->type->block)
-            var->base_location = VARYING_SLOT_VAR0;
-      }
-
       break;
    }