From 7ac6ec2ceb0b858dfb440e55961fcf7b0528d9a9 Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Wed, 5 Feb 2025 07:25:13 -0600
Subject: [PATCH] nir: Add a get_io_index_src() helper

Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33402>
---
 src/compiler/nir/nir.h          |  2 ++
 src/compiler/nir/nir_lower_io.c | 40 +++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index b8506163e90..a293beef8af 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -5154,9 +5154,11 @@ nir_lower_shader_calls(nir_shader *shader,
                        void *mem_ctx);
 
 int nir_get_io_offset_src_number(const nir_intrinsic_instr *instr);
+int nir_get_io_index_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_index_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 9ea8f39ff7e..7ab2c63330e 100644
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -2873,6 +2873,46 @@ nir_get_io_offset_src(nir_intrinsic_instr *instr)
    return idx >= 0 ? &instr->src[idx] : NULL;
 }
 
+/**
+ * Return the index or handle source number for a load/store intrinsic or -1
+ * if there's no index or handle.
+ */
+int
+nir_get_io_index_src_number(const nir_intrinsic_instr *instr)
+{
+   switch (instr->intrinsic) {
+   case nir_intrinsic_load_ubo:
+   case nir_intrinsic_load_ssbo:
+   case nir_intrinsic_load_input_vertex:
+   case nir_intrinsic_load_per_vertex_input:
+   case nir_intrinsic_load_per_vertex_output:
+   case nir_intrinsic_load_per_view_output:
+   case nir_intrinsic_load_per_primitive_output:
+   case nir_intrinsic_load_interpolated_input:
+   case nir_intrinsic_load_smem_amd:
+   case nir_intrinsic_ldc_nv:
+   case nir_intrinsic_ldcx_nv:
+      return 0;
+   case nir_intrinsic_store_ssbo:
+   case nir_intrinsic_store_per_vertex_output:
+   case nir_intrinsic_store_per_view_output:
+   case nir_intrinsic_store_per_primitive_output:
+      return 1;
+   default:
+      return -1;
+   }
+}
+
+/**
+ * Return the offset or handle source for a load/store intrinsic.
+ */
+nir_src *
+nir_get_io_index_src(nir_intrinsic_instr *instr)
+{
+   const int idx = nir_get_io_index_src_number(instr);
+   return idx >= 0 ? &instr->src[idx] : NULL;
+}
+
 /**
  * Return the array index source number for an arrayed load/store intrinsic or -1 if there's no offset.
  */