nir: switch to a normal sampler for ARB program with not depth textures

It is undefined behavior when an ARB assembly or shadow2d GLSL func
uses SHADOW2D target with a texture in not depth format.
In this case AMD and NVIDIA automatically replaces SHADOW sampler
with a normal sampler and some games like Penumbra Overture which abuses
this UB works fine but breaks with mesa.

Replace the shadow sampler with a normal one here by recompiling
the ARB program variant

Closes: https://gitlab.freedesktop.org/mesa/mesa/-/issues/8425
Reviewed-by: Emma Anholt <emma@anholt.net>
Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Signed-off-by: Illia Polishchuk <illia.a.polishchuk@globallogic.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/22147>

src/compiler/nir/meson.build

@@ -276,6 +276,7 @@ files_libnir = files(
   'nir_range_analysis.c',
   'nir_range_analysis.h',
   'nir_remove_dead_variables.c',
+  'nir_remove_tex_shadow.c',
   'nir_repair_ssa.c',
   'nir_scale_fdiv.c',
   'nir_schedule.c',

src/compiler/nir/nir.h

@@ -5962,6 +5962,9 @@ bool nir_lower_poly_line_smooth(nir_shader *shader, unsigned num_smooth_aa_sampl
 
 bool nir_mod_analysis(nir_ssa_scalar val, nir_alu_type val_type, unsigned div, unsigned *mod);
 
+bool
+nir_remove_tex_shadow(nir_shader *shader, unsigned textures_bitmask);
+
 #include "nir_inline_helpers.h"
 
 #ifdef __cplusplus

src/compiler/nir/nir_remove_tex_shadow.c

@@ -0,0 +1,59 @@
+/*
+ * Copyright © 2023 Intel Corporation
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice (including the next
+ * paragraph) shall be included in all copies or substantial portions of the
+ * Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include "nir.h"
+#include "nir_builder.h"
+
+static bool
+remove_tex_shadow(struct nir_builder *b, nir_instr *instr, void *data)
+{
+   if (instr->type != nir_instr_type_tex)
+      return false;
+
+   nir_tex_instr *tex = nir_instr_as_tex(instr);
+
+   if (!tex->is_shadow)
+      return false;
+
+   unsigned *textures_bitmask = data;
+
+   if (BITFIELD_BIT(tex->texture_index) & ~*textures_bitmask)
+      return false;
+
+   int index = nir_tex_instr_src_index(tex, nir_tex_src_comparator);
+
+   if (index != -1) {
+      tex->is_shadow = false;
+      nir_tex_instr_remove_src(tex, index);
+      return true;
+   }
+
+   return false;
+}
+
+bool
+nir_remove_tex_shadow(nir_shader *shader, unsigned textures_bitmask)
+{
+   return nir_shader_instructions_pass(shader, remove_tex_shadow,
+                                       nir_metadata_none, &textures_bitmask);
+}

src/mesa/state_tracker/st_atom_shader.c

@@ -40,6 +40,7 @@
 #include "main/framebuffer.h"
 #include "main/state.h"
 #include "main/texobj.h"
+#include "main/teximage.h"
 #include "main/texstate.h"
 #include "program/program.h"
 
@@ -121,7 +122,8 @@ st_update_fp( struct st_context *st )
 
    if (st->shader_has_one_variant[MESA_SHADER_FRAGMENT] &&
        !fp->ati_fs && /* ATI_fragment_shader always has multiple variants */
-       !fp->ExternalSamplersUsed /* external samplers need variants */) {
+       !fp->ExternalSamplersUsed && /* external samplers need variants */
+       !(!fp->shader_program && fp->ShadowSamplers)) {
       shader = fp->variants->driver_shader;
    } else {
       struct st_fp_variant_key key;
@@ -163,6 +165,18 @@ st_update_fp( struct st_context *st )
          }
       }
 
