asahi,nir: Stop relying on zero and scratch page in GS/TESS code

Introduce new NIR intrinsics to handle getting a "sink" read-only
address and another intrinsic to handle conversion of address to
read-write (allowing implementation to replace the "sink" read-only with
another address like required for Asahi)

Signed-off-by: Mary Guillemard <mary.guillemard@collabora.com>
Reviewed-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37914>

src/asahi/compiler/agx_compile.c

@@ -3301,6 +3301,30 @@ lower_load_from_texture_handle(nir_builder *b, nir_intrinsic_instr *intr,
    return true;
 }
 
+static bool
+lower_sink_address(nir_builder *b, nir_intrinsic_instr *intr, UNUSED void *data)
+{
+   switch (intr->intrinsic) {
+   case nir_intrinsic_load_ro_sink_address_poly:
+      b->cursor = nir_before_instr(&intr->instr);
+      nir_def_replace(&intr->def, nir_imm_int64(b, AGX_ZERO_PAGE_ADDRESS));
+      return true;
+   case nir_intrinsic_ro_to_rw_poly: {
+      b->cursor = nir_before_instr(&intr->instr);
+      nir_def_replace(
+         &intr->def,
+         nir_bcsel(b, nir_ieq_imm(b, intr->src[0].ssa, AGX_ZERO_PAGE_ADDRESS),
+                   nir_imm_int64(b, AGX_SCRATCH_PAGE_ADDRESS),
+                   intr->src[0].ssa));
+      return true;
+   }
+   default:
+      break;
+   }
+
+   return false;
+}
+
 static void
 agx_remove_unreachable_block(agx_block *block)
 {
@@ -3823,6 +3847,9 @@ agx_compile_shader_nir(nir_shader *nir, struct agx_shader_key *key,
    NIR_PASS(_, nir, nir_shader_intrinsics_pass, lower_load_from_texture_handle,
             nir_metadata_control_flow, NULL);
 
+   NIR_PASS(_, nir, nir_shader_intrinsics_pass, lower_sink_address,
+            nir_metadata_control_flow, NULL);
+
    info->push_count = key->reserved_preamble;
    agx_optimize_nir(
       nir, key->dev.soft_fault, key->secondary ? NULL : &info->push_count,

src/asahi/libagx/geometry.cl

@@ -9,10 +9,11 @@
 #include "util/macros.h"
 #include "util/u_math.h"
 #include "geometry.h"
-#include "libagx_intrinsics.h"
 #include "query.h"
 #include "tessellator.h"
 
+uint64_t nir_ro_to_rw_poly(uint64_t address);
+
 /* Swap the two non-provoking vertices in odd triangles. This generates a vertex
  * ID list with a consistent winding order.
  *
@@ -941,10 +942,8 @@ libagx_pre_gs(global struct agx_geometry_params *p, uint streams,
          unsigned stream = buffer_to_stream[i];
 
          global uint *ptr = p->xfb_offs_ptrs[i];
-         if ((uintptr_t)ptr == AGX_ZERO_PAGE_ADDRESS) {
-            ptr = (global uint *)AGX_SCRATCH_PAGE_ADDRESS;
-         }
 
+         ptr = (global uint *)nir_ro_to_rw_poly((uint64_t)ptr);
          *ptr += prims[stream] * prim_stride_B;
       }
    }

src/asahi/libagx/geometry.h

@@ -154,6 +154,18 @@ agx_heap_alloc_nonatomic(global struct agx_heap *heap, uint size_B)
 {
    return heap->base + agx_heap_alloc_nonatomic_offs(heap, size_B);
 }
+
+uint64_t nir_load_ro_sink_address_poly(void);
+
+static inline uint64_t
+libagx_index_buffer(uint64_t index_buffer, uint size_el, uint offset_el,
+                    uint elsize_B)
+{
+   if (offset_el < size_el)
+      return index_buffer + (offset_el * elsize_B);
+   else
+      return nir_load_ro_sink_address_poly();
+}
 #endif
 
 struct agx_ia_state {
@@ -170,16 +182,6 @@ struct agx_ia_state {
 } PACKED;
 static_assert(sizeof(struct agx_ia_state) == 4 * 4);
 
-static inline uint64_t
-libagx_index_buffer(uint64_t index_buffer, uint size_el, uint offset_el,
-                    uint elsize_B)
-{
-   if (offset_el < size_el)
-      return index_buffer + (offset_el * elsize_B);
-   else
-      return AGX_ZERO_PAGE_ADDRESS;
-}
-
 static inline uint
 libagx_index_buffer_range_el(uint size_el, uint offset_el)
 {

src/asahi/vulkan/hk_cmd_draw.c

@@ -1030,6 +1030,16 @@ hk_heap(struct hk_cmd_buffer *cmd)
    return dev->rodata.heap;
 }
 
+static uint64_t
+hk_index_buffer(uint64_t index_buffer, uint size_el, uint offset_el,
+                uint elsize_B)
+{
+   if (offset_el < size_el)
+      return index_buffer + (offset_el * elsize_B);
+   else
+      return AGX_ZERO_PAGE_ADDRESS;
+}
+
 static uint64_t
 hk_upload_ia_params(struct hk_cmd_buffer *cmd, struct agx_draw draw)
 {
@@ -1041,8 +1051,8 @@ hk_upload_ia_params(struct hk_cmd_buffer *cmd, struct agx_draw draw)
       unsigned index_size_B = agx_index_size_to_B(draw.index_size);
       unsigned range_el = agx_draw_index_range_el(draw);
 
-      ia.index_buffer = libagx_index_buffer(agx_draw_index_buffer(draw),
-                                            range_el, 0, index_size_B);
+      ia.index_buffer = hk_index_buffer(agx_draw_index_buffer(draw), range_el,
+                                        0, index_size_B);
 
       ia.index_buffer_range_el = range_el;
    }

src/compiler/nir/nir_intrinsics.py

@@ -1407,6 +1407,13 @@ intrinsic("select_vertex_poly", src_comp=[1], indices=[STREAM_ID])
 # Sources: (index offset, first vertex, number of vertices, # of XFB primitives before).
 intrinsic("emit_primitive_poly", src_comp=[1, 1, 1, 1], indices=[STREAM_ID])
 
+# Transform a given address to be used as read-write,
+# allowing to transition a "sink" read-only address to the read-write address.
+intrinsic("ro_to_rw_poly", src_comp=[1], dest_comp=1, flags=[CAN_ELIMINATE, CAN_REORDER], bit_sizes=[64])
+
+# Get the address representing the read-only "sink" address (always read 0)
+system_value("ro_sink_address_poly", 1, bit_sizes=[64])
+
 # mesa_prim for the input topology (in a geometry shader)
 system_value("input_topology_poly", 1)