+static bool
+find_descriptor_for_index_src(nir_src src,
+ struct apply_pipeline_layout_state *state)
+{
+ nir_intrinsic_instr *intrin = nir_src_as_intrinsic(src);
+
+ while (intrin && intrin->intrinsic == nir_intrinsic_vulkan_resource_reindex)
+ intrin = nir_src_as_intrinsic(intrin->src[0]);
+
+ if (!intrin || intrin->intrinsic != nir_intrinsic_vulkan_resource_index)
+ return false;
+
+ uint32_t set = nir_intrinsic_desc_set(intrin);
+ uint32_t binding = nir_intrinsic_binding(intrin);
+ uint32_t surface_index = state->set[set].surface_offsets[binding];
+
+ /* Only lower to a BTI message if we have a valid binding table index. */
+ return surface_index < MAX_BINDING_TABLE_SIZE;
+}
+
+static bool
+nir_deref_find_descriptor(nir_deref_instr *deref,
+ struct apply_pipeline_layout_state *state)
+{
+ while (1) {
+ /* Nothing we will use this on has a variable */
+ assert(deref->deref_type != nir_deref_type_var);
+
+ nir_deref_instr *parent = nir_src_as_deref(deref->parent);
+ if (!parent)
+ break;
+
+ deref = parent;
+ }
+ assert(deref->deref_type == nir_deref_type_cast);
+
+ nir_intrinsic_instr *intrin = nir_src_as_intrinsic(deref->parent);
+ if (!intrin || intrin->intrinsic != nir_intrinsic_load_vulkan_descriptor)
+ return false;
+
+ return find_descriptor_for_index_src(intrin->src[0], state);
+}
+
+static nir_ssa_def *
+build_index_for_res_reindex(nir_intrinsic_instr *intrin,
+ struct apply_pipeline_layout_state *state)
+{
+ nir_builder *b = &state->builder;
+
+ if (intrin->intrinsic == nir_intrinsic_vulkan_resource_reindex) {
+ nir_ssa_def *bti =
+ build_index_for_res_reindex(nir_src_as_intrinsic(intrin->src[0]), state);
+
+ b->cursor = nir_before_instr(&intrin->instr);
+ return nir_iadd(b, bti, nir_ssa_for_src(b, intrin->src[1], 1));
+ }
+
+ assert(intrin->intrinsic == nir_intrinsic_vulkan_resource_index);
+
+ uint32_t set = nir_intrinsic_desc_set(intrin);
+ uint32_t binding = nir_intrinsic_binding(intrin);
+
+ const struct anv_descriptor_set_binding_layout *bind_layout =
+ &state->layout->set[set].layout->binding[binding];
+
+ uint32_t surface_index = state->set[set].surface_offsets[binding];
+ uint32_t array_size = bind_layout->array_size;
+
+ b->cursor = nir_before_instr(&intrin->instr);
+
+ nir_ssa_def *array_index = nir_ssa_for_src(b, intrin->src[0], 1);
+ if (nir_src_is_const(intrin->src[0]) || state->add_bounds_checks)
+ array_index = nir_umin(b, array_index, nir_imm_int(b, array_size - 1));
+
+ return nir_iadd_imm(b, array_index, surface_index);
+}
+
+static nir_ssa_def *
+build_index_offset_for_deref(nir_deref_instr *deref,
+ struct apply_pipeline_layout_state *state)
+{
+ nir_builder *b = &state->builder;
+
+ nir_deref_instr *parent = nir_deref_instr_parent(deref);
+ if (parent) {
+ nir_ssa_def *addr = build_index_offset_for_deref(parent, state);
+
+ b->cursor = nir_before_instr(&deref->instr);
+ return nir_explicit_io_address_from_deref(b, deref, addr,
+ nir_address_format_32bit_index_offset);
+ }
+
+ nir_intrinsic_instr *load_desc = nir_src_as_intrinsic(deref->parent);
+ assert(load_desc->intrinsic == nir_intrinsic_load_vulkan_descriptor);
+
+ nir_ssa_def *index =
+ build_index_for_res_reindex(nir_src_as_intrinsic(load_desc->src[0]), state);
+
+ /* Return a 0 offset which will get picked up by the recursion */
+ b->cursor = nir_before_instr(&deref->instr);
+ return nir_vec2(b, index, nir_imm_int(b, 0));
+}
+
+static bool
+try_lower_direct_buffer_intrinsic(nir_intrinsic_instr *intrin, bool is_atomic,
+ struct apply_pipeline_layout_state *state)
+{
+ nir_builder *b = &state->builder;
+
+ nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
+ if (deref->mode != nir_var_mem_ssbo)
+ return false;
+
+ /* 64-bit atomics only support A64 messages so we can't lower them to the
+ * index+offset model.
+ */
+ if (is_atomic && nir_dest_bit_size(intrin->dest) == 64)
+ return false;
+
+ /* Normal binding table-based messages can't handle non-uniform access so
+ * we have to fall back to A64.
+ */
+ if (nir_intrinsic_access(intrin) & ACCESS_NON_UNIFORM)
+ return false;
+
+ if (!nir_deref_find_descriptor(deref, state))
+ return false;
+
+ nir_ssa_def *addr = build_index_offset_for_deref(deref, state);
+
+ b->cursor = nir_before_instr(&intrin->instr);
+ nir_lower_explicit_io_instr(b, intrin, addr,
+ nir_address_format_32bit_index_offset);
+ return true;
+}
+
+static void
+lower_direct_buffer_access(nir_function_impl *impl,
+ struct apply_pipeline_layout_state *state)
+{
+ nir_foreach_block(block, impl) {
+ nir_foreach_instr_safe(instr, block) {
+ if (instr->type != nir_instr_type_intrinsic)
+ continue;
+
+ nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+ switch (intrin->intrinsic) {
+ case nir_intrinsic_load_deref:
+ case nir_intrinsic_store_deref:
+ try_lower_direct_buffer_intrinsic(intrin, false, state);
+ break;
+ case nir_intrinsic_deref_atomic_add:
+ case nir_intrinsic_deref_atomic_imin:
+ case nir_intrinsic_deref_atomic_umin:
+ case nir_intrinsic_deref_atomic_imax:
+ case nir_intrinsic_deref_atomic_umax:
+ case nir_intrinsic_deref_atomic_and:
+ case nir_intrinsic_deref_atomic_or:
+ case nir_intrinsic_deref_atomic_xor:
+ case nir_intrinsic_deref_atomic_exchange:
+ case nir_intrinsic_deref_atomic_comp_swap:
+ case nir_intrinsic_deref_atomic_fmin:
+ case nir_intrinsic_deref_atomic_fmax:
+ case nir_intrinsic_deref_atomic_fcomp_swap:
+ try_lower_direct_buffer_intrinsic(intrin, true, state);
+ break;
+
+ case nir_intrinsic_get_buffer_size: {
+ /* The get_buffer_size intrinsic always just takes a
+ * index/reindex intrinsic.
+ */
+ if (!find_descriptor_for_index_src(intrin->src[0], state))
+ break;
+
+ nir_ssa_def *index =
+ build_index_for_res_reindex(nir_src_as_intrinsic(intrin->src[0]),
+ state);
+ nir_instr_rewrite_src(&intrin->instr, &intrin->src[0],
+ nir_src_for_ssa(index));
+ _mesa_set_add(state->lowered_instrs, intrin);
+ }
+
+ default:
+ break;
+ }
+ }
+ }
+}
+
+static nir_address_format
+desc_addr_format(VkDescriptorType desc_type,
+ struct apply_pipeline_layout_state *state)
+{
+ return (desc_type == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER ||
+ desc_type == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC) ?
+ state->ssbo_addr_format : nir_address_format_32bit_index_offset;
+}
+