zink: move image aoa access to nir pass

this brings zink's output nir closer to the emitted spirv

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30001>
This commit is contained in:
Mike Blumenkrantz
2024-07-04 10:28:20 -04:00
committed by Marge Bot
parent 586d0c4a9b
commit 235eb1af96
2 changed files with 53 additions and 20 deletions

View File

@@ -4053,26 +4053,6 @@ emit_deref_array(struct ntv_context *ctx, nir_deref_instr *deref)
if (itype == nir_type_float)
index = emit_bitcast(ctx, get_uvec_type(ctx, 32, 1), index);
if (var->data.mode == nir_var_uniform || var->data.mode == nir_var_image) {
nir_deref_instr *aoa_deref = nir_src_as_deref(deref->parent);
uint32_t inner_stride = glsl_array_size(aoa_deref->type);
while (aoa_deref->deref_type != nir_deref_type_var) {
assert(aoa_deref->deref_type == nir_deref_type_array);
SpvId aoa_index = get_src(ctx, &aoa_deref->arr.index, &itype);
if (itype == nir_type_float)
aoa_index = emit_bitcast(ctx, get_uvec_type(ctx, 32, 1), aoa_index);
aoa_deref = nir_src_as_deref(aoa_deref->parent);
uint32_t stride = glsl_get_aoa_size(aoa_deref->type) / inner_stride;
aoa_index = emit_binop(ctx, SpvOpIMul, get_uvec_type(ctx, 32, 1), aoa_index,
emit_uint_const(ctx, 32, stride));
index = emit_binop(ctx, SpvOpIAdd, get_uvec_type(ctx, 32, 1), index, aoa_index);
}
}
SpvId ptr_type = spirv_builder_type_pointer(&ctx->builder,
storage_class,
type);

View File

@@ -6108,6 +6108,58 @@ trivial_revectorize(nir_shader *nir)
return progress;
}
static bool
flatten_image_arrays_intr(struct nir_builder *b, nir_instr *instr, void *data)
{
if (instr->type != nir_instr_type_deref)
return false;
nir_deref_instr *deref = nir_instr_as_deref(instr);
if (deref->deref_type != nir_deref_type_array)
return false;
nir_deref_instr *parent = nir_deref_instr_parent(deref);
if (!parent || parent->deref_type != nir_deref_type_array)
return false;
nir_variable *var = nir_deref_instr_get_variable(deref);
const struct glsl_type *type = glsl_without_array(var->type);
if (type == var->type || (!glsl_type_is_sampler(type) && !glsl_type_is_image(type)))
return false;
nir_deref_instr *parent_parent = nir_deref_instr_parent(parent);
int parent_size = glsl_array_size(parent->type);
b->cursor = nir_after_instr(instr);
nir_deref_instr *new_deref = nir_build_deref_array(b, parent_parent, nir_iadd(b, nir_imul_imm(b, parent->arr.index.ssa, parent_size), deref->arr.index.ssa));
nir_def_rewrite_uses_after(&deref->def, &new_deref->def, &new_deref->instr);
_mesa_set_add(data, instr);
_mesa_set_add(data, &parent->instr);
return true;
}
static bool
flatten_image_arrays(nir_shader *nir)
{
bool progress = false;
nir_foreach_variable_with_modes(var, nir, nir_var_uniform | nir_var_image) {
const struct glsl_type *type = glsl_without_array(var->type);
if (!glsl_type_is_sampler(type) && !glsl_type_is_image(type))
continue;
if (type == var->type)
continue;
var->type = glsl_array_type(type, glsl_get_aoa_size(var->type), sizeof(void*));
progress = true;
}
struct set *deletions = _mesa_set_create(NULL, _mesa_hash_pointer, _mesa_key_pointer_equal);
progress |= nir_shader_instructions_pass(nir, flatten_image_arrays_intr, nir_metadata_dominance, deletions);
set_foreach_remove(deletions, he) {
nir_instr *instr = (void*)he->key;
nir_instr_remove_v(instr);
}
_mesa_set_destroy(deletions, NULL);
if (progress)
nir_fixup_deref_types(nir);
return progress;
}
struct zink_shader *
zink_shader_create(struct zink_screen *screen, struct nir_shader *nir)
{
@@ -6180,6 +6232,7 @@ zink_shader_init(struct zink_screen *screen, struct zink_shader *zs)
}
NIR_PASS_V(nir, nir_lower_io_to_scalar, nir_var_shader_in | nir_var_shader_out, NULL, NULL);
NIR_PASS_V(nir, flatten_image_arrays);
optimize_nir(nir, NULL, true);
nir_foreach_variable_with_modes(var, nir, nir_var_shader_in | nir_var_shader_out) {
if (glsl_type_is_image(var->type) || glsl_type_is_sampler(var->type)) {