nir: add nir_get_io_data_src

This complements our existing nir_get_io_index_src helper. Most, but annoyingly
not all, stores put their data source in source 0. Having a helper for this lets
us reduce special casing in a bunch of random places.

Signed-off-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Reviewed-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Job Noorman <jnoorman@igalia.com>
(cherry picked from commit 8fb1d65426)

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/40752>
This commit is contained in:
Alyssa Rosenzweig 2026-02-17 09:17:27 -05:00 committed by Eric Engestrom
parent ca4f356d37
commit ef136c2687
3 changed files with 82 additions and 3 deletions

View file

@ -17474,7 +17474,7 @@
"description": "nir: add nir_get_io_data_src",
"nominated": false,
"nomination_type": 0,
"resolution": 4,
"resolution": 1,
"main_sha": null,
"because_sha": null,
"notes": null

View file

@ -5555,10 +5555,12 @@ nir_lower_shader_calls(nir_shader *shader,
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_data_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_data_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);

View file

@ -1078,6 +1078,11 @@ nir_get_io_offset_src(nir_intrinsic_instr *instr)
return idx >= 0 ? &instr->src[idx] : NULL;
}
#define IMG_CASE(name) \
case nir_intrinsic_image_##name: \
case nir_intrinsic_image_deref_##name: \
case nir_intrinsic_bindless_image_##name
/**
* Return the index or handle source number for a load/store intrinsic or -1
* if there's no index or handle.
@ -1107,7 +1112,6 @@ nir_get_io_index_src_number(const nir_intrinsic_instr *instr)
case nir_intrinsic_store_shared_block_intel:
case nir_intrinsic_load_ubo_uniform_block_intel:
case nir_intrinsic_load_ssbo_uniform_block_intel:
#define IMG_CASE(name) case nir_intrinsic_image_##name: case nir_intrinsic_bindless_image_##name
IMG_CASE(load):
IMG_CASE(store):
IMG_CASE(sparse_load):
@ -1123,7 +1127,6 @@ nir_get_io_index_src_number(const nir_intrinsic_instr *instr)
IMG_CASE(order):
IMG_CASE(fragment_mask_load_amd):
return 0;
#undef IMG_CASE
case nir_intrinsic_store_ssbo:
case nir_intrinsic_store_per_vertex_output:
case nir_intrinsic_store_per_view_output:
@ -1137,6 +1140,69 @@ nir_get_io_index_src_number(const nir_intrinsic_instr *instr)
}
}
int
nir_get_io_data_src_number(const nir_intrinsic_instr *intr)
{
switch (intr->intrinsic) {
case nir_intrinsic_store_output:
case nir_intrinsic_store_pixel_local:
case nir_intrinsic_store_per_vertex_output:
case nir_intrinsic_store_per_primitive_output:
case nir_intrinsic_store_ssbo:
case nir_intrinsic_store_ssbo_block_intel:
case nir_intrinsic_store_ssbo_intel:
case nir_intrinsic_store_ssbo_ir3:
case nir_intrinsic_store_shared:
case nir_intrinsic_store_shared_block_intel:
case nir_intrinsic_store_shared_ir3:
case nir_intrinsic_store_task_payload:
case nir_intrinsic_store_global:
case nir_intrinsic_store_global_block_intel:
case nir_intrinsic_store_global_amd:
case nir_intrinsic_store_global_2x32:
case nir_intrinsic_store_global_ir3:
case nir_intrinsic_store_global_etna:
case nir_intrinsic_store_scratch:
case nir_intrinsic_store_raw_output_pan:
case nir_intrinsic_store_combined_output_pan:
case nir_intrinsic_store_tile_pan:
case nir_intrinsic_store_converted_mem_pan:
case nir_intrinsic_store_tlb_sample_color_v3d:
case nir_intrinsic_store_uvs_agx:
case nir_intrinsic_store_local_pixel_agx:
case nir_intrinsic_store_agx:
return 0;
case nir_intrinsic_store_deref:
case nir_intrinsic_shared_atomic:
case nir_intrinsic_shared_atomic_swap:
case nir_intrinsic_deref_atomic:
case nir_intrinsic_deref_atomic_swap:
case nir_intrinsic_global_atomic:
case nir_intrinsic_global_atomic_amd:
case nir_intrinsic_global_atomic_swap:
return 1;
case nir_intrinsic_ssbo_atomic:
case nir_intrinsic_ssbo_atomic_ir3:
case nir_intrinsic_ssbo_atomic_swap:
case nir_intrinsic_ssbo_atomic_swap_ir3:
return 2;
/* clang-format off */
IMG_CASE(store):
IMG_CASE(atomic):
IMG_CASE(atomic_swap):
return 3;
/* clang-format on */
default:
return -1;
}
}
#undef IMG_CASE
/**
* Return the offset or handle source for a load/store intrinsic.
*/
@ -1147,6 +1213,17 @@ nir_get_io_index_src(nir_intrinsic_instr *instr)
return idx >= 0 ? &instr->src[idx] : NULL;
}
/**
* Return the data source for a store intrinsic (including an atomic). For
* atomic swaps, this returns the first of the two contiguous sources.
*/
nir_src *
nir_get_io_data_src(nir_intrinsic_instr *intr)
{
const int idx = nir_get_io_data_src_number(intr);
return idx >= 0 ? &intr->src[idx] : NULL;
}
/**
* Return the array index source number for an arrayed load/store intrinsic or -1 if there's no offset.
*/