agx: Lower sampler LOD bias

G13 does not support sampler descriptor LOD biasing, so this needs to be lowered
to shader code for APIs that require this functionality. Add an option to do
this lowering while doing our other backend texture lowerings. This generates
lod_bias_agx texture instructions which the driver is expected to lower
according to its binding model.

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

src/asahi/compiler/agx_compile.c

@@ -2189,7 +2189,7 @@ agx_compile_function_nir(nir_shader *nir, nir_function_impl *impl,
  * lowered here to avoid duplicate work with shader variants.
  */
 void
-agx_preprocess_nir(nir_shader *nir)
+agx_preprocess_nir(nir_shader *nir, bool support_lod_bias)
 {
    NIR_PASS_V(nir, nir_lower_vars_to_ssa);
 
@@ -2233,7 +2233,7 @@ agx_preprocess_nir(nir_shader *nir)
 
    /* Clean up deref gunk after lowering I/O */
    NIR_PASS_V(nir, nir_opt_dce);
-   NIR_PASS_V(nir, agx_nir_lower_texture);
+   NIR_PASS_V(nir, agx_nir_lower_texture, support_lod_bias);
 
    nir_lower_idiv_options idiv_options = {
       .allow_fp16 = true,

src/asahi/compiler/agx_compile.h

@@ -173,7 +173,7 @@ struct agx_shader_key {
    };
 };
 
-void agx_preprocess_nir(nir_shader *nir);
+void agx_preprocess_nir(nir_shader *nir, bool support_lod_bias);
 
 void agx_compile_shader_nir(nir_shader *nir, struct agx_shader_key *key,
                             struct util_debug_callback *debug,

src/asahi/compiler/agx_compiler.h

@@ -826,7 +826,7 @@ void agx_compute_liveness(agx_context *ctx);
 void agx_liveness_ins_update(BITSET_WORD *live, agx_instr *I);
 
 bool agx_nir_lower_zs_emit(nir_shader *s);
-bool agx_nir_lower_texture(nir_shader *s);
+bool agx_nir_lower_texture(nir_shader *s, bool support_lod_bias);
 bool agx_nir_opt_preamble(nir_shader *s, unsigned *preamble_size);
 bool agx_nir_lower_load_mask(nir_shader *shader);
 bool agx_nir_lower_address(nir_shader *shader);

src/asahi/compiler/agx_nir_lower_texture.c

@@ -233,8 +233,88 @@ lower_regular_texture(nir_builder *b, nir_instr *instr, UNUSED void *data)
    return true;
 }
 
+static nir_ssa_def *
+bias_for_tex(nir_builder *b, nir_tex_instr *tex)
+{
+   nir_instr *instr = nir_get_texture_size(b, tex)->parent_instr;
+   nir_tex_instr *query = nir_instr_as_tex(instr);
+
+   query->op = nir_texop_lod_bias_agx;
+   query->dest_type = nir_type_float16;
+
+   nir_ssa_dest_init(instr, &query->dest, 1, 16, NULL);
+   return &query->dest.ssa;
+}
+
+static bool
+lower_sampler_bias(nir_builder *b, nir_instr *instr, UNUSED void *data)
+{
+   if (instr->type != nir_instr_type_tex)
+      return false;
+
+   nir_tex_instr *tex = nir_instr_as_tex(instr);
+   b->cursor = nir_before_instr(instr);
+
+   switch (tex->op) {
+   case nir_texop_tex: {
+      tex->op = nir_texop_txb;
+      nir_tex_instr_add_src(tex, nir_tex_src_bias,
+                            nir_src_for_ssa(bias_for_tex(b, tex)));
+      return true;
+   }
+
+   case nir_texop_txb:
+   case nir_texop_txl: {
+      nir_tex_src_type src =
+         tex->op == nir_texop_txl ? nir_tex_src_lod : nir_tex_src_bias;
+
+      nir_ssa_def *orig = steal_tex_src(tex, src);
+      assert(orig != NULL && "invalid NIR");
+
+      if (orig->bit_size != 16)
+         orig = nir_f2f16(b, orig);
+
+      nir_tex_instr_add_src(
+         tex, src, nir_src_for_ssa(nir_fadd(b, orig, bias_for_tex(b, tex))));
+      return true;
+   }
+
+   case nir_texop_txd: {
+      /* For txd, the computed level-of-detail is log2(rho)
+       * where rho should scale proportionally to all
+       * derivatives. So scale derivatives by exp2(bias) to
+       * get level-of-detail log2(exp2(bias) * rho) = bias + log2(rho).
+       */
+      nir_ssa_def *scale = nir_fexp2(b, nir_f2f32(b, bias_for_tex(b, tex)));
+      nir_tex_src_type src[] = {nir_tex_src_ddx, nir_tex_src_ddy};
+
+      for (unsigned s = 0; s < ARRAY_SIZE(src); ++s) {
+         nir_ssa_def *orig = steal_tex_src(tex, src[s]);
+         assert(orig != NULL && "invalid");
+
+         nir_ssa_def *scaled = nir_fmul(b, nir_f2f32(b, orig), scale);
+         nir_tex_instr_add_src(tex, src[s], nir_src_for_ssa(scaled));
+      }
+
+      return true;
+   }
+
+   case nir_texop_txf:
+   case nir_texop_txf_ms:
+   case nir_texop_txs:
+   case nir_texop_tg4:
+   case nir_texop_texture_samples:
+   case nir_texop_samples_identical:
+      /* These operations do not use a sampler */
+      return false;
+
+   default:
+      unreachable("Unhandled texture operation");
+   }
+}
+
 bool
-agx_nir_lower_texture(nir_shader *s)
+agx_nir_lower_texture(nir_shader *s, bool support_lod_bias)
 {
    bool progress = false;
 
@@ -256,6 +336,15 @@ agx_nir_lower_texture(nir_shader *s)
    };
 
    NIR_PASS(progress, s, nir_lower_tex, &lower_tex_options);
+
+   /* Lower bias after nir_lower_tex (to get rid of txd) but before
+    * lower_regular_texture (which will shuffle around the sources)
+    */
+   if (support_lod_bias) {
+      NIR_PASS(progress, s, nir_shader_instructions_pass, lower_sampler_bias,
+               nir_metadata_block_index | nir_metadata_dominance, NULL);
+   }
+
    NIR_PASS(progress, s, nir_legalize_16bit_sampler_srcs, tex_constraints);
 
    /* Lower texture sources after legalizing types (as the lowering depends on

src/asahi/lib/agx_meta.c

@@ -17,7 +17,7 @@ agx_compile_meta_shader(struct agx_meta_cache *cache, nir_shader *shader,
    struct util_dynarray binary;
    util_dynarray_init(&binary, NULL);
 
-   agx_preprocess_nir(shader);
+   agx_preprocess_nir(shader, false);
    if (tib)
       agx_nir_lower_tilebuffer(shader, tib, NULL, NULL);
 

src/gallium/drivers/asahi/agx_state.c

@@ -1531,7 +1531,7 @@ agx_create_shader_state(struct pipe_context *pctx,
    blob_finish(&blob);
 
    so->nir = nir;
-   agx_preprocess_nir(nir);
+   agx_preprocess_nir(nir, true);
 
    /* For shader-db, precompile a shader with a default key. This could be
     * improved but hopefully this is acceptable for now.
@@ -1607,7 +1607,7 @@ agx_create_compute_state(struct pipe_context *pctx,
    blob_finish(&blob);
 
    so->nir = nir;
-   agx_preprocess_nir(nir);
+   agx_preprocess_nir(nir, true);
    agx_get_shader_variant(agx_screen(pctx->screen), so, &pctx->debug, &key);
 
    /* We're done with the NIR, throw it away */