mesa: Move ATI_fragment_shader fog code emit to a NIR lowering pass.

Now it's implemented as a RMW of the FRAG_RESULT_COLOR output var (or
adjusting the store_output intrinsic's value for lowered i/o), which
should be reusable other places we might want to emit shader code for fog
(ARB_fp, fixed function fragment shaders).

Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/23111>

src/mesa/meson.build

@@ -358,6 +358,7 @@ files_libmesa = files(
   'state_tracker/st_nir.h',
   'state_tracker/st_nir_builtins.c',
   'state_tracker/st_nir_lower_builtin.c',
+  'state_tracker/st_nir_lower_fog.c',
   'state_tracker/st_nir_lower_tex_src_plane.c',
   'state_tracker/st_pbo.c',
   'state_tracker/st_pbo_compute.c',

src/mesa/state_tracker/st_atifs_to_nir.c

@@ -32,9 +32,6 @@
 #include "st_atifs_to_nir.h"
 #include "compiler/nir/nir_builder.h"
 
-#define FOG_PARAMS_UNIFORM (MAX_NUM_FRAGMENT_CONSTANTS_ATI + 0)
-#define FOG_COLOR_UNIFORM (MAX_NUM_FRAGMENT_CONSTANTS_ATI + 1)
-
 /**
  * Intermediate state used during shader translation.
  */
@@ -119,8 +116,7 @@ load_input(struct st_translate *t, gl_varying_slot slot)
 {
    if (!t->inputs[slot]) {
       nir_variable *var = nir_create_variable_with_location(t->b->shader, nir_var_shader_in, slot,
-                                                            slot == VARYING_SLOT_FOGC ?
-                                                            glsl_float_type() : glsl_vec4_type());
+                                                            glsl_vec4_type());
       var->data.interpolation = INTERP_MODE_NONE;
 
       t->inputs[slot] = nir_load_var(t->b, var);
@@ -418,8 +414,7 @@ compile_instruction(struct st_translate *t,
 }
 
 
-/* Creates the uniform variable referencing the ATI_fragment_shader constants
- * plus the optimized fog state.
+/* Creates the uniform variable referencing the ATI_fragment_shader constants.
  */
 static void
 st_atifs_setup_uniforms(struct st_translate *t, struct gl_program *program)
@@ -475,51 +470,8 @@ st_translate_atifs_program(struct ati_fragment_shader *atifs,
       }
    }
 
-   if (t->regs_written[atifs->NumPasses-1][0]) {
-      nir_ssa_def *color = t->temps[0];
-
-      if (key->fog) {
-         nir_ssa_def *fogc = load_input(t, VARYING_SLOT_FOGC);
-
-         nir_ssa_def *params = atifs_load_uniform(t, FOG_PARAMS_UNIFORM);
-
-         /* compute the 1 component fog factor f */
-         nir_ssa_def *f = NULL;
-         if (key->fog == FOG_LINEAR) {
-            f = nir_ffma(t->b, fogc,
-                         nir_channel(t->b, params, 0),
-                         nir_channel(t->b, params, 1));
-         } else if (key->fog == FOG_EXP) {
-            /* EXP formula: f = exp(-dens * z)
-             * with optimized parameters:
-             *    f = MUL(fogcoord, oparams.z); f= EX2(-f)
-             */
-            f = nir_fmul(t->b, fogc, nir_channel(t->b, params, 2));
-            f = nir_fexp2(t->b, nir_fneg(t->b, f));
-         } else if (key->fog == FOG_EXP2) {
-            /* EXP2 formula: f = exp(-(dens * z)^2)
-             * with optimized parameters:
-             *    f = MUL(fogcoord, oparams.w); f=MUL(f, f); f= EX2(-f)
-             */
-            f = nir_fmul(t->b, fogc, nir_channel(t->b, params, 3));
-            f = nir_fmul(t->b, f, f);
-            f = nir_fexp2(t->b, nir_fneg(t->b, f));
-         }
-         f = nir_fsat(t->b, f);
-
-         nir_ssa_def *fog_color = nir_flrp(t->b,
-                                           atifs_load_uniform(t, FOG_COLOR_UNIFORM),
-                                           color,
-                                           f);
-         color = nir_vec4(t->b,
-                          nir_channel(t->b, fog_color, 0),
-                          nir_channel(t->b, fog_color, 1),
-                          nir_channel(t->b, fog_color, 2),
-                          nir_channel(t->b, color, 3));
-      }
-
-      nir_store_var(t->b, t->fragcolor, color, 0xf);
-   }
+   if (t->regs_written[atifs->NumPasses-1][0])
+      nir_store_var(t->b, t->fragcolor, t->temps[0], 0xf);
 
    return b.shader;
 }
