From 7d94f2ee49caa3201e8a6605a9103d719d2c5ffe Mon Sep 17 00:00:00 2001
From: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Date: Mon, 28 Aug 2023 08:14:58 -0400
Subject: [PATCH] agx: Add pass to lower layer ID writes

The hardware needs the layer ID and the viewport index packed together. That
consumes an entire varying slot, if we want those available in the frag shader
we need a separate slot. Add a pass to insert the extra packed write.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
---
 src/asahi/compiler/agx_nir.h             |  1 +
 src/asahi/compiler/agx_nir_lower_layer.c | 67 ++++++++++++++++++++++++
 src/asahi/compiler/meson.build           |  1 +
 3 files changed, 69 insertions(+)
 create mode 100644 src/asahi/compiler/agx_nir_lower_layer.c

diff --git a/src/asahi/compiler/agx_nir.h b/src/asahi/compiler/agx_nir.h
index db1ddea1111..4ef3a326299 100644
--- a/src/asahi/compiler/agx_nir.h
+++ b/src/asahi/compiler/agx_nir.h
@@ -15,5 +15,6 @@ bool agx_nir_lower_algebraic_late(struct nir_shader *shader);
 bool agx_nir_fuse_algebraic_late(struct nir_shader *shader);
 bool agx_nir_fence_images(struct nir_shader *shader);
 bool agx_nir_lower_multisampled_image_store(struct nir_shader *s);
+void agx_nir_lower_layer(struct nir_shader *s);
 
 #endif
diff --git a/src/asahi/compiler/agx_nir_lower_layer.c b/src/asahi/compiler/agx_nir_lower_layer.c
new file mode 100644
index 00000000000..4d606ce952b
--- /dev/null
+++ b/src/asahi/compiler/agx_nir_lower_layer.c
@@ -0,0 +1,67 @@
+/*
+ * Copyright 2023 Valve Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "compiler/nir/nir.h"
+#include "compiler/nir/nir_builder.h"
+#include "agx_nir.h"
+
+void
+agx_nir_lower_layer(nir_shader *s)
+{
+   assert(s->info.stage == MESA_SHADER_VERTEX);
+   if (!(s->info.outputs_written & (VARYING_BIT_LAYER | VARYING_BIT_VIEWPORT)))
+      return;
+
+   /* Writes are in the last block, search */
+   nir_function_impl *impl = nir_shader_get_entrypoint(s);
+   nir_block *last = nir_impl_last_block(impl);
+
+   nir_def *layer = NULL, *viewport = NULL;
+   nir_cursor last_cursor;
+
+   nir_foreach_instr(instr, last) {
+      if (instr->type != nir_instr_type_intrinsic)
+         continue;
+
+      nir_intrinsic_instr *store = nir_instr_as_intrinsic(instr);
+      if (store->intrinsic != nir_intrinsic_store_output)
+         continue;
+
+      nir_io_semantics sem = nir_intrinsic_io_semantics(store);
+      nir_def *value = store->src[0].ssa;
+
+      if (sem.location == VARYING_SLOT_LAYER) {
+         assert(layer == NULL && "only written once");
+         layer = value;
+      } else if (sem.location == VARYING_SLOT_VIEWPORT) {
+         assert(viewport == NULL && "only written once");
+         viewport = value;
+      } else {
+         continue;
+      }
+
+      last_cursor = nir_after_instr(&store->instr);
+
+      /* Leave the store as a varying-only, no sysval output */
+      sem.no_sysval_output = true;
+      nir_intrinsic_set_io_semantics(store, sem);
+   }
+
+   assert((layer || viewport) && "metadata inconsistent with program");
+
+   /* Pack together and write out */
+   nir_builder b = nir_builder_at(last_cursor);
+
+   nir_def *zero = nir_imm_intN_t(&b, 0, 16);
+   nir_def *packed =
+      nir_pack_32_2x16_split(&b, layer ? nir_u2u16(&b, layer) : zero,
+                             viewport ? nir_u2u16(&b, viewport) : zero);
+
+   /* Written with a sysval-only store, no varying output */
+   nir_store_output(&b, packed, nir_imm_int(&b, 0),
+                    .io_semantics.location = VARYING_SLOT_LAYER,
+                    .io_semantics.num_slots = 1,
+                    .io_semantics.no_varying = true);
+}
diff --git a/src/asahi/compiler/meson.build b/src/asahi/compiler/meson.build
index 3a985fbdb60..8130dbb3eb8 100644
--- a/src/asahi/compiler/meson.build
+++ b/src/asahi/compiler/meson.build
@@ -14,6 +14,7 @@ libasahi_agx_files = files(
   'agx_nir_lower_discard_zs_emit.c',
   'agx_nir_lower_texture.c',
   'agx_nir_lower_interpolation.c',
+  'agx_nir_lower_layer.c',
   'agx_nir_lower_load_mask.c',
   'agx_nir_lower_shared_bitsize.c',
   'agx_nir_opt_preamble.c',