radv: Use mapped driver locations for determining I/O strides.

This will allow us to more accurately determine the
input and output strides, because the I/O locations mapped
by RADV don't match the locations in NIR.
As a result, ESO will use less LDS.

It also fixes the per-patch output stride of tess control
shaders, because previously we omitted tess factors from them.

Signed-off-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/28021>

3 files changed, 45 insertions, 4 deletions

diff --git a/src/amd/vulkan/nir/radv_nir.h b/src/amd/vulkan/nir/radv_nir.h
index e6315d4a2f0..ba4ff31ece5 100644
--- a/src/amd/vulkan/nir/radv_nir.h
+++ b/src/amd/vulkan/nir/radv_nir.h
@@ -76,6 +76,8 @@ bool radv_nir_export_multiview(nir_shader *nir);
 
 void radv_nir_lower_io_to_scalar_early(nir_shader *nir, nir_variable_mode mask);
 
+unsigned radv_map_io_driver_location(unsigned semantic);
+
 void radv_nir_lower_io(struct radv_device *device, nir_shader *nir);
 
 bool radv_nir_lower_io_to_mem(struct radv_device *device, struct radv_shader_stage *stage);
diff --git a/src/amd/vulkan/nir/radv_nir_lower_io.c b/src/amd/vulkan/nir/radv_nir_lower_io.c
index 839378fb1bb..49a18375af9 100644
--- a/src/amd/vulkan/nir/radv_nir_lower_io.c
+++ b/src/amd/vulkan/nir/radv_nir_lower_io.c
@@ -108,7 +108,7 @@ enum {
    RADV_IO_SLOT_VAR0, /* 0..31 */
 };
 
-static unsigned
+unsigned
 radv_map_io_driver_location(unsigned semantic)
 {
    if ((semantic >= VARYING_SLOT_PATCH0 && semantic < VARYING_SLOT_TESS_MAX) ||
diff --git a/src/amd/vulkan/radv_shader_info.c b/src/amd/vulkan/radv_shader_info.c
index 42accf60102..4122a9585c5 100644
--- a/src/amd/vulkan/radv_shader_info.c
+++ b/src/amd/vulkan/radv_shader_info.c
@@ -22,6 +22,7 @@
  */
 #include "nir/nir.h"
 #include "nir/nir_xfb_info.h"
+#include "nir/radv_nir.h"
 #include "radv_private.h"
 #include "radv_shader.h"
 
@@ -422,13 +423,49 @@ gather_shader_info_ngg_query(struct radv_device *device, struct radv_shader_info
    info->has_prim_query = device->cache_key.primitives_generated_query || info->has_xfb_query;
 }
 
+static uint64_t
+gather_io_mask(const uint64_t nir_mask, const uint64_t nir_patch_mask, const bool per_patch)
+{
+   /* Select the correct NIR IO mask.
+    * Exclude per-patch built-in variables, because in NIR they are not part of the patch I/O masks.
+    */
+   const uint64_t nir_io_mask =
+      (per_patch ? nir_patch_mask : nir_mask) & ~(VARYING_BIT_TESS_LEVEL_OUTER | VARYING_BIT_TESS_LEVEL_INNER);
+
+   /* Create a mask of driver locations mapped from NIR semantics. */
+   uint64_t radv_io_mask = 0;
+   u_foreach_bit64 (semantic, nir_io_mask) {
+      /* These outputs are not used when fixed output slots are needed. */
+      if (semantic == VARYING_SLOT_LAYER || semantic == VARYING_SLOT_VIEWPORT ||
+          semantic == VARYING_SLOT_PRIMITIVE_ID || semantic == VARYING_SLOT_PRIMITIVE_SHADING_RATE)
+         continue;
+
+      /* Generic per-patch outputs start at an offset. */
+      if (per_patch)
+         semantic += VARYING_SLOT_PATCH0;
+
+      radv_io_mask |= BITFIELD64_BIT(radv_map_io_driver_location(semantic));
+   }
+
+   /* Include per-patch built-in variables. */
+   if (per_patch) {
+      if (nir_mask & VARYING_BIT_TESS_LEVEL_OUTER)
+         radv_io_mask |= BITFIELD64_BIT(radv_map_io_driver_location(VARYING_SLOT_TESS_LEVEL_OUTER));
+      if (nir_mask & VARYING_BIT_TESS_LEVEL_INNER)
+         radv_io_mask |= BITFIELD64_BIT(radv_map_io_driver_location(VARYING_SLOT_TESS_LEVEL_INNER));
+   }
+
+   return radv_io_mask;
+}
+
 static void
 gather_info_unlinked_input(struct radv_shader_info *info, const nir_shader *nir)
 {
    if (info->inputs_linked)
       return;
 
-   const unsigned num_linked_inputs = util_last_bit64(nir->info.inputs_read);
+   const uint64_t io_mask = gather_io_mask(nir->info.inputs_read, 0, false);
+   const unsigned num_linked_inputs = util_last_bit64(io_mask);
 
    switch (nir->info.stage) {
    case MESA_SHADER_TESS_CTRL:
@@ -451,7 +488,9 @@ gather_info_unlinked_output(struct radv_shader_info *info, const nir_shader *nir
    if (info->outputs_linked)
       return;
 
-   const unsigned num_linked_outputs = util_last_bit64(nir->info.outputs_written);
+   const uint64_t io_mask = gather_io_mask(nir->info.outputs_written, 0, false);
+   const uint64_t patch_io_mask = gather_io_mask(nir->info.outputs_written, nir->info.patch_outputs_written, true);
+   const unsigned num_linked_outputs = util_last_bit64(io_mask);
 
    switch (nir->info.stage) {
    case MESA_SHADER_VERTEX:
@@ -459,7 +498,7 @@ gather_info_unlinked_output(struct radv_shader_info *info, const nir_shader *nir
       break;
    case MESA_SHADER_TESS_CTRL:
       info->tcs.num_linked_outputs = num_linked_outputs;
-      info->tcs.num_linked_patch_outputs = util_last_bit64(nir->info.patch_outputs_written);
+      info->tcs.num_linked_patch_outputs = util_last_bit64(patch_io_mask);
       break;
    case MESA_SHADER_TESS_EVAL:
       info->tes.num_linked_outputs = num_linked_outputs;