nir: Make a const-friendly way to get the offset_src and arrayed_io_src from intrinsic

The existing helper returns a `nir_src *` so expects a non-const instr.

We plan to use this function in queries that don't modify the shader, so
create (and use internally) a variant that returns the index instead.

Reviewed-by: Jesse Natalie <jenatali@microsoft.com>
Acked-by: Emma Anholt <emma@anholt.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/23562>

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index ac522e7..665a46f 100644 (file)
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -5226,6 +5226,9 @@ nir_lower_shader_calls(nir_shader *shader,
                        uint32_t *num_resume_shaders_out,
                        void *mem_ctx);
 
+int nir_get_io_offset_src_number(const nir_intrinsic_instr *instr);
+int nir_get_io_arrayed_index_src_number(const nir_intrinsic_instr *instr);
+
 nir_src *nir_get_io_offset_src(nir_intrinsic_instr *instr);
 nir_src *nir_get_io_arrayed_index_src(nir_intrinsic_instr *instr);
 nir_src *nir_get_shader_call_payload_src(nir_intrinsic_instr *call);

diff --git a/src/compiler/nir/nir_lower_io.c b/src/compiler/nir/nir_lower_io.c
index a5a524b..db42666 100644 (file)
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -2597,10 +2597,10 @@ nir_gather_explicit_io_initializers(nir_shader *shader,
 }
 
 /**
- * Return the offset source for a load/store intrinsic.
+ * Return the offset source number for a load/store intrinsic or -1 if there's no offset.
  */
-nir_src *
-nir_get_io_offset_src(nir_intrinsic_instr *instr)
+int
+nir_get_io_offset_src_number(const nir_intrinsic_instr *instr)
 {
    switch (instr->intrinsic) {
    case nir_intrinsic_load_input:
@@ -2620,7 +2620,7 @@ nir_get_io_offset_src(nir_intrinsic_instr *instr)
    case nir_intrinsic_task_payload_atomic_swap:
    case nir_intrinsic_global_atomic:
    case nir_intrinsic_global_atomic_swap:
-      return &instr->src[0];
+      return 0;
    case nir_intrinsic_load_ubo:
    case nir_intrinsic_load_ssbo:
    case nir_intrinsic_load_input_vertex:
@@ -2636,36 +2636,56 @@ nir_get_io_offset_src(nir_intrinsic_instr *instr)
    case nir_intrinsic_store_scratch:
    case nir_intrinsic_ssbo_atomic:
    case nir_intrinsic_ssbo_atomic_swap:
-      return &instr->src[1];
+      return 1;
    case nir_intrinsic_store_ssbo:
    case nir_intrinsic_store_per_vertex_output:
    case nir_intrinsic_store_per_primitive_output:
-      return &instr->src[2];
+      return 2;
    default:
-      return NULL;
+      return -1;
    }
 }
 
 /**
- * Return the vertex index source for a load/store per_vertex intrinsic.
+ * Return the offset source for a load/store intrinsic.
  */
 nir_src *
-nir_get_io_arrayed_index_src(nir_intrinsic_instr *instr)
+nir_get_io_offset_src(nir_intrinsic_instr *instr)
+{
+   const int idx = nir_get_io_offset_src_number(instr);
+   return idx >= 0 ? &instr->src[idx] : NULL;
+}
+
+/**
+ * Return the vertex index source number for a load/store per_vertex intrinsic or -1 if there's no offset.
+ */
+int
+nir_get_io_arrayed_index_src_number(const nir_intrinsic_instr *instr)
 {
    switch (instr->intrinsic) {
    case nir_intrinsic_load_per_vertex_input:
    case nir_intrinsic_load_per_vertex_output:
    case nir_intrinsic_load_per_primitive_output:
-      return &instr->src[0];
+      return 0;
    case nir_intrinsic_store_per_vertex_output:
    case nir_intrinsic_store_per_primitive_output:
-      return &instr->src[1];
+      return 1;
    default:
-      return NULL;
+      return -1;
    }
 }
 
 /**
+ * Return the vertex index source for a load/store per_vertex intrinsic.
+ */
+nir_src *
+nir_get_io_arrayed_index_src(nir_intrinsic_instr *instr)
+{
+   const int idx = nir_get_io_arrayed_index_src_number(instr);
+   return idx >= 0 ? &instr->src[idx] : NULL;
+}
+
+/**
  * Return the numeric constant that identify a NULL pointer for each address
  * format.
  */