asahi: Call nir_lower_blend with selected key

Also need to key to pipe formats.

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

2 files changed, 36 insertions, 4 deletions

diff --git a/src/gallium/drivers/asahi/agx_state.c b/src/gallium/drivers/asahi/agx_state.c
index 4dfdf366862..ab54109edb7 100644
--- a/src/gallium/drivers/asahi/agx_state.c
+++ b/src/gallium/drivers/asahi/agx_state.c
@@ -626,6 +626,25 @@ agx_update_shader(struct agx_context *ctx, struct agx_compiled_shader **out,
    util_dynarray_init(&binary, NULL);
 
    nir_shader *nir = nir_shader_clone(NULL, so->nir);
+
+   if (key->blend.blend_enable) {
+      nir_lower_blend_options opts = {
+         .format = { key->rt_formats[0] },
+         .scalar_blend_const = true
+      };
+
+      memcpy(opts.rt, key->blend.rt, sizeof(opts.rt));
+      NIR_PASS_V(nir, nir_lower_blend, opts);
+   } else if (key->blend.logicop_enable) {
+      nir_lower_blend_options opts = {
+         .format = { key->rt_formats[0] },
+         .logicop_enable = true,
+         .logicop_func = key->blend.logicop_func,
+      };
+
+      NIR_PASS_V(nir, nir_lower_blend, opts);
+   }
+
    agx_compile_shader_nir(nir, &key->base, &binary, &compiled->info);
 
    /* TODO: emit this properly */
@@ -715,12 +734,23 @@ agx_update_vs(struct agx_context *ctx)
 static bool
 agx_update_fs(struct agx_context *ctx)
 {
-   struct agx_fs_shader_key key = {
+   struct agx_fs_shader_key base_key = {
       .tib_formats = { AGX_FORMAT_U8NORM }
    };
 
-   return agx_update_shader(ctx, &ctx->fs, PIPE_SHADER_FRAGMENT,
-                            (struct asahi_shader_key *) &key);
+   struct asahi_shader_key key = {
+      .base.fs = base_key,
+      .nr_cbufs = ctx->batch->nr_cbufs,
+   };
+
+   for (unsigned i = 0; i < key.nr_cbufs; ++i) {
+      key.rt_formats[i] = ctx->batch->cbufs[i] ?
+         ctx->batch->cbufs[i]->format : PIPE_FORMAT_NONE;
+   }
+
+   memcpy(&key.blend, ctx->blend, sizeof(key.blend));
+
+   return agx_update_shader(ctx, &ctx->fs, PIPE_SHADER_FRAGMENT, &key);
 }
 
 static void
diff --git a/src/gallium/drivers/asahi/agx_state.h b/src/gallium/drivers/asahi/agx_state.h
index 0cbabbaf2fc..a8d0aebcc10 100644
--- a/src/gallium/drivers/asahi/agx_state.h
+++ b/src/gallium/drivers/asahi/agx_state.h
@@ -99,7 +99,7 @@ struct agx_zsa {
 };
 
 struct agx_blend {
-   bool logicop_enable;
+   bool logicop_enable, blend_enable;
 
    union {
       nir_lower_blend_rt rt[8];
@@ -110,6 +110,8 @@ struct agx_blend {
 struct asahi_shader_key {
    struct agx_shader_key base;
    struct agx_blend blend;
+   unsigned nr_cbufs;
+   enum pipe_format rt_formats[PIPE_MAX_COLOR_BUFS];
 };
 
 #define AGX_DIRTY_VERTEX (1 << 0)