nir/divergence_analysis: Handle Task/Mesh shaders

Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/10600>

diff --git a/src/compiler/nir/nir_divergence_analysis.c b/src/compiler/nir/nir_divergence_analysis.c
index 9aecf8aee116b552b136ace0176d5a5ef91d2fa9..6a6ef29ef917bc969e3858a874420011f8d71211 100644 (file)
--- a/src/compiler/nir/nir_divergence_analysis.c
+++ b/src/compiler/nir/nir_divergence_analysis.c
@@ -167,7 +167,7 @@ visit_intrinsic(nir_shader *shader, nir_intrinsic_instr *instr)
          is_divergent |= !(options & nir_divergence_single_prim_per_subgroup);
       else if (stage == MESA_SHADER_TESS_EVAL)
          is_divergent |= !(options & nir_divergence_single_patch_per_tes_subgroup);
-      else
+      else if (stage != MESA_SHADER_MESH)
          is_divergent = true;
       break;
    case nir_intrinsic_load_per_vertex_input:
@@ -186,18 +186,33 @@ visit_intrinsic(nir_shader *shader, nir_intrinsic_instr *instr)
       is_divergent |= !(options & nir_divergence_single_prim_per_subgroup);
       break;
    case nir_intrinsic_load_output:
-      assert(stage == MESA_SHADER_TESS_CTRL || stage == MESA_SHADER_FRAGMENT);
       is_divergent = instr->src[0].ssa->divergent;
-      if (stage == MESA_SHADER_TESS_CTRL)
+      switch (stage) {
+      case MESA_SHADER_TESS_CTRL:
          is_divergent |= !(options & nir_divergence_single_patch_per_tcs_subgroup);
-      else
+         break;
+      case MESA_SHADER_FRAGMENT:
          is_divergent = true;
+         break;
+      case MESA_SHADER_TASK:
+      case MESA_SHADER_MESH:
+         /* Divergent if src[0] is, so nothing else to do. */
+         break;
+      default:
+         unreachable("Invalid stage for load_output");
+      }
       break;
    case nir_intrinsic_load_per_vertex_output:
-      assert(stage == MESA_SHADER_TESS_CTRL);
+      assert(stage == MESA_SHADER_TESS_CTRL || stage == MESA_SHADER_MESH);
       is_divergent = instr->src[0].ssa->divergent ||
                      instr->src[1].ssa->divergent ||
-                     !(options & nir_divergence_single_patch_per_tcs_subgroup);
+                     (stage == MESA_SHADER_TESS_CTRL &&
+                      !(options & nir_divergence_single_patch_per_tcs_subgroup));
+      break;
+   case nir_intrinsic_load_per_primitive_output:
+      assert(stage == MESA_SHADER_MESH);
+      is_divergent = instr->src[0].ssa->divergent ||
+                     instr->src[1].ssa->divergent;
       break;
    case nir_intrinsic_load_layer_id:
    case nir_intrinsic_load_front_face:
@@ -245,8 +260,9 @@ visit_intrinsic(nir_shader *shader, nir_intrinsic_instr *instr)
       break;
 
    case nir_intrinsic_load_workgroup_id:
-      assert(stage == MESA_SHADER_COMPUTE);
-      is_divergent |= (options & nir_divergence_multiple_workgroup_per_compute_subgroup);
+      assert(gl_shader_stage_uses_workgroup(stage));
+      if (stage == MESA_SHADER_COMPUTE)
+         is_divergent |= (options & nir_divergence_multiple_workgroup_per_compute_subgroup);
       break;
 
    /* Clustered reductions are uniform if cluster_size == subgroup_size or