From e2a0d64d52415994cfe2af8a830da3c52fa62d2a Mon Sep 17 00:00:00 2001
From: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Date: Wed, 16 Aug 2023 17:58:06 -0400
Subject: [PATCH] asahi: Add pass to predicate layer ID reads

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
---
 src/asahi/lib/agx_nir_predicate_layer_id.c | 54 ++++++++++++++++++++++
 src/asahi/lib/agx_tilebuffer.h             |  2 +
 src/asahi/lib/meson.build                  |  1 +
 3 files changed, 57 insertions(+)
 create mode 100644 src/asahi/lib/agx_nir_predicate_layer_id.c

diff --git a/src/asahi/lib/agx_nir_predicate_layer_id.c b/src/asahi/lib/agx_nir_predicate_layer_id.c
new file mode 100644
index 00000000000..2dba0ccf5bb
--- /dev/null
+++ b/src/asahi/lib/agx_nir_predicate_layer_id.c
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2023 Alyssa Rosenzweig
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "compiler/shader_enums.h"
+#include "agx_tilebuffer.h"
+#include "nir.h"
+#include "nir_builder.h"
+#include "nir_builder_opcodes.h"
+#include "nir_intrinsics_indices.h"
+
+/*
+ * If a fragment shader reads the layer ID but the vertex shader does not write
+ * the layer ID, the fragment shader is supposed to read zero. However, in our
+ * hardware, if the vertex shader does not write the layer ID, the value read by
+ * the fragment shader is UNDEFINED. To reconcile, the driver passes in whether
+ * the layer ID value is written, and this pass predicates layer ID on that
+ * system value. This handles both cases without shader variants, at the cost of
+ * a single instruction.
+ */
+static bool
+lower(nir_builder *b, nir_intrinsic_instr *intr, void *_)
+{
+   if (intr->intrinsic != nir_intrinsic_load_input)
+      return false;
+
+   if (nir_intrinsic_io_semantics(intr).location != VARYING_SLOT_LAYER)
+      return false;
+
+   b->cursor = nir_after_instr(&intr->instr);
+   nir_def *written = nir_load_layer_id_written_agx(b);
+
+   /* Zero extend the mask since layer IDs are 16-bits, so upper bits of the
+    * layer ID are necessarily zero.
+    */
+   nir_def *mask = nir_u2uN(b, written, intr->def.bit_size);
+   nir_def *repl = nir_iand(b, &intr->def, mask);
+   nir_def_rewrite_uses_after(&intr->def, repl, repl->parent_instr);
+   return true;
+}
+
+void
+agx_nir_predicate_layer_id(nir_shader *shader)
+{
+   assert(shader->info.stage == MESA_SHADER_FRAGMENT);
+
+   /* If layer is not read, there's nothing to lower */
+   if (shader->info.inputs_read & VARYING_BIT_LAYER) {
+      nir_shader_intrinsics_pass(
+         shader, lower, nir_metadata_block_index | nir_metadata_dominance,
+         NULL);
+   }
+}
diff --git a/src/asahi/lib/agx_tilebuffer.h b/src/asahi/lib/agx_tilebuffer.h
index 3ad3dea74ae..e0911642af0 100644
--- a/src/asahi/lib/agx_tilebuffer.h
+++ b/src/asahi/lib/agx_tilebuffer.h
@@ -112,6 +112,8 @@ void agx_nir_lower_alpha_to_coverage(struct nir_shader *shader,
 
 void agx_nir_lower_alpha_to_one(struct nir_shader *shader);
 
+void agx_nir_predicate_layer_id(struct nir_shader *shader);
+
 void agx_usc_tilebuffer(struct agx_usc_builder *b,
                         struct agx_tilebuffer_layout *tib);
 
diff --git a/src/asahi/lib/meson.build b/src/asahi/lib/meson.build
index 7a91d3a1970..e595f5ca055 100644
--- a/src/asahi/lib/meson.build
+++ b/src/asahi/lib/meson.build
@@ -16,6 +16,7 @@ libasahi_lib_files = files(
   'agx_nir_lower_sample_intrinsics.c',
   'agx_nir_lower_tilebuffer.c',
   'agx_nir_lower_vbo.c',
+  'agx_nir_predicate_layer_id.c',
   'agx_ppp.h',
   'pool.c',
 )