nir/deref: Add helpers for getting offsets

These are very similar to the related function in nir_lower_io except
that they don't handle per-vertex or packed things (that could be added,
in theory) and they take a more detailed size/align function pointer.
One day, we should consider switching nir_lower_io over to using the
more detailed size/align functions and then we could make it use these
helpers instead of having its own.

Reviewed-by: Timothy Arceri <[email protected]>
Reviewed-by: Iago Toral Quiroga <[email protected]>
Reviewed-by: Kenneth Graunke <[email protected]>

2 files changed, 95 insertions, 0 deletions

diff --git a/src/compiler/nir/nir_deref.c b/src/compiler/nir/nir_deref.c
index 22ecde4ecca..c03acf83597 100644
--- a/src/compiler/nir/nir_deref.c
+++ b/src/compiler/nir/nir_deref.c
@@ -120,6 +120,95 @@ nir_deref_instr_has_indirect(nir_deref_instr *instr)
    return false;
 }
 
+static unsigned
+type_get_array_stride(const struct glsl_type *elem_type,
+                      glsl_type_size_align_func size_align)
+{
+   unsigned elem_size, elem_align;
+   glsl_get_natural_size_align_bytes(elem_type, &elem_size, &elem_align);
+   return ALIGN_POT(elem_size, elem_align);
+}
+
+static unsigned
+struct_type_get_field_offset(const struct glsl_type *struct_type,
+                             glsl_type_size_align_func size_align,
+                             unsigned field_idx)
+{
+   assert(glsl_type_is_struct(struct_type));
+   unsigned offset = 0;
+   for (unsigned i = 0; i <= field_idx; i++) {
+      unsigned elem_size, elem_align;
+      glsl_get_natural_size_align_bytes(glsl_get_struct_field(struct_type, i),
+                                        &elem_size, &elem_align);
+      offset = ALIGN_POT(offset, elem_align);
+      if (i < field_idx)
+         offset += elem_size;
+   }
+   return offset;
+}
+
+unsigned
+nir_deref_instr_get_const_offset(nir_deref_instr *deref,
+                                 glsl_type_size_align_func size_align)
+{
+   nir_deref_path path;
+   nir_deref_path_init(&path, deref, NULL);
+
+   assert(path.path[0]->deref_type == nir_deref_type_var);
+
+   unsigned offset = 0;
+   for (nir_deref_instr **p = &path.path[1]; *p; p++) {
+      if ((*p)->deref_type == nir_deref_type_array) {
+         offset += nir_src_as_const_value((*p)->arr.index)->u32[0] *
+                   type_get_array_stride((*p)->type, size_align);
+      } else if ((*p)->deref_type == nir_deref_type_struct) {
+         /* p starts at path[1], so this is safe */
+         nir_deref_instr *parent = *(p - 1);
+         offset += struct_type_get_field_offset(parent->type, size_align,
+                                                (*p)->strct.index);
+      } else {
+         unreachable("Unsupported deref type");
+      }
+   }
+
+   nir_deref_path_finish(&path);
+
+   return offset;
+}
+
+nir_ssa_def *
+nir_build_deref_offset(nir_builder *b, nir_deref_instr *deref,
+                       glsl_type_size_align_func size_align)
+{
+   nir_deref_path path;
+   nir_deref_path_init(&path, deref, NULL);
+
+   assert(path.path[0]->deref_type == nir_deref_type_var);
+
+   nir_ssa_def *offset = nir_imm_int(b, 0);
+   for (nir_deref_instr **p = &path.path[1]; *p; p++) {
+      if ((*p)->deref_type == nir_deref_type_array) {
+         nir_ssa_def *index = nir_ssa_for_src(b, (*p)->arr.index, 1);
+         nir_ssa_def *stride =
+            nir_imm_int(b, type_get_array_stride((*p)->type, size_align));
+         offset = nir_iadd(b, offset, nir_imul(b, index, stride));
+      } else if ((*p)->deref_type == nir_deref_type_struct) {
+         /* p starts at path[1], so this is safe */
+         nir_deref_instr *parent = *(p - 1);
+         unsigned field_offset =
+            struct_type_get_field_offset(parent->type, size_align,
+                                         (*p)->strct.index);
+         nir_iadd(b, offset, nir_imm_int(b, field_offset));
+      } else {
+         unreachable("Unsupported deref type");
+      }
+   }
+
+   nir_deref_path_finish(&path);
+
+   return offset;
+}
+
 bool
 nir_remove_dead_derefs_impl(nir_function_impl *impl)
 {
diff --git a/src/compiler/nir/nir_deref.h b/src/compiler/nir/nir_deref.h
index 0980bae7215..6f4141aaf82 100644
--- a/src/compiler/nir/nir_deref.h
+++ b/src/compiler/nir/nir_deref.h
@@ -48,6 +48,12 @@ void nir_deref_path_init(nir_deref_path *path,
                          nir_deref_instr *deref, void *mem_ctx);
 void nir_deref_path_finish(nir_deref_path *path);
 
+unsigned nir_deref_instr_get_const_offset(nir_deref_instr *deref,
+                                          glsl_type_size_align_func size_align);
+
+nir_ssa_def *nir_build_deref_offset(nir_builder *b, nir_deref_instr *deref,
+                                    glsl_type_size_align_func size_align);
+
 #ifdef __cplusplus
 } /* extern "C" */
 #endif