From bd89279dd4c0470620ba8574638235a5bd4d966e Mon Sep 17 00:00:00 2001 From: Alyssa Rosenzweig Date: Mon, 9 Dec 2024 17:05:54 -0500 Subject: [PATCH] nir: add lower_scratch_to_var pass to ease opencl pain. Signed-off-by: Alyssa Rosenzweig Reviewed-by: Mary Guillemard Part-of: --- src/compiler/nir/meson.build | 1 + src/compiler/nir/nir.h | 2 + src/compiler/nir/nir_lower_scratch_to_var.c | 115 ++++++++++++++++++++ src/compiler/nir/nir_opcodes.py | 3 + src/compiler/nir/nir_opt_algebraic.py | 9 ++ 5 files changed, 130 insertions(+) create mode 100644 src/compiler/nir/nir_lower_scratch_to_var.c diff --git a/src/compiler/nir/meson.build b/src/compiler/nir/meson.build index 4b7f91dbdd2..c2173157b3b 100644 --- a/src/compiler/nir/meson.build +++ b/src/compiler/nir/meson.build @@ -188,6 +188,7 @@ files_libnir = files( 'nir_lower_robust_access.c', 'nir_lower_samplers.c', 'nir_lower_scratch.c', + 'nir_lower_scratch_to_var.c', 'nir_lower_shader_calls.c', 'nir_lower_single_sampled.c', 'nir_lower_ssbo.c', diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h index 5379442a122..1eea39602fd 100644 --- a/src/compiler/nir/nir.h +++ b/src/compiler/nir/nir.h @@ -5597,6 +5597,8 @@ bool nir_lower_vars_to_scratch(nir_shader *shader, glsl_type_size_align_func variable_size_align, glsl_type_size_align_func scratch_layout_size_align); +bool nir_lower_scratch_to_var(nir_shader *nir); + void nir_lower_clip_halfz(nir_shader *shader); void nir_shader_gather_info(nir_shader *shader, nir_function_impl *entrypoint); diff --git a/src/compiler/nir/nir_lower_scratch_to_var.c b/src/compiler/nir/nir_lower_scratch_to_var.c new file mode 100644 index 00000000000..0e01c960f6f --- /dev/null +++ b/src/compiler/nir/nir_lower_scratch_to_var.c @@ -0,0 +1,115 @@ +/* + * Copyright 2024 Valve Corporation + * SPDX-License-Identifier: MIT + */ +#include "nir.h" +#include "nir_builder.h" + +/* + * It is challenging to optimize the complex deref chains resulting from + * nontrivial OpenCL C constructs. nir_opt_deref generally does a good job, but + * occassionally we are forced to lower temporaries to scratch anyway. LLVM's + * recent embrace of opaque pointers have exacerbated this problem. + * + * The "proper" solutions here are to smarten nir_opt_deref and/or to use LLVM's + * own optimization passes to clean up the input IR. Both of these are + * challenging projects for the medium-term. + * + * In the short term, this pass is a stopgap. After lowering away all derefs to + * scratch, this pass can "unlower" scratch memory back into nir_variable + * access. The lower->unlower pair is lossy. The point is not to reconstruct the + * original derefs (that we failed to optimize), but instead just to model array + * access that other NIR passes can optimize. The resulting array accesses will + * generally optimize out if there are no indirects, or can be lowered to bcsel + * instead of scratch if that's preferable for a driver. + */ + +/* + * This pass operates only on 32-bit scalars, so this callback instructs + * nir_lower_mem_access_bit_sizes_options to turn all scratch access into + * 32-bit scalars. We don't want to use 8-bit accesses, since that would be + * challenging to optimize the resulting pack/unpack on some drivers. Larger + * 32-bit access however requires nontrivial tracking to extract/insert. Since + * nir_lower_mem_access_bit_sizes already has that code, we use it in this pass + * instead of NIH'ing it here. + */ +static nir_mem_access_size_align +mem_access_cb(nir_intrinsic_op intrin, uint8_t bytes, uint8_t bit_size, + uint32_t align, uint32_t align_offset, bool offset_is_const, + enum gl_access_qualifier access, const void *cb_data) +{ + return (nir_mem_access_size_align){ + .num_components = 1, + .bit_size = 32, + .align = 4, + .shift = nir_mem_access_shift_method_scalar, + }; +} + +/* + * Thanks to nir_lower_mem_access_bit_sizes, we can lower scratch intrinsics 1:1 + * to word-based array access. + */ +static bool +lower_scratch_to_var(nir_builder *b, nir_intrinsic_instr *intr, void *data) +{ + nir_variable *scratch = data; + b->cursor = nir_before_instr(&intr->instr); + + if (intr->intrinsic == nir_intrinsic_store_scratch) { + nir_def *index = nir_udiv_aligned_4(b, intr->src[1].ssa); + nir_def *value = intr->src[0].ssa; + + nir_store_array_var(b, scratch, index, value, nir_component_mask(1)); + } else if (intr->intrinsic == nir_intrinsic_load_scratch) { + nir_def *index = nir_udiv_aligned_4(b, intr->src[0].ssa); + nir_def_rewrite_uses(&intr->def, nir_load_array_var(b, scratch, index)); + } else { + return false; + } + + nir_instr_remove(&intr->instr); + return true; +} + +bool +nir_lower_scratch_to_var(nir_shader *nir) +{ + unsigned words = DIV_ROUND_UP(nir->scratch_size, 4); + + /* Early exit in the common case that scratch is not used. */ + if (words == 0) { + return false; + } + + /* First, lower bit sizes and vectors as required by lower_scratch_to_var */ + nir_lower_mem_access_bit_sizes_options lower_mem_access_options = { + .modes = nir_var_shader_temp | nir_var_function_temp, + .callback = mem_access_cb, + }; + NIR_PASS(_, nir, nir_lower_mem_access_bit_sizes, &lower_mem_access_options); + + /* Then, back scratch by an array of words and turn all scratch access into + * array access. + */ + nir_function_impl *entry = nir_shader_get_entrypoint(nir); + const glsl_type *type_ = glsl_array_type(glsl_uint_type(), words, 1); + nir_variable *var = nir_local_variable_create(entry, type_, "scratch"); + nir_shader_intrinsics_pass(nir, lower_scratch_to_var, + nir_metadata_control_flow, var); + nir->scratch_size = 0; + + /* Now clean up the mess we made */ + bool progress; + do { + progress = false; + NIR_PASS(progress, nir, nir_lower_vars_to_ssa); + NIR_PASS(progress, nir, nir_opt_constant_folding); + NIR_PASS(progress, nir, nir_opt_algebraic); + NIR_PASS(progress, nir, nir_copy_prop); + NIR_PASS(progress, nir, nir_opt_cse); + NIR_PASS(progress, nir, nir_opt_dce); + } while (progress); + + return true; +} diff --git a/src/compiler/nir/nir_opcodes.py b/src/compiler/nir/nir_opcodes.py index e315022c49b..d7e302ef619 100644 --- a/src/compiler/nir/nir_opcodes.py +++ b/src/compiler/nir/nir_opcodes.py @@ -872,6 +872,9 @@ opcode("ushr", 0, tuint, [0, 0], [tuint, tuint32], False, "", "src0 >> (src1 & (sizeof(src0) * 8 - 1))", description = "Unsigned right-shift." + shift_note) +opcode("udiv_aligned_4", 0, tuint, [0], [tuint], False, "", + "src0 >> 2", description = "Divide a multiple of 4 by 4") + opcode("urol", 0, tuint, [0, 0], [tuint, tuint32], False, "", """ uint32_t rotate_mask = sizeof(src0) * 8 - 1; dst = (src0 << (src1 & rotate_mask)) | diff --git a/src/compiler/nir/nir_opt_algebraic.py b/src/compiler/nir/nir_opt_algebraic.py index e8172690d41..de6a4dd622d 100644 --- a/src/compiler/nir/nir_opt_algebraic.py +++ b/src/compiler/nir/nir_opt_algebraic.py @@ -2539,6 +2539,11 @@ optimizations.extend([ # rule converts everyone else to imul: (('amul', a, b), ('imul', a, b), '!options->has_imul24 && !options->has_amul'), + # udiv_aligned_4 assumes the source is a multiple of 4 specifically to enable + # this identity. Usually this transform would require masking. + (('amul', ('udiv_aligned_4', a), 4), a), + (('imul', ('udiv_aligned_4', a), 4), a), + (('umul24', a, b), ('imul', ('iand', a, 0xffffff), ('iand', b, 0xffffff)), '!options->has_umul24'), @@ -3300,6 +3305,10 @@ late_optimizations = [ # combine imul and iadd to imad (('iadd@32', ('imul(is_only_used_by_iadd)', a, b), c), ('imad', a, b, c), 'options->has_imad32'), + + # Drivers do not actually implement udiv_aligned_4, it is just used to + # optimize scratch lowering. + (('udiv_aligned_4', a), ('ushr', a, 2)), ] # re-combine inexact mul+add to ffma. Do this before fsub so that a * b - c