+      if (!fp->shader_program && fp->ShadowSamplers) {
+         u_foreach_bit(i, fp->ShadowSamplers) {
+            struct gl_texture_object *tex_obj =
+                _mesa_get_tex_unit(st->ctx, fp->SamplerUnits[i])->_Current;
+            GLenum16 baseFormat = _mesa_base_tex_image(tex_obj)->_BaseFormat;
+
+            if (baseFormat == GL_DEPTH_COMPONENT ||
+                baseFormat == GL_DEPTH_STENCIL)
+               key.depth_textures |= BITFIELD_BIT(i);
+         }
+      }
+
       key.external = st_get_external_sampler_key(st, fp);
       update_gl_clamp(st, st->ctx->FragmentProgram._Current, key.gl_clamp);
 

src/mesa/state_tracker/st_context.c

@@ -177,7 +177,8 @@ st_invalidate_state(struct gl_context *ctx)
       if (ctx->FragmentProgram._Current) {
          struct gl_program *fp = ctx->FragmentProgram._Current;
 
-         if (fp->ExternalSamplersUsed || fp->ati_fs)
+         if (fp->ExternalSamplersUsed || fp->ati_fs ||
+            (!fp->shader_program && fp->ShadowSamplers))
             ctx->NewDriverState |= ST_NEW_FS_STATE;
       }
    }

src/mesa/state_tracker/st_program.c

@@ -1063,6 +1063,18 @@ st_create_fp_variant(struct st_context *st,
       finalize = true;
    }
 
+   /* It is undefined behavior when an ARB assembly uses SHADOW2D target
+    * with a texture in not depth format. In this case NVIDIA automatically
+    * replaces SHADOW sampler with a normal sampler and some games like
+    * Penumbra Overture which abuses this UB (issues/8425) works fine but
+    * breaks with mesa. Replace the shadow sampler with a normal one here
+    */
+   if (!fp->shader_program && ~key->depth_textures & fp->ShadowSamplers) {
+      NIR_PASS_V(state.ir.nir, nir_remove_tex_shadow,
+                 ~key->depth_textures & fp->ShadowSamplers);
+      finalize = true;
+   }
+
    if (finalize || !st->allow_st_finalize_nir_twice) {
       /* Some of the lowering above may have introduced new varyings */
       nir_shader_gather_info(state.ir.nir,
@@ -1104,7 +1116,7 @@ st_get_fp_variant(struct st_context *st,
 
       if (fp->variants != NULL) {
          _mesa_perf_debug(st->ctx, MESA_DEBUG_SEVERITY_MEDIUM,
-                          "Compiling fragment shader variant (%s%s%s%s%s%s%s%s%s%s%s%s)",
+                          "Compiling fragment shader variant (%s%s%s%s%s%s%s%s%s%s%s%s%s%d)",
                           key->bitmap ? "bitmap," : "",
                           key->drawpixels ? "drawpixels," : "",
                           key->scaleAndBias ? "scale_bias," : "",
@@ -1117,7 +1129,8 @@ st_get_fp_variant(struct st_context *st,
                           key->lower_alpha_func ? "alpha_compare," : "",
                           /* skipped ATI_fs targets */
                           fp->ExternalSamplersUsed ? "external?," : "",
-                          key->gl_clamp[0] || key->gl_clamp[1] || key->gl_clamp[2] ? "GL_CLAMP," : "");
+                          key->gl_clamp[0] || key->gl_clamp[1] || key->gl_clamp[2] ? "GL_CLAMP," : "",
+                          "depth_textures=", key->depth_textures);
       }
 
       fpv = st_create_fp_variant(st, fp, key);
@@ -1307,6 +1320,10 @@ st_precompile_shader_variant(struct st_context *st,
          for (int i = 0; i < ARRAY_SIZE(key.texture_index); i++)
             key.texture_index[i] = TEXTURE_2D_INDEX;
       }
+
+      /* Shadow samplers require texture in depth format */
+      key.depth_textures = prog->ShadowSamplers;
+
       st_get_fp_variant(st, prog, &key);
       break;
    }

src/mesa/state_tracker/st_program.h

@@ -186,6 +186,9 @@ struct st_fp_variant_key
 
    /* bitmask of sampler units; PIPE_CAP_GL_CLAMP */
    uint32_t gl_clamp[3];
+
+   /* bitmask of texture depth units; */
+   GLbitfield depth_textures;
 };
 
 /**