@@ -536,11 +488,6 @@ st_init_atifs_prog(struct gl_context *ctx, struct gl_program *prog)
 
    unsigned pass, i, r, optype, arg;
 
-   static const gl_state_index16 fog_params_state[STATE_LENGTH] =
-      {STATE_FOG_PARAMS_OPTIMIZED, 0, 0};
-   static const gl_state_index16 fog_color[STATE_LENGTH] =
-      {STATE_FOG_COLOR, 0, 0, 0};
-
    prog->info.inputs_read = 0;
    prog->info.outputs_written = BITFIELD64_BIT(FRAG_RESULT_COLOR);
    prog->SamplersUsed = 0;
@@ -587,17 +534,13 @@ st_init_atifs_prog(struct gl_context *ctx, struct gl_program *prog)
          }
       }
    }
+
    /* we may need fog */
    prog->info.inputs_read |= BITFIELD64_BIT(VARYING_SLOT_FOGC);
 
-   /* we always have the ATI_fs constants, and the fog params */
+   /* we always have the ATI_fs constants */
    for (i = 0; i < MAX_NUM_FRAGMENT_CONSTANTS_ATI; i++) {
       _mesa_add_parameter(prog->Parameters, PROGRAM_UNIFORM,
                           NULL, 4, GL_FLOAT, NULL, NULL, true);
    }
-   ASSERTED uint32_t ref;
-   ref = _mesa_add_state_reference(prog->Parameters, fog_params_state);
-   assert(ref == FOG_PARAMS_UNIFORM);
-   ref = _mesa_add_state_reference(prog->Parameters, fog_color);
-   assert(ref == FOG_COLOR_UNIFORM);
 }

src/mesa/state_tracker/st_nir.h

@@ -87,6 +87,9 @@ st_nir_state_variable_create(struct nir_shader *shader,
                              const struct glsl_type *type,
                              const gl_state_index16 state[STATE_LENGTH]);
 
+bool st_nir_lower_fog(struct nir_shader *s, enum gl_fog_mode fog_mode,
+                      struct gl_program_parameter_list *paramList);
+
 #ifdef __cplusplus
 }
 #endif

src/mesa/state_tracker/st_nir_lower_fog.c

