xref: /aosp_15_r20/external/mesa3d/src/mesa/state_tracker/st_nir_lower_fog.c (revision 6104692788411f58d303aa86923a9ff6ecaded22)
1 /*
2  * Copyright © 2023 Google LLC
3  * SPDX-License-Identifier: MIT
4  */
5 
6 #include "nir_builder.h"
7 #include "nir_builtin_builder.h"
8 #include "st_nir.h"
9 
10 static nir_def *
fog_result(nir_builder * b,nir_def * color,enum gl_fog_mode fog_mode,struct gl_program_parameter_list * paramList)11 fog_result(nir_builder *b, nir_def *color, enum gl_fog_mode fog_mode, struct gl_program_parameter_list *paramList)
12 {
13    nir_shader *s = b->shader;
14    nir_variable *fogc_var =
15       nir_create_variable_with_location(s, nir_var_shader_in, VARYING_SLOT_FOGC, glsl_float_type());
16    nir_def *fogc = nir_load_var(b, fogc_var);
17    s->info.inputs_read |= VARYING_BIT_FOGC;
18 
19    static const gl_state_index16 fog_params_tokens[STATE_LENGTH] = {STATE_FOG_PARAMS_OPTIMIZED};
20    static const gl_state_index16 fog_color_tokens[STATE_LENGTH] = {STATE_FOG_COLOR};
21 
22    nir_variable *fog_params_var = st_nir_state_variable_create(s, glsl_vec4_type(), fog_params_tokens);
23    fog_params_var->data.driver_location = _mesa_add_state_reference(paramList, fog_params_tokens);
24    nir_def *params = nir_load_var(b, fog_params_var);
25 
26    nir_variable *fog_color_var = st_nir_state_variable_create(s, glsl_vec4_type(), fog_color_tokens);
27    fog_color_var->data.driver_location = _mesa_add_state_reference(paramList, fog_color_tokens);
28    nir_def *fog_color = nir_load_var(b, fog_color_var);
29 
30    /* compute the 1 component fog factor f */
31    nir_def *f = NULL;
32    switch (fog_mode) {
33    case FOG_LINEAR:
34       /* f = (end - z) / (end - start)
35        *
36        * gl_MesaFogParamsOptimized gives us (-1 / (end - start)) and
37        * (end / (end - start)) so we can generate a single MAD.
38        */
39       f = nir_fmad(b, fogc,
40                    nir_channel(b, params, 0),
41                    nir_channel(b, params, 1));
42       break;
43   case FOG_EXP:
44       /* f = e^(-(density * fogcoord))
45        *
46        * gl_MesaFogParamsOptimized gives us density/ln(2) so we can
47        * use EXP2 which is generally the native instruction without
48        * having to do any further math on the fog density uniform.
49        */
50       f = nir_fmul(b, fogc, nir_channel(b, params, 2));
51       f = nir_fexp2(b, nir_fneg(b, f));
52       break;
53   case FOG_EXP2:
54       /* f = e^(-(density * fogcoord)^2)
55        *
56        * gl_MesaFogParamsOptimized gives us density/sqrt(ln(2)) so we
57        * can do this like FOG_EXP but with a squaring after the
58        * multiply by density.
59        */
60       f = nir_fmul(b, fogc, nir_channel(b, params, 3));
61       f = nir_fmul(b, f, f);
62       f = nir_fexp2(b, nir_fneg(b, f));
63       break;
64    default:
65       unreachable("unsupported fog mode");
66    }
67    f = nir_fsat(b, f);
68 
69    /* Not using flrp because we may end up lowering fog after driver lowering
70     * that meant to remove all lrps.
71     */
72    return nir_fmad(b, color, f, nir_fmul(b, fog_color, nir_fsub_imm(b, 1.0, f)));
73 }
74 
75 struct lower_fog_state {
76    enum gl_fog_mode fog_mode;
77    struct gl_program_parameter_list *paramList;
78 };
79 
80 static bool
st_nir_lower_fog_instr(nir_builder * b,nir_instr * instr,void * _state)81 st_nir_lower_fog_instr(nir_builder *b, nir_instr *instr, void *_state)
82 {
83    const struct lower_fog_state *state = _state;
84 
85    if (instr->type != nir_instr_type_intrinsic)
86       return false;
87    nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
88    if (intr->intrinsic != nir_intrinsic_store_output)
89       return false;
90 
91    int loc = nir_intrinsic_io_semantics(intr).location;
92    if (loc != FRAG_RESULT_COLOR && loc != FRAG_RESULT_DATA0)
93       return false;
94 
95    b->cursor = nir_before_instr(instr);
96 
97    nir_def *color = intr->src[0].ssa;
98    color = nir_resize_vector(b, color, 4);
99 
100    nir_def *fog = fog_result(b, color, state->fog_mode, state->paramList);
101 
102    /* retain the non-fog-blended alpha value for color */
103    color = nir_vector_insert_imm(b, fog, nir_channel(b, color, 3), 3);
104 
105    nir_src_rewrite(&intr->src[0],
106                    nir_resize_vector(b, color, intr->num_components));
107 
108    return true;
109 }
110 
111 bool
st_nir_lower_fog(nir_shader * s,enum gl_fog_mode fog_mode,struct gl_program_parameter_list * paramList)112 st_nir_lower_fog(nir_shader *s, enum gl_fog_mode fog_mode, struct gl_program_parameter_list *paramList)
113 {
114    if (s->info.io_lowered) {
115       struct lower_fog_state state = {
116          .fog_mode = fog_mode,
117          .paramList = paramList,
118       };
119       nir_shader_instructions_pass(s, st_nir_lower_fog_instr,
120                                    nir_metadata_control_flow,
121                                    &state);
122    } else {
123       nir_variable *color_var = nir_find_variable_with_location(s, nir_var_shader_out, FRAG_RESULT_COLOR);
124       if (!color_var) {
125          color_var = nir_find_variable_with_location(s, nir_var_shader_out, FRAG_RESULT_DATA0);
126          /* Fog result would be undefined if we had no output color (in ARB_fragment_program) */
127          if (!color_var)
128             return false;
129       }
130 
131       nir_function_impl *impl = nir_shader_get_entrypoint(s);
132       nir_builder b = nir_builder_at(nir_after_impl(impl));
133 
134       /* Note: while ARB_fragment_program plus ARB_draw_buffers allows an array
135        * of result colors, prog_to_nir generates separate vars per slot so we
136        * don't have to handle that.  Fog only applies to the first color result.
137        */
138       assert(!glsl_type_is_array(color_var->type));
139 
140       nir_def *color = nir_load_var(&b, color_var);
141       color = fog_result(&b, color, fog_mode, paramList);
142       nir_store_var(&b, color_var, color, 0x7);
143 
144       nir_metadata_preserve(b.impl, nir_metadata_control_flow);   }
145 
146    return true;
147 }
148