diff --git a/src/compiler/nir/nir_lower_explicit_io.c b/src/compiler/nir/nir_lower_explicit_io.c index 87d9ce770ea..498a56312c5 100644 --- a/src/compiler/nir/nir_lower_explicit_io.c +++ b/src/compiler/nir/nir_lower_explicit_io.c @@ -118,6 +118,26 @@ nir_build_addr_iadd_imm(nir_builder *b, nir_def *addr, addr_get_offset_bit_size(addr, addr_format))); } +static nir_def * +build_addr_ushr_imm(nir_builder *b, nir_def *addr, + nir_address_format addr_format, unsigned shift) +{ + switch (addr_format) { + case nir_address_format_vec2_index_32bit_offset: + assert(addr->num_components == 3); + return nir_vector_insert_imm( + b, addr, nir_ushr_imm(b, nir_channel(b, addr, 2), shift), 2); + + case nir_address_format_32bit_index_offset: + assert(addr->num_components == 2); + return nir_vector_insert_imm( + b, addr, nir_ushr_imm(b, nir_channel(b, addr, 1), shift), 1); + + default: + UNREACHABLE("Unsupported address format"); + } +} + static nir_def * build_addr_for_var(nir_builder *b, nir_variable *var, nir_address_format addr_format) @@ -768,6 +788,9 @@ build_explicit_io_load(nir_builder *b, nir_intrinsic_instr *intrin, nir_intrinsic_set_range(load, range); } + if (addr_shift.shift) + nir_intrinsic_set_offset_shift(load, addr_shift.shift); + load->num_components = num_components; nir_def_init(&load->instr, &load->def, num_components, bit_size); @@ -966,6 +989,9 @@ build_explicit_io_store(nir_builder *b, nir_intrinsic_instr *intrin, nir_intrinsic_set_align(store, align_mul, align_offset); + if (addr_shift.shift) + nir_intrinsic_set_offset_shift(store, addr_shift.shift); + assert(value->num_components == 1 || value->num_components == intrin->num_components); store->num_components = value->num_components; @@ -1081,6 +1107,9 @@ build_explicit_io_atomic(nir_builder *b, nir_intrinsic_instr *intrin, if (nir_intrinsic_has_access(atomic)) nir_intrinsic_set_access(atomic, nir_intrinsic_access(intrin)); + if (addr_shift.shift) + nir_intrinsic_set_offset_shift(atomic, addr_shift.shift); + assert(intrin->def.num_components == 1); nir_def_init(&atomic->instr, &atomic->def, 1, intrin->def.bit_size); @@ -1102,9 +1131,14 @@ build_explicit_io_atomic(nir_builder *b, nir_intrinsic_instr *intrin, } } +/* The resulting address will be right-shifted by `shift` and any bits that + * got shifted-out by that will be put in `leftover`. So the final byte + * address will be: (ret << shift) + *leftover. + */ static nir_def * explicit_io_offset_from_deref(nir_builder *b, nir_deref_instr *deref, - unsigned offset_bit_size) + unsigned offset_bit_size, unsigned shift, + unsigned *leftover) { switch (deref->deref_type) { case nir_deref_type_var: @@ -1118,6 +1152,35 @@ explicit_io_offset_from_deref(nir_builder *b, nir_deref_instr *deref, nir_def *index = deref->arr.index.ssa; nir_def *offset; + unsigned max_stride_shift = ffs(stride) - 1; + unsigned stride_shift = MIN2(shift, max_stride_shift); + + if (stride_shift < shift) { + /* The stride isn't aligned enough to fully shift right. Try to apply + * the leftover shift to the index. We can only do this (without + * losing precision) if the index is constant. + */ + assert(nir_src_is_const(deref->arr.index)); + + unsigned index_shift = shift - stride_shift; + int64_t const_index = nir_src_as_int(deref->arr.index); + + if (!util_is_aligned(const_index, (uintmax_t)1 << index_shift)) { + assert(leftover); + + /* The index isn't aligned enough either. Just put the full offset + * in `leftover` and return zero. + */ + *leftover = stride * const_index; + return nir_imm_intN_t(b, 0, deref->arr.index.ssa->bit_size); + } + + index = nir_imm_intN_t(b, const_index >> index_shift, + deref->arr.index.ssa->bit_size); + } + + stride >>= stride_shift; + /* If the access chain has been declared in-bounds, then we know it doesn't * overflow the type. For nir_deref_type_array, this implies it cannot be * negative. Also, since types in NIR have a maximum 32-bit size, we know the @@ -1144,6 +1207,18 @@ explicit_io_offset_from_deref(nir_builder *b, nir_deref_instr *deref, int offset = glsl_get_struct_field_offset(parent->type, deref->strct.index); assert(offset >= 0); + + if (!util_is_aligned(offset, (uintmax_t)1 << shift)) { + assert(leftover); + + /* The offset isn't aligned enough to fully shift right. Just put it + * in leftover and return zero. + */ + *leftover = offset; + return nir_imm_intN_t(b, 0, offset_bit_size); + } + + offset >>= shift; return nir_imm_intN_t(b, offset, offset_bit_size); } @@ -1165,12 +1240,114 @@ nir_explicit_io_address_from_deref(nir_builder *b, nir_deref_instr *deref, } nir_def *offset = explicit_io_offset_from_deref( - b, deref, addr_get_offset_bit_size(base_addr, addr_format)); + b, deref, addr_get_offset_bit_size(base_addr, addr_format), 0, NULL); return offset ? nir_build_addr_iadd(b, base_addr, addr_format, deref->modes, offset) : base_addr; } +/* Walk the full deref chain and return the resulting address. The resulting + * address will be right-shifted by `shift` and any bits that got shifted-out + * by that will be put in `leftover`. So the final byte address will be: + * (ret << shift) + *leftover. + * + * Note on `shift` and `leftover`: we try to shift-right the resulting address + * by shifting-right the array strides and struct offsets. As long as `shift` + * isn't larger than the alignment, this should generally work without + * requiring the `leftover` value (i.e., all intermediary strides and offsets + * should be properly aligned). However, `leftover` is necessary in cases like + * this: + * + * struct { + * uint16_t a; + * uint8_t b; + * uint8_t c[3]; + * } s; + * use s.c[1]; + * + * s.c[1] is 2-byte aligned but neither s.c[]'s base offset (3) nor its stride + * (1) are. In this case, the returned address will be zero while `leftover` + * is set to 4. + */ +static nir_def * +explicit_io_address_from_deref_chain_aux(nir_builder *b, nir_deref_instr *deref, + nir_address_format addr_format, + nir_variable_mode modes, + unsigned shift, unsigned *leftover) +{ + unsigned offset_bit_size = + addr_get_offset_bit_size(&deref->def, addr_format); + + switch (deref->deref_type) { + case nir_deref_type_var: + UNREACHABLE("Unsupported deref type"); + return NULL; + + case nir_deref_type_cast: + if (nir_src_is_deref(deref->parent)) { + return explicit_io_address_from_deref_chain_aux( + b, nir_deref_instr_parent(deref), addr_format, modes, shift, + leftover); + } else { + /* For casts of non-deref instructions, there's nothing we can do + * besides simply right-shifting the result. As long as the cast's + * alignment is large enough, this should be fine. This is ensured in + * build_deref_addr by taking the alignment information into account. + * Note that it's correct to use a logical shift as the cast has to + * be in-bounds which means its result cannot be negative. + */ + return build_addr_ushr_imm(b, deref->parent.ssa, addr_format, shift); + } + + default: { + unsigned offset_leftover = 0; + nir_def *offset = explicit_io_offset_from_deref(b, deref, offset_bit_size, + shift, &offset_leftover); + assert(offset); + + unsigned parent_leftover = 0; + nir_def *parent_addr = explicit_io_address_from_deref_chain_aux( + b, nir_deref_instr_parent(deref), addr_format, modes, shift, + &parent_leftover); + + if (offset_leftover || parent_leftover) { + assert(leftover && *leftover == 0); + *leftover = offset_leftover + parent_leftover; + } + + return nir_build_addr_iadd(b, parent_addr, addr_format, modes, offset); + } + } +} + +static nir_def * +explicit_io_address_from_deref_chain(nir_builder *b, nir_deref_instr *deref, + nir_address_format addr_format, + nir_variable_mode modes, unsigned shift, + unsigned comp_offset) +{ + unsigned leftover = 0; + nir_def *addr = explicit_io_address_from_deref_chain_aux( + b, deref, addr_format, modes, shift, &leftover); + + unsigned extra_offset = comp_offset + leftover; + assert(util_is_aligned(extra_offset, (uintmax_t)1 << shift)); + + return nir_build_addr_iadd_imm(b, addr, addr_format, modes, + extra_offset >> shift); +} + +static unsigned +get_max_shift(nir_intrinsic_instr *intrin, + const nir_shader_compiler_options *options) +{ + if (options->max_offset_shift) { + return options->max_offset_shift(intrin, options->cb_data); + } + + return 0; +} + static nir_io_offset build_addr(nir_builder *b, nir_intrinsic_instr *intrin, nir_def *base_addr, nir_address_format addr_format, unsigned comp_offset, @@ -1180,9 +1357,28 @@ build_addr(nir_builder *b, nir_intrinsic_instr *intrin, nir_def *base_addr, assert(deref); nir_io_offset addr; - addr.def = nir_build_addr_iadd_imm(b, base_addr, addr_format, deref->modes, - comp_offset); - addr.shift = 0; + unsigned max_shift = get_max_shift(intrin, b->shader->options); + + if (max_shift) { + /* In order to calculate a shifted address, we have to walk the full + * deref chain. In that case, providing a base_addr that's not the + * actual deref probably won't produce the desired result. + */ + assert(base_addr == &deref->def); + + /* Don't try to shift more than the alignment would allow. This ensures + * we can just right-shift casts. + */ + unsigned align = nir_combined_align(align_mul, align_offset); + addr.shift = MIN2(max_shift, util_logbase2(align)); + addr.def = explicit_io_address_from_deref_chain( + b, deref, addr_format, deref->modes, addr.shift, comp_offset); + } else { + addr.def = nir_build_addr_iadd_imm(b, base_addr, addr_format, + deref->modes, comp_offset); + addr.shift = 0; + } + return addr; } @@ -1645,10 +1841,21 @@ nir_lower_explicit_io_impl(nir_function_impl *impl, nir_variable_mode modes, * consistent format, pointers can safely be conjured from thin air by the * driver, stored to variables, passed through phis, etc. * - * The one exception to the simple algorithm described above is for handling + * One exception to the simple algorithm described above is for handling * row-major matrices in which case we may look down one additional level of * the deref chain. * + * Another exception is when lowering accesses using offset_shift. The goal + * here is to generate addresses that are a right-shifted version of the + * actual byte address and record the shift amount in the offset_shift index. + * While we could just insert a ushr at the end of deref chains, this will + * prevent the shift to be optimized away in many cases. Instead, we try to + * extract the shift from the array strides and struct offsets that make up + * the deref chain, and only insert a ushr when absolutely necessary (i.e., + * for casts). This means we have to walk the entire deref chain at once for + * accesses that support offset_shift and we don't use the algorithm described + * above. + * * This pass is also capable of handling OpenCL generic pointers. If the * address mode is global, it will lower any ambiguous (more than one mode) * access to global and pass through the deref_mode_is run-time checks as diff --git a/src/compiler/nir/nir_shader_compiler_options.h b/src/compiler/nir/nir_shader_compiler_options.h index 575a368634d..423a1f3b517 100644 --- a/src/compiler/nir/nir_shader_compiler_options.h +++ b/src/compiler/nir/nir_shader_compiler_options.h @@ -840,6 +840,17 @@ typedef struct nir_shader_compiler_options { * the next shader according to default_varying_estimate_instr_cost. */ unsigned max_varying_expression_cost; + + /** + * Used by nir_lower_explicit_io to determine the maximum offset_shift to + * use when lowering the deref address of the given intrinsic. + */ + unsigned (*max_offset_shift)(nir_intrinsic_instr *, const void *); + + /** + * Passed to the callbacks that accept a data pointer. + */ + const void *cb_data; } nir_shader_compiler_options; #ifdef __cplusplus