nir: add nir_gather_tcs_info, new gathering/analysis pass

This does shader analysis that is more niche than regular shader info.

It's planned to be used by nir_restructure_tcs_flow as discussed here:
    https://gitlab.freedesktop.org/mesa/mesa/-/issues/11910

It's also useful for driver-specific passes.

The code for gathering "all_invocations_define_tess_levels" is copied
from radeonsi. The rest is new.

Acked-by: Pierre-Eric Pelloux-Prayer <pierre-eric.pelloux-prayer@amd.com>
Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Rhys Perry <pendingchaos02@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31447>

src/compiler/nir/meson.build

@@ -98,6 +98,7 @@ files_libnir = files(
   'nir_from_ssa.c',
   'nir_functions.c',
   'nir_gather_info.c',
+  'nir_gather_tcs_info.c',
   'nir_gather_types.c',
   'nir_gather_xfb_info.c',
   'nir_group_loads.c',

src/compiler/nir/nir.h

@@ -5567,6 +5567,44 @@ void nir_gather_types(nir_function_impl *impl,
                       BITSET_WORD *float_types,
                       BITSET_WORD *int_types);
 
+typedef struct {
+   /* Whether all invocations write tess level outputs.
+    *
+    * This is useful when a pass wants to read tess level values at the end
+    * of the shader. If this is true, the pass doesn't have to insert a barrier
+    * and use output loads, it can just use the SSA defs that are being stored
+    * (or phis thereof) to get the tess level output values.
+    */
+   bool all_invocations_define_tess_levels;
+
+   /* Whether any of the outer tess level components is effectively 0, meaning
+    * that the shader discards the patch. NaNs and negative values are included
+    * in this. If the patch is discarded, inner tess levels have no effect.
+    */
+   bool all_tess_levels_are_effectively_zero;
+
+   /* Whether all tess levels are effectively 1, meaning that the tessellator
+    * behaves as if they were 1. There is a range of values that lead to that
+    * behavior depending on the tessellation spacing.
+    */
+   bool all_tess_levels_are_effectively_one;
+
+   /* Whether the shader uses a barrier synchronizing TCS output stores.
+    * For example, passes that write an output at the beginning of the shader
+    * and load it at the end can use this to determine whether they have to
+    * insert a barrier or whether the shader already contains a barrier.
+    */
+   bool always_executes_barrier;
+
+   /* Whether outer tess levels <= 0 are written anywhere in the shader. */
+   bool discards_patches;
+} nir_tcs_info;
+
+void
+nir_gather_tcs_info(const nir_shader *nir, nir_tcs_info *info,
+                    enum tess_primitive_mode prim,
+                    enum gl_tess_spacing spacing);
+
 void nir_assign_var_locations(nir_shader *shader, nir_variable_mode mode,
                               unsigned *size,
                               int (*type_size)(const struct glsl_type *, bool));

src/compiler/nir/nir_gather_tcs_info.c

@@ -0,0 +1,334 @@
+/*
+ * Copyright © 2024 Advanced Micro Devices, Inc.
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+#include <math.h>
+
+static unsigned
+get_tess_level_component(nir_intrinsic_instr *intr)
+{
+   unsigned location = nir_intrinsic_io_semantics(intr).location;
+
+   return (location == VARYING_SLOT_TESS_LEVEL_INNER ? 4 : 0) +
+          nir_intrinsic_component(intr);
+}
+
+static unsigned
+get_inst_tesslevel_writemask(nir_intrinsic_instr *intr)
+{
+   if (intr->intrinsic != nir_intrinsic_store_output)
+      return 0;
+
+   unsigned location = nir_intrinsic_io_semantics(intr).location;
+   if (location != VARYING_SLOT_TESS_LEVEL_OUTER &&
+       location != VARYING_SLOT_TESS_LEVEL_INNER)
+      return 0;
+
+   return nir_intrinsic_write_mask(intr) << get_tess_level_component(intr);
+}
+
+static bool
+is_tcs_output_barrier(nir_intrinsic_instr *intr)
+{
+   return intr->intrinsic == nir_intrinsic_barrier &&
+          nir_intrinsic_memory_modes(intr) & nir_var_shader_out &&
+          nir_intrinsic_memory_scope(intr) >= SCOPE_WORKGROUP &&
+          nir_intrinsic_execution_scope(intr) >= SCOPE_WORKGROUP;
+}
+
+static void
+scan_tess_levels(struct exec_list *cf_list, unsigned *upper_block_tl_writemask,
+                 unsigned *cond_block_tl_writemask,
+                 bool *all_invocs_define_tess_levels, bool is_nested_cf)
+{
+   foreach_list_typed(nir_cf_node, cf_node, node, cf_list) {
+      switch (cf_node->type) {
+      case nir_cf_node_block: {
+         nir_block *block = nir_cf_node_as_block(cf_node);
+         nir_foreach_instr (instr, block) {
+            if (instr->type != nir_instr_type_intrinsic)
+               continue;
+
+            nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+
+            if (!is_tcs_output_barrier(intrin)) {
+               *upper_block_tl_writemask |= get_inst_tesslevel_writemask(intrin);
+               continue;
+            }
+
+            /* This is a barrier. If it's in nested control flow, put this
+             * in the too hard basket. In GLSL this is not possible but it is
+             * in SPIR-V.
+             */
+            if (is_nested_cf) {
+               *all_invocs_define_tess_levels = false;
+               return;
+            }
+
+            /* The following case must be prevented:
+             *    gl_TessLevelInner = ...;
+             *    barrier();
+             *    if (gl_InvocationID == 1)
+             *       gl_TessLevelInner = ...;
+             *
+             * If you consider disjoint code segments separated by barriers,
+             * each such segment that writes tess level channels should write
+             * the same channels in all codepaths within that segment.
+             */
+            if (*upper_block_tl_writemask || *cond_block_tl_writemask) {
+               /* Accumulate the result: */
+               *all_invocs_define_tess_levels &=
+                  !(*cond_block_tl_writemask & ~(*upper_block_tl_writemask));
+
+               /* Analyze the next code segment from scratch. */
+               *upper_block_tl_writemask = 0;
+               *cond_block_tl_writemask = 0;
+            }
+         }
+         break;
+      }
+      case nir_cf_node_if: {
+         unsigned then_tesslevel_writemask = 0;
+         unsigned else_tesslevel_writemask = 0;
+         nir_if *if_stmt = nir_cf_node_as_if(cf_node);
+
+         scan_tess_levels(&if_stmt->then_list, &then_tesslevel_writemask,
+                          cond_block_tl_writemask,
+                          all_invocs_define_tess_levels, true);
+
+         scan_tess_levels(&if_stmt->else_list, &else_tesslevel_writemask,
+                          cond_block_tl_writemask,
+                          all_invocs_define_tess_levels, true);
+
+         if (then_tesslevel_writemask || else_tesslevel_writemask) {
+            /* If both statements write the same tess level channels,
+             * we can say that the upper block writes them too.
+             */
+            *upper_block_tl_writemask |= then_tesslevel_writemask &
+                                         else_tesslevel_writemask;
+            *cond_block_tl_writemask |= then_tesslevel_writemask |
+                                        else_tesslevel_writemask;
+         }
+         break;
+      }
+      case nir_cf_node_loop: {
+         nir_loop *loop = nir_cf_node_as_loop(cf_node);
+         assert(!nir_loop_has_continue_construct(loop));
+
+         scan_tess_levels(&loop->body, cond_block_tl_writemask,
+                          cond_block_tl_writemask,
+                          all_invocs_define_tess_levels, true);
+         break;
+      }
+      default:
+         unreachable("unknown cf node type");
+      }
+   }
+}
+
+static bool
+all_invocations_define_tess_levels(const struct nir_shader *nir)
+{
+   assert(nir->info.stage == MESA_SHADER_TESS_CTRL);
+
+   /* The pass works as follows:
+    *
+    * If all codepaths write tess levels, we can say that all invocations
+    * define tess level values. Whether a tess level value is defined is
+    * determined for each component separately.
+    */
+   unsigned main_block_tl_writemask = 0; /* if main block writes tess levels */
+   unsigned cond_block_tl_writemask = 0; /* if cond block writes tess levels */
+
+   /* Initial value = true. Here the pass will accumulate results from
+    * multiple segments surrounded by barriers. If tess levels aren't
+    * written at all, it's a shader bug and we don't care if this will be
+    * true.
+    */
+   bool result = true;
+
+   nir_foreach_function_impl(impl, nir) {
+      scan_tess_levels(&impl->body, &main_block_tl_writemask,
+                       &cond_block_tl_writemask,
+                       &result, false);
+   }
+
+   /* Accumulate the result for the last code segment separated by a
+    * barrier.
+    */
+   if (main_block_tl_writemask || cond_block_tl_writemask)
+      result &= !(cond_block_tl_writemask & ~main_block_tl_writemask);
+
+   return result;
+}
+
+/* It's OK to pass UNSPECIFIED to prim and spacing. */
+void
+nir_gather_tcs_info(const nir_shader *nir, nir_tcs_info *info,
+                    enum tess_primitive_mode prim,
+                    enum gl_tess_spacing spacing)
+{
+   memset(info, 0, sizeof(*info));
+   info->all_invocations_define_tess_levels =
+      all_invocations_define_tess_levels(nir);
+
+   unsigned tess_level_writes_le_zero = 0;
+   unsigned tess_level_writes_le_one = 0;
+   unsigned tess_level_writes_le_two = 0;
+   unsigned tess_level_writes_other = 0;
+
+   /* Gather barriers and which values are written to tess level outputs. */
+   nir_foreach_function_impl(impl, nir) {
+      nir_foreach_block(block, impl) {
+         nir_foreach_instr(instr, block) {
+            if (instr->type != nir_instr_type_intrinsic)
+               continue;
+
+            nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+
+            if (is_tcs_output_barrier(intr)) {
+               /* Only gather barriers outside control flow. */
+               if (block->cf_node.parent->type == nir_cf_node_function)
+                  info->always_executes_barrier = true;
+               continue;
+            }
+
+            if (intr->intrinsic != nir_intrinsic_store_output)
+               continue;
+
+            unsigned location = nir_intrinsic_io_semantics(intr).location;
+            if (location != VARYING_SLOT_TESS_LEVEL_OUTER &&
+                location != VARYING_SLOT_TESS_LEVEL_INNER)
+               continue;
+
+            unsigned base_shift = get_tess_level_component(intr);
+            unsigned writemask = nir_intrinsic_write_mask(intr);
+
+            u_foreach_bit(i, writemask) {
+               nir_scalar scalar = nir_scalar_resolved(intr->src[0].ssa, i);
+               unsigned shift = base_shift + i;
+
+               if (nir_scalar_is_const(scalar)) {
+                  float f = nir_scalar_as_float(scalar);
+
+                  if (f <= 0 || isnan(f))
+                     tess_level_writes_le_zero |= BITFIELD_BIT(shift);
+                  else if (f <= 1)
+                     tess_level_writes_le_one |= BITFIELD_BIT(shift);
+                  else if (f <= 2)
+                     tess_level_writes_le_two |= BITFIELD_BIT(shift);
+                  else
+                     tess_level_writes_other |= BITFIELD_BIT(shift);
+               } else {
+                  /* TODO: This could use range analysis. */
+                  tess_level_writes_other |= BITFIELD_BIT(shift);
+               }
+            }
+         }
+      }
+   }
+
+   /* Determine which outer tess level components can discard patches.
+    * If the primitive type is unspecified, we have to assume the worst case.
+    */
+   unsigned min_outer, min_inner, max_outer, max_inner;
+   mesa_count_tess_level_components(prim == TESS_PRIMITIVE_UNSPECIFIED ?
+                                       TESS_PRIMITIVE_ISOLINES : prim,
+                                    &min_outer, &min_inner);
+   mesa_count_tess_level_components(prim, &max_outer, &max_inner);
+   const unsigned min_valid_outer_comp_mask = BITFIELD_RANGE(0, min_outer);
+   const unsigned max_valid_outer_comp_mask = BITFIELD_RANGE(0, max_outer);
+   const unsigned max_valid_inner_comp_mask = BITFIELD_RANGE(4, max_inner);
+
+   /* All tessellation levels are effectively 0 if the patch has at least one
+    * outer tess level component either in the [-inf, 0] range or equal to NaN,
+    * causing it to be discarded. Inner tess levels have no effect.
+    */
+   info->all_tess_levels_are_effectively_zero =
+      tess_level_writes_le_zero & ~tess_level_writes_le_one &
+      ~tess_level_writes_le_two & ~tess_level_writes_other &
+      min_valid_outer_comp_mask;
+
+   const unsigned tess_level_writes_any =
+      tess_level_writes_le_zero | tess_level_writes_le_one |
+      tess_level_writes_le_two | tess_level_writes_other;
+
+   const bool outer_is_gt_zero_le_one =
+      (tess_level_writes_le_one & ~tess_level_writes_le_zero &
+       ~tess_level_writes_le_two & ~tess_level_writes_other &
+       max_valid_outer_comp_mask) ==
+      (tess_level_writes_any & max_valid_outer_comp_mask);
+
+   /* Whether the inner tess levels are in the [-inf, 1] range. */
+   const bool inner_is_le_one =
+      ((tess_level_writes_le_zero | tess_level_writes_le_one) &
+       ~tess_level_writes_le_two & ~tess_level_writes_other &
+       max_valid_inner_comp_mask) ==
+      (tess_level_writes_any & max_valid_inner_comp_mask);
+
+   /* If the patch has tess level values set to 1 or equivalent numbers, it's
+    * not discarded, but different things happen depending on the spacing.
+    */
+   switch (spacing) {
+   case TESS_SPACING_EQUAL:
+   case TESS_SPACING_FRACTIONAL_ODD:
+   case TESS_SPACING_UNSPECIFIED:
+      /* The tessellator clamps all tess levels greater than 0 to 1.
+       * If all outer and inner tess levels are in the (0, 1] range, which is
+       * effectively 1, untessellated patches are drawn.
+       */
+      info->all_tess_levels_are_effectively_one = outer_is_gt_zero_le_one &&
+                                                  inner_is_le_one;
+      break;
+
+   case TESS_SPACING_FRACTIONAL_EVEN: {
+      /* The tessellator clamps all tess levels to 2 (both outer and inner)
+       * except outer tess level component 0 of isolines, which is clamped
+       * to 1. If all outer tess levels are in the (0, 2] or (0, 1] range
+       * (for outer[0] of isolines) and all inner tess levels are
+       * in the [-inf, 2] range, it's the same as writing 1 to all tess
+       * levels.
+       */
+      bool isolines_are_eff_one =
+         /* The (0, 1] range of outer[0]. */
+         (tess_level_writes_le_one & ~tess_level_writes_le_zero &
+          ~tess_level_writes_le_two & ~tess_level_writes_other & 0x1) ==
+         (tess_level_writes_any & 0x1) &&
+         /* The (0, 2] range of outer[1]. */
+         ((tess_level_writes_le_one | tess_level_writes_le_two) &
+          ~tess_level_writes_le_zero & ~tess_level_writes_other & 0x2) ==
+         (tess_level_writes_any & 0x2);
+
+      bool triquads_are_eff_one =
+         /* The (0, 2] outer range. */
+         ((tess_level_writes_le_one | tess_level_writes_le_two) &
+          ~tess_level_writes_le_zero & ~tess_level_writes_other &
+          max_valid_outer_comp_mask) ==
+         (tess_level_writes_any & max_valid_outer_comp_mask) &&
+         /* The [-inf, 2] inner range. */
+         ((tess_level_writes_le_zero | tess_level_writes_le_one |
+           tess_level_writes_le_two) & ~tess_level_writes_other &
+          max_valid_inner_comp_mask) ==
+         (tess_level_writes_any & max_valid_inner_comp_mask);
+
+      if (prim == TESS_PRIMITIVE_UNSPECIFIED) {
+         info->all_tess_levels_are_effectively_one = isolines_are_eff_one &&
+                                                     triquads_are_eff_one;
+      } else if (prim == TESS_PRIMITIVE_ISOLINES) {
+         info->all_tess_levels_are_effectively_one = isolines_are_eff_one;
+      } else {
+         info->all_tess_levels_are_effectively_one = triquads_are_eff_one;
+      }
+      break;
+   }
+   }
+
+   assert(!info->all_tess_levels_are_effectively_zero ||
+          !info->all_tess_levels_are_effectively_one);
+
+   info->discards_patches =
+      (tess_level_writes_le_zero & min_valid_outer_comp_mask) != 0;
+}