intel/nir: Extract add_const_offset_to_base

Pretty much every driver using nir_lower_io_to_temporaries followed by
nir_lower_io is going to want this. In particular, radv and radeonsi in
the next commits.

Reviewed-by: Bas Nieuwenhuizen <[email protected]>

3 files changed, 100 insertions, 74 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index a2b3d393f12..55a39a493e9 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -3171,6 +3171,8 @@ bool nir_lower_io(nir_shader *shader,
                   int (*type_size)(const struct glsl_type *, bool),
                   nir_lower_io_options);
 
+bool nir_io_add_const_offset_to_base(nir_shader *nir, nir_variable_mode mode);
+
 typedef enum {
    /**
     * An address format which is a simple 32-bit global GPU address.
diff --git a/src/compiler/nir/nir_lower_io.c b/src/compiler/nir/nir_lower_io.c
index a9b3b95ef49..9275803d6f5 100644
--- a/src/compiler/nir/nir_lower_io.c
+++ b/src/compiler/nir/nir_lower_io.c
@@ -1307,3 +1307,78 @@ nir_build_addr_isub(nir_builder *b, nir_ssa_def *addr0, nir_ssa_def *addr1,
 
    unreachable("Invalid address format");
 }
+
+static bool
+is_input(nir_intrinsic_instr *intrin)
+{
+   return intrin->intrinsic == nir_intrinsic_load_input ||
+          intrin->intrinsic == nir_intrinsic_load_per_vertex_input ||
+          intrin->intrinsic == nir_intrinsic_load_interpolated_input ||
+          intrin->intrinsic == nir_intrinsic_load_fs_input_interp_deltas;
+}
+
+static bool
+is_output(nir_intrinsic_instr *intrin)
+{
+   return intrin->intrinsic == nir_intrinsic_load_output ||
+          intrin->intrinsic == nir_intrinsic_load_per_vertex_output ||
+          intrin->intrinsic == nir_intrinsic_store_output ||
+          intrin->intrinsic == nir_intrinsic_store_per_vertex_output;
+}
+
+
+/**
+ * This pass adds constant offsets to instr->const_index[0] for input/output
+ * intrinsics, and resets the offset source to 0.  Non-constant offsets remain
+ * unchanged - since we don't know what part of a compound variable is
+ * accessed, we allocate storage for the entire thing. For drivers that use
+ * nir_lower_io_to_temporaries() before nir_lower_io(), this guarantees that
+ * the offset source will be 0, so that they don't have to add it in manually.
+ */
+
+static bool
+add_const_offset_to_base_block(nir_block *block, nir_builder *b,
+                               nir_variable_mode mode)
+{
+   bool progress = false;
+   nir_foreach_instr_safe(instr, block) {
+      if (instr->type != nir_instr_type_intrinsic)
+         continue;
+
+      nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+
+      if ((mode == nir_var_shader_in && is_input(intrin)) ||
+          (mode == nir_var_shader_out && is_output(intrin))) {
+         nir_src *offset = nir_get_io_offset_src(intrin);
+
+         if (nir_src_is_const(*offset)) {
+            intrin->const_index[0] += nir_src_as_uint(*offset);
+            b->cursor = nir_before_instr(&intrin->instr);
+            nir_instr_rewrite_src(&intrin->instr, offset,
+                                  nir_src_for_ssa(nir_imm_int(b, 0)));
+            progress = true;
+         }
+      }
+   }
+
+   return progress;
+}
+
+bool
+nir_io_add_const_offset_to_base(nir_shader *nir, nir_variable_mode mode)
+{
+   bool progress = false;
+
+   nir_foreach_function(f, nir) {
+      if (f->impl) {
+         nir_builder b;
+         nir_builder_init(&b, f->impl);
+         nir_foreach_block(block, f->impl) {
+            progress |= add_const_offset_to_base_block(block, &b, mode);
+         }
+      }
+   }
+
+   return progress;
+}
+
diff --git a/src/intel/compiler/brw_nir.c b/src/intel/compiler/brw_nir.c
index 4dea234dd4b..2d46b87de6c 100644
--- a/src/intel/compiler/brw_nir.c
+++ b/src/intel/compiler/brw_nir.c
@@ -29,75 +29,6 @@
 #include "util/u_math.h"
 
 static bool
-is_input(nir_intrinsic_instr *intrin)
-{
-   return intrin->intrinsic == nir_intrinsic_load_input ||
-          intrin->intrinsic == nir_intrinsic_load_per_vertex_input ||
-          intrin->intrinsic == nir_intrinsic_load_interpolated_input ||
-          intrin->intrinsic == nir_intrinsic_load_fs_input_interp_deltas;
-}
-
-static bool
-is_output(nir_intrinsic_instr *intrin)
-{
-   return intrin->intrinsic == nir_intrinsic_load_output ||
-          intrin->intrinsic == nir_intrinsic_load_per_vertex_output ||
-          intrin->intrinsic == nir_intrinsic_store_output ||
-          intrin->intrinsic == nir_intrinsic_store_per_vertex_output;
-}
-
-/**
- * In many cases, we just add the base and offset together, so there's no
- * reason to keep them separate.  Sometimes, combining them is essential:
- * if a shader only accesses part of a compound variable (such as a matrix
- * or array), the variable's base may not actually exist in the VUE map.
- *
- * This pass adds constant offsets to instr->const_index[0], and resets
- * the offset source to 0.  Non-constant offsets remain unchanged - since
- * we don't know what part of a compound variable is accessed, we allocate
- * storage for the entire thing.
- */
-
-static bool
-add_const_offset_to_base_block(nir_block *block, nir_builder *b,
-                               nir_variable_mode mode)
-{
-   nir_foreach_instr_safe(instr, block) {
-      if (instr->type != nir_instr_type_intrinsic)
-         continue;
-
-      nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
-
-      if ((mode == nir_var_shader_in && is_input(intrin)) ||
-          (mode == nir_var_shader_out && is_output(intrin))) {
-         nir_src *offset = nir_get_io_offset_src(intrin);
-
-         if (nir_src_is_const(*offset)) {
-            intrin->const_index[0] += nir_src_as_uint(*offset);
-            b->cursor = nir_before_instr(&intrin->instr);
-            nir_instr_rewrite_src(&intrin->instr, offset,
-                                  nir_src_for_ssa(nir_imm_int(b, 0)));
-         }
-      }
-   }
-   return true;
-}
-
-static void
-add_const_offset_to_base(nir_shader *nir, nir_variable_mode mode)
-{
-   nir_foreach_function(f, nir) {
-      if (f->impl) {
-         nir_builder b;
-         nir_builder_init(&b, f->impl);
-         nir_foreach_block(block, f->impl) {
-            add_const_offset_to_base_block(block, &b, mode);
-         }
-      }
-   }
-}
-
-static bool
 remap_tess_levels(nir_builder *b, nir_intrinsic_instr *intr,
                   GLenum primitive_mode)
 {
@@ -153,6 +84,24 @@ remap_tess_levels(nir_builder *b, nir_intrinsic_instr *intr,
 }
 
 static bool
+is_input(nir_intrinsic_instr *intrin)
+{
+   return intrin->intrinsic == nir_intrinsic_load_input ||
+          intrin->intrinsic == nir_intrinsic_load_per_vertex_input ||
+          intrin->intrinsic == nir_intrinsic_load_interpolated_input;
+}
+
+static bool
+is_output(nir_intrinsic_instr *intrin)
+{
+   return intrin->intrinsic == nir_intrinsic_load_output ||
+          intrin->intrinsic == nir_intrinsic_load_per_vertex_output ||
+          intrin->intrinsic == nir_intrinsic_store_output ||
+          intrin->intrinsic == nir_intrinsic_store_per_vertex_output;
+}
+
+
+static bool
 remap_patch_urb_offsets(nir_block *block, nir_builder *b,
                         const struct brw_vue_map *vue_map,
                         GLenum tes_primitive_mode)
@@ -227,7 +176,7 @@ brw_nir_lower_vs_inputs(nir_shader *nir,
    /* This pass needs actual constants */
    nir_opt_constant_folding(nir);
 
-   add_const_offset_to_base(nir, nir_var_shader_in);
+   nir_io_add_const_offset_to_base(nir, nir_var_shader_in);
 
    brw_nir_apply_attribute_workarounds(nir, vs_attrib_wa_flags);
 
@@ -350,7 +299,7 @@ brw_nir_lower_vue_inputs(nir_shader *nir,
    /* This pass needs actual constants */
    nir_opt_constant_folding(nir);
 
-   add_const_offset_to_base(nir, nir_var_shader_in);
+   nir_io_add_const_offset_to_base(nir, nir_var_shader_in);
 
    nir_foreach_function(function, nir) {
       if (!function->impl)
@@ -401,7 +350,7 @@ brw_nir_lower_tes_inputs(nir_shader *nir, const struct brw_vue_map *vue_map)
    /* This pass needs actual constants */
    nir_opt_constant_folding(nir);
 
-   add_const_offset_to_base(nir, nir_var_shader_in);
+   nir_io_add_const_offset_to_base(nir, nir_var_shader_in);
 
    nir_foreach_function(function, nir) {
       if (function->impl) {
@@ -458,7 +407,7 @@ brw_nir_lower_fs_inputs(nir_shader *nir,
    /* This pass needs actual constants */
    nir_opt_constant_folding(nir);
 
-   add_const_offset_to_base(nir, nir_var_shader_in);
+   nir_io_add_const_offset_to_base(nir, nir_var_shader_in);
 }
 
 void
@@ -484,7 +433,7 @@ brw_nir_lower_tcs_outputs(nir_shader *nir, const struct brw_vue_map *vue_map,
    /* This pass needs actual constants */
    nir_opt_constant_folding(nir);
 
-   add_const_offset_to_base(nir, nir_var_shader_out);
+   nir_io_add_const_offset_to_base(nir, nir_var_shader_out);
 
    nir_foreach_function(function, nir) {
       if (function->impl) {