@@ -0,0 +1,137 @@
+/*
+ * Copyright © 2023 Google LLC
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir_builder.h"
+#include "nir_builtin_builder.h"
+#include "st_nir.h"
+
+static nir_ssa_def *
+fog_result(nir_builder *b, nir_ssa_def *color, enum gl_fog_mode fog_mode, struct gl_program_parameter_list *paramList)
+{
+   nir_shader *s = b->shader;
+   nir_variable *fogc_var =
+      nir_create_variable_with_location(s, nir_var_shader_in, VARYING_SLOT_FOGC, glsl_float_type());
+   nir_ssa_def *fogc = nir_load_var(b, fogc_var);
+   s->info.inputs_read |= VARYING_BIT_FOGC;
+
+   static const gl_state_index16 fog_params_tokens[STATE_LENGTH] = {STATE_FOG_PARAMS_OPTIMIZED};
+   static const gl_state_index16 fog_color_tokens[STATE_LENGTH] = {STATE_FOG_COLOR};
+
+   nir_variable *fog_params_var = st_nir_state_variable_create(s, glsl_vec4_type(), fog_params_tokens);
+   fog_params_var->data.driver_location = _mesa_add_state_reference(paramList, fog_params_tokens);
+   nir_ssa_def *params = nir_load_var(b, fog_params_var);
+
+   nir_variable *fog_color_var = st_nir_state_variable_create(s, glsl_vec4_type(), fog_color_tokens);
+   fog_color_var->data.driver_location = _mesa_add_state_reference(paramList, fog_color_tokens);
+   nir_ssa_def *fog_color = nir_load_var(b, fog_color_var);
+
+   /* compute the 1 component fog factor f */
+   nir_ssa_def *f = NULL;
+   switch (fog_mode) {
+   case FOG_LINEAR:
+      /* f = (end - z) / (end - start)
+       *
+       * gl_MesaFogParamsOptimized gives us (-1 / (end - start)) and
+       * (end / (end - start)) so we can generate a single MAD.
+       */
+      f = nir_fmad(b, fogc,
+                   nir_channel(b, params, 0),
+                   nir_channel(b, params, 1));
+      break;
+  case FOG_EXP:
+      /* f = e^(-(density * fogcoord))
+       *
+       * gl_MesaFogParamsOptimized gives us density/ln(2) so we can
+       * use EXP2 which is generally the native instruction without
+       * having to do any further math on the fog density uniform.
+       */
+      f = nir_fmul(b, fogc, nir_channel(b, params, 2));
+      f = nir_fexp2(b, nir_fneg(b, f));
+      break;
+  case FOG_EXP2:
+      /* f = e^(-(density * fogcoord)^2)
+       *
+       * gl_MesaFogParamsOptimized gives us density/sqrt(ln(2)) so we
+       * can do this like FOG_EXP but with a squaring after the
+       * multiply by density.
+       */
+      f = nir_fmul(b, fogc, nir_channel(b, params, 3));
+      f = nir_fmul(b, f, f);
+      f = nir_fexp2(b, nir_fneg(b, f));
+      break;
+   default:
+      unreachable("unsupported fog mode");
+   }
+   f = nir_fsat(b, f);
+
+   /* Not using flrp because we may end up lowering fog after driver lowering
+    * that meant to remove all lrps.
+    */
+   return nir_fmad(b, color, f, nir_fmul(b, fog_color, nir_fsub_imm(b, 1.0, f)));
+}
+
+struct lower_fog_state {
+   enum gl_fog_mode fog_mode;
+   struct gl_program_parameter_list *paramList;
+};
+
+static bool
+st_nir_lower_fog_instr(nir_builder *b, nir_instr *instr, void *_state)
+{
+   const struct lower_fog_state *state = _state;
+
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+   if (intr->intrinsic != nir_intrinsic_store_output)
+      return false;
+
+   int loc = nir_intrinsic_io_semantics(intr).location;
+   if (loc != FRAG_RESULT_COLOR)
+      return false;
+
+   b->cursor = nir_before_instr(instr);
+
+   nir_ssa_def *color = nir_ssa_for_src(b, intr->src[0], intr->num_components);
+   color = nir_resize_vector(b, color, 4);
+
+   nir_ssa_def *fog = fog_result(b, color, state->fog_mode, state->paramList);
+
+   /* retain the non-fog-blended alpha value for color */
+   color = nir_vector_insert_imm(b, fog, nir_channel(b, color, 3), 3);
+
+   nir_instr_rewrite_src_ssa(instr, &intr->src[0], nir_resize_vector(b, color, intr->num_components));
+
+   return true;
+}
+
+bool
+st_nir_lower_fog(nir_shader *s, enum gl_fog_mode fog_mode, struct gl_program_parameter_list *paramList)
+{
+   if (s->info.io_lowered) {
+      struct lower_fog_state state = {
+         .fog_mode = fog_mode,
+         .paramList = paramList,
+      };
+      nir_shader_instructions_pass(s, st_nir_lower_fog_instr,
+                                   nir_metadata_block_index |
+                                   nir_metadata_dominance,
+                                   &state);
+   } else {
+      nir_variable *color_var = nir_find_variable_with_location(s, nir_var_shader_out, FRAG_RESULT_COLOR);
+
+      nir_builder b;
+      nir_builder_init(&b, nir_shader_get_entrypoint(s));
+      b.cursor = nir_after_block(nir_impl_last_block(b.impl));
+
+      nir_ssa_def *color = nir_load_var(&b, color_var);
+      color = fog_result(&b, color, fog_mode, paramList);
+      nir_store_var(&b, color_var, color, 0x7);
+
+      nir_metadata_preserve(b.impl, nir_metadata_block_index |
+                                    nir_metadata_dominance);   }
+
+   return true;
+}

src/mesa/state_tracker/st_program.c

@@ -53,6 +53,7 @@
 #include "tgsi/tgsi_dump.h"
 #include "tgsi/tgsi_parse.h"
 #include "tgsi/tgsi_ureg.h"
+#include "nir_builder.h"
 #include "nir/nir_to_tgsi.h"
 
 #include "util/u_memory.h"
@@ -926,6 +927,17 @@ st_create_fp_variant(struct st_context *st,
 
    bool finalize = false;
 
+   if (fp->ati_fs) {
+      if (key->fog) {
+         NIR_PASS_V(state.ir.nir, st_nir_lower_fog, key->fog, fp->Parameters);
+         NIR_PASS_V(state.ir.nir, nir_lower_io_to_temporaries,
+            nir_shader_get_entrypoint(state.ir.nir),
+            true, false);
+         nir_lower_global_vars_to_local(state.ir.nir);
+      }
+      finalize = true;
+   }
+
    if (key->clamp_color) {
       NIR_PASS_V(state.ir.nir, nir_lower_clamp_color_outputs);
       finalize = true;