ac/nir: Add helper for I/O location mapping.

Map I/O locations based on a prefix sum (for linked shaders),
or based on the provided callback.

Signed-off-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Reviewed-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/29812>

src/amd/common/ac_nir.c

@@ -594,6 +594,26 @@ ac_nir_export_parameters(nir_builder *b,
    }
 }
 
+unsigned
+ac_nir_map_io_location(unsigned location,
+                       uint64_t mask,
+                       ac_nir_map_io_driver_location map_io)
+{
+   /* Unlinked shaders:
+    * We are unaware of the inputs of the next stage while lowering outputs.
+    * The driver needs to pass a callback to map varyings to a fixed location.
+    */
+   if (map_io)
+      return map_io(location);
+
+   /* Linked shaders:
+    * Take advantage of knowledge of the inputs of the next stage when lowering outputs.
+    * Map varyings to a prefix sum of the IO mask to save space in LDS or VRAM.
+    */
+   assert(mask & BITFIELD64_BIT(location));
+   return util_bitcount64(mask & BITFIELD64_MASK(location));
+}
+
 /**
  * This function takes an I/O intrinsic like load/store_input,
  * and emits a sequence that calculates the full offset of that instruction,

src/amd/common/ac_nir_helpers.h

@@ -125,6 +125,11 @@ ac_nir_calc_io_offset_mapped(nir_builder *b,
                              unsigned component_stride,
                              unsigned mapped_location);
 
+unsigned
+ac_nir_map_io_location(unsigned location,
+                       uint64_t mask,
+                       ac_nir_map_io_driver_location map_io);
+
 nir_def *
 ac_nir_cull_primitive(nir_builder *b,
                       nir_def *initially_accepted,