ir3: make ir3_const_state less error-prone to use

ir3_const_state is shared between the binning and non-binning variants.
The non-binning variant is compiled first and sets up ir3_const_state
after which the binning variant is not supposed to modify it anymore. If
it would, things may go haywire since the layout of the constant state
will change after the non-binning variant already finished compiling.

Currently, the ir3_const_state() accessor takes care of the sharing
(i.e., it returns the non-binning const state for the binning variant)
but nothing would be prevent the binning variant from accidentally
modifying the state. This is handled by restraint from its users.

This commit tries to make it more difficult to accidentally modify the
const state by the binning shader by making the following changes:
- ir3_const_state(): the same logic as before but now returns a const
  pointer to prevent the binning variant from (accidentally) modifying
  the const state.
- ir3_const_state_mut(): returns a non-const pointer but asserts that it
  is not called by the binning variant.

As a corollary ir3_get_driver_ubo() also had to be split in two variants
(const and non-const) as it is called with a pointer to one of the
fields of ir3_const_state.

Signed-off-by: Job Noorman <jnoorman@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30256>

src/freedreno/ir3/ir3.c

@@ -69,7 +69,7 @@ ir3_destroy(struct ir3 *shader)
 
 static bool
 is_shared_consts(struct ir3_compiler *compiler,
-                 struct ir3_const_state *const_state,
+                 const struct ir3_const_state *const_state,
                  struct ir3_register *reg)
 {
    if (const_state->push_consts_type == IR3_PUSH_CONSTS_SHARED &&

src/freedreno/ir3/ir3_compiler_nir.c

@@ -193,7 +193,7 @@ create_driver_param(struct ir3_context *ctx, enum ir3_driver_param dp)
 {
    /* first four vec4 sysval's reserved for UBOs: */
    /* NOTE: dp is in scalar, but there can be >4 dp components: */
-   struct ir3_const_state *const_state = ir3_const_state(ctx->so);
+   const struct ir3_const_state *const_state = ir3_const_state(ctx->so);
    unsigned n = const_state->offsets.driver_param;
    unsigned r = regid(n + dp / 4, dp % 4);
    return create_uniform(ctx->block, r);
@@ -205,7 +205,7 @@ create_driver_param_indirect(struct ir3_context *ctx, enum ir3_driver_param dp,
 {
    /* first four vec4 sysval's reserved for UBOs: */
    /* NOTE: dp is in scalar, but there can be >4 dp components: */
-   struct ir3_const_state *const_state = ir3_const_state(ctx->so);
+   const struct ir3_const_state *const_state = ir3_const_state(ctx->so);
    unsigned n = const_state->offsets.driver_param;
    return create_uniform_indirect(ctx->block, n * 4 + dp, TYPE_U32, address);
 }

src/freedreno/ir3/ir3_nir.c

@@ -31,6 +31,15 @@
 #include "ir3_nir.h"
 #include "ir3_shader.h"
 
+nir_def *
+ir3_get_shared_driver_ubo(nir_builder *b, const struct ir3_driver_ubo *ubo)
+{
+   assert(ubo->idx > 0);
+
+   /* Binning shader shared ir3_driver_ubo definitions but not shader info */
+   b->shader->info.num_ubos = MAX2(b->shader->info.num_ubos, ubo->idx + 1);
+   return nir_imm_int(b, ubo->idx);
+}
 
 nir_def *
 ir3_get_driver_ubo(nir_builder *b, struct ir3_driver_ubo *ubo)
@@ -42,13 +51,18 @@ ir3_get_driver_ubo(nir_builder *b, struct ir3_driver_ubo *ubo)
       if (b->shader->info.num_ubos == 0)
          b->shader->info.num_ubos++;
       ubo->idx = b->shader->info.num_ubos++;
-   } else {
-      assert(ubo->idx != 0);
-      /* Binning shader shared ir3_driver_ubo definitions but not shader info */
-      b->shader->info.num_ubos = MAX2(b->shader->info.num_ubos, ubo->idx + 1);
+      return nir_imm_int(b, ubo->idx);
    }
 
-   return nir_imm_int(b, ubo->idx);
+   return ir3_get_shared_driver_ubo(b, ubo);
+}
+
+nir_def *
+ir3_get_driver_consts_ubo(nir_builder *b, struct ir3_shader_variant *v)
+{
+   if (v->binning_pass)
+      return ir3_get_shared_driver_ubo(b, &ir3_const_state(v)->consts_ubo);
+   return ir3_get_driver_ubo(b, &ir3_const_state_mut(v)->consts_ubo);
 }
 
 nir_def *
@@ -1068,7 +1082,7 @@ ir3_nir_lower_variant(struct ir3_shader_variant *so, nir_shader *s)
     * passes:
     */
    if (!so->binning_pass)
-      ir3_setup_const_state(s, so, ir3_const_state(so));
+      ir3_setup_const_state(s, so, ir3_const_state_mut(so));
 }
 
 bool

src/freedreno/ir3/ir3_nir.h

@@ -96,7 +96,11 @@ nir_def *ir3_nir_try_propagate_bit_shift(nir_builder *b,
 
 bool ir3_nir_opt_subgroups(nir_shader *nir, struct ir3_shader_variant *v);
 
+nir_def *ir3_get_shared_driver_ubo(nir_builder *b,
+                                   const struct ir3_driver_ubo *ubo);
 nir_def *ir3_get_driver_ubo(nir_builder *b, struct ir3_driver_ubo *ubo);
+nir_def *ir3_get_driver_consts_ubo(nir_builder *b,
+                                   struct ir3_shader_variant *v);
 nir_def *ir3_load_driver_ubo(nir_builder *b, unsigned components,
                              struct ir3_driver_ubo *ubo,
                              unsigned offset);

src/freedreno/ir3/ir3_nir_analyze_ubo_ranges.c

@@ -564,7 +564,7 @@ assign_offsets(struct ir3_ubo_analysis_state *state, unsigned start,
 bool
 ir3_nir_lower_const_global_loads(nir_shader *nir, struct ir3_shader_variant *v)
 {
-   struct ir3_const_state *const_state = ir3_const_state(v);
+   const struct ir3_const_state *const_state = ir3_const_state(v);
    struct ir3_compiler *compiler = v->compiler;
 
    if (ir3_shader_debug & IR3_DBG_NOUBOOPT)
@@ -630,7 +630,7 @@ ir3_nir_lower_const_global_loads(nir_shader *nir, struct ir3_shader_variant *v)
    }
 
    if (!v->binning_pass)
-      const_state->global_size = DIV_ROUND_UP(state.size, 16);
+      ir3_const_state_mut(v)->global_size = DIV_ROUND_UP(state.size, 16);
 
    return progress;
 }
@@ -638,7 +638,7 @@ ir3_nir_lower_const_global_loads(nir_shader *nir, struct ir3_shader_variant *v)
 void
 ir3_nir_analyze_ubo_ranges(nir_shader *nir, struct ir3_shader_variant *v)
 {
-   struct ir3_const_state *const_state = ir3_const_state(v);
+   struct ir3_const_state *const_state = ir3_const_state_mut(v);
    struct ir3_ubo_analysis_state *state = &const_state->ubo_state;
    struct ir3_compiler *compiler = v->compiler;
 
@@ -807,7 +807,7 @@ ir3_nir_fixup_load_const_ir3(nir_shader *nir)
 static nir_def *
 ir3_nir_lower_load_const_instr(nir_builder *b, nir_instr *in_instr, void *data)
 {
-   struct ir3_const_state *const_state = data;
+   struct ir3_shader_variant *v = data;
    nir_intrinsic_instr *instr = nir_instr_as_intrinsic(in_instr);
 
    unsigned num_components = instr->num_components;
@@ -823,7 +823,7 @@ ir3_nir_lower_load_const_instr(nir_builder *b, nir_instr *in_instr, void *data)
       bit_size = 32;
    }
    unsigned base = nir_intrinsic_base(instr);
-   nir_def *index = ir3_get_driver_ubo(b, &const_state->consts_ubo);
+   nir_def *index = ir3_get_driver_consts_ubo(b, v);
    nir_def *offset =
       nir_iadd_imm(b, instr->src[0].ssa, base);
 
@@ -855,11 +855,9 @@ ir3_lower_load_const_filter(const nir_instr *instr, const void *data)
 bool
 ir3_nir_lower_load_constant(nir_shader *nir, struct ir3_shader_variant *v)
 {
-   struct ir3_const_state *const_state = ir3_const_state(v);
-
    bool progress = nir_shader_lower_instructions(
       nir, ir3_lower_load_const_filter, ir3_nir_lower_load_const_instr,
-      const_state);
+      v);
 
    if (progress) {
       struct ir3_compiler *compiler = v->compiler;

src/freedreno/ir3/ir3_nir_lower_driver_params_to_ubo.c

@@ -12,7 +12,8 @@
 static bool
 lower_driver_param_to_ubo(nir_builder *b, nir_intrinsic_instr *intr, void *in)
 {
-   struct ir3_const_state *const_state = in;
+   struct ir3_shader_variant *v = in;
+   struct ir3_const_state *const_state = ir3_const_state_mut(v);
 
    unsigned components = nir_intrinsic_dest_components(intr);
 
@@ -80,7 +81,7 @@ ir3_nir_lower_driver_params_to_ubo(nir_shader *nir,
 {
    bool result = nir_shader_intrinsics_pass(
       nir, lower_driver_param_to_ubo,
-      nir_metadata_control_flow, ir3_const_state(v));
+      nir_metadata_control_flow, v);
 
    return result;
 }

src/freedreno/ir3/ir3_nir_opt_preamble.c

@@ -299,10 +299,9 @@ set_speculate(nir_builder *b, nir_intrinsic_instr *intr, UNUSED void *_)
 bool
 ir3_nir_opt_preamble(nir_shader *nir, struct ir3_shader_variant *v)
 {
-   struct ir3_const_state *const_state = ir3_const_state(v);
-
    unsigned max_size;
    if (v->binning_pass) {
+      const struct ir3_const_state *const_state = ir3_const_state(v);
       max_size = const_state->preamble_size * 4;
    } else {
       struct ir3_const_state worst_case_const_state = {};
@@ -330,8 +329,10 @@ ir3_nir_opt_preamble(nir_shader *nir, struct ir3_shader_variant *v)
    unsigned size = 0;
    progress |= nir_opt_preamble(nir, &options, &size);
 
-   if (!v->binning_pass)
+   if (!v->binning_pass) {
+      struct ir3_const_state *const_state = ir3_const_state_mut(v);
       const_state->preamble_size = DIV_ROUND_UP(size, 4);
+   }
 
    return progress;
 }
@@ -613,7 +614,7 @@ get_preamble_offset(nir_def *def)
 bool
 ir3_nir_opt_prefetch_descriptors(nir_shader *nir, struct ir3_shader_variant *v)
 {
-   struct ir3_const_state *const_state = ir3_const_state(v);
+   const struct ir3_const_state *const_state = ir3_const_state(v);
 
    nir_function_impl *main = nir_shader_get_entrypoint(nir);
    struct set *instr_set = nir_instr_set_create(NULL);

src/freedreno/ir3/ir3_parser.y

@@ -217,7 +217,7 @@ static void fixup_cat5_s2en(void)
 
 static void add_const(unsigned reg, unsigned c0, unsigned c1, unsigned c2, unsigned c3)
 {
-	struct ir3_const_state *const_state = ir3_const_state(variant);
+	struct ir3_const_state *const_state = ir3_const_state_mut(variant);
 	assert((reg & 0x7) == 0);
 	int idx = reg >> (1 + 2); /* low bit is half vs full, next two bits are swiz */
 	if (idx * 4 + 4 > const_state->immediates_size) {

src/freedreno/ir3/ir3_shader.c

@@ -65,7 +65,7 @@ ir3_const_find_imm(struct ir3_shader_variant *v, uint32_t imm)
 uint16_t
 ir3_const_add_imm(struct ir3_shader_variant *v, uint32_t imm)
 {
-   struct ir3_const_state *const_state = ir3_const_state(v);
+   struct ir3_const_state *const_state = ir3_const_state_mut(v);
 
    /* Reallocate for 4 more elements whenever it's necessary.  Note that ir3
     * printing relies on having groups of 4 dwords, so we fill the unused

src/freedreno/ir3/ir3_shader.h

@@ -937,7 +937,7 @@ struct ir3_shader {
  * emit, for both binning and draw pass (a6xx+), the binning pass re-uses it's
  * corresponding draw pass shaders const_state.
  */
-static inline struct ir3_const_state *
+static inline const struct ir3_const_state *
 ir3_const_state(const struct ir3_shader_variant *v)
 {
    if (v->binning_pass)
@@ -945,6 +945,13 @@ ir3_const_state(const struct ir3_shader_variant *v)
    return v->const_state;
 }
 
+static inline struct ir3_const_state *
+ir3_const_state_mut(const struct ir3_shader_variant *v)
+{
+   assert(!v->binning_pass);
+   return v->const_state;
+}
+
 static inline unsigned
 _ir3_max_const(const struct ir3_shader_variant *v, bool safe_constlen)
 {

src/gallium/drivers/freedreno/a6xx/fd6_const.cc

@@ -234,8 +234,8 @@ fd6_user_consts_cmdstream_size(const struct ir3_shader_variant *v)
    if (!v)
       return 0;
 
-   struct ir3_const_state *const_state = ir3_const_state(v);
-   struct ir3_ubo_analysis_state *ubo_state = &const_state->ubo_state;
+   const struct ir3_const_state *const_state = ir3_const_state(v);
+   const struct ir3_ubo_analysis_state *ubo_state = &const_state->ubo_state;
    unsigned packets, size;
 
    /* pre-calculate size required for userconst stateobj: */

src/gallium/drivers/freedreno/ir3/ir3_const.h

@@ -93,7 +93,7 @@ ring_wfi(struct fd_batch *batch, struct fd_ringbuffer *ring) assert_dt
  * Returns size in dwords.
  */
 static inline void
-ir3_user_consts_size(struct ir3_ubo_analysis_state *state, unsigned *packets,
+ir3_user_consts_size(const struct ir3_ubo_analysis_state *state, unsigned *packets,
                      unsigned *size)
 {
    *packets = *size = 0;