nir: Use derefs in nir_lower_samplers

We change glsl_to_nir to provide derefs for bot textures and samplers
while we're at it.  This makes the lowering much easier since we only
either replace sources or remove them.

Acked-by: Rob Clark <[email protected]>
Acked-by: Bas Nieuwenhuizen <[email protected]>
Acked-by: Dave Airlie <[email protected]>
Reviewed-by: Kenneth Graunke <[email protected]>

2 files changed, 88 insertions, 69 deletions

diff --git a/src/compiler/glsl/gl_nir_lower_samplers.c b/src/compiler/glsl/gl_nir_lower_samplers.c
index a55652c5ae1..4847db9bfc6 100644
--- a/src/compiler/glsl/gl_nir_lower_samplers.c
+++ b/src/compiler/glsl/gl_nir_lower_samplers.c
@@ -35,99 +35,116 @@
  * calculate the base uniform location for struct members.
  */
 static void
-calc_sampler_offsets(nir_deref *tail, nir_tex_instr *instr,
-                     unsigned *array_elements, nir_ssa_def **indirect,
-                     nir_builder *b, unsigned *location)
+calc_sampler_offsets(nir_builder *b, nir_ssa_def *ptr,
+                     const struct gl_shader_program *shader_program,
+                     unsigned *base_index, nir_ssa_def **index,
+                     unsigned *array_elements)
 {
-   if (tail->child == NULL)
-      return;
-
-   switch (tail->child->deref_type) {
-   case nir_deref_type_array: {
-      nir_deref_array *deref_array = nir_deref_as_array(tail->child);
+   *base_index = 0;
+   *index = NULL;
+   *array_elements = 1;
+   unsigned location = 0;
+
+   nir_deref_instr *deref = nir_instr_as_deref(ptr->parent_instr);
+   while (deref->deref_type != nir_deref_type_var) {
+      assert(deref->parent.is_ssa);
+      nir_deref_instr *parent =
+         nir_instr_as_deref(deref->parent.ssa->parent_instr);
+
+      switch (deref->deref_type) {
+      case nir_deref_type_struct:
+         location += glsl_get_record_location_offset(parent->type,
+                                                     deref->strct.index);
+         break;
+
+      case nir_deref_type_array: {
+         nir_const_value *const_deref_index =
+            nir_src_as_const_value(deref->arr.index);
+
+         if (const_deref_index && *index == NULL) {
+            /* We're still building a direct index */
+            *base_index += const_deref_index->u32[0] * *array_elements;
+         } else {
+            if (*index == NULL) {
+               /* We used to be direct but not anymore */
+               *index = nir_imm_int(b, *base_index);
+               *base_index = 0;
+            }
+
+            *index = nir_iadd(b, *index,
+                     nir_imul(b, nir_imm_int(b, *array_elements),
+                              nir_ssa_for_src(b, deref->arr.index, 1)));
+         }
 
-      assert(deref_array->deref_array_type != nir_deref_array_type_wildcard);
+         *array_elements *= glsl_get_length(parent->type);
+         break;
+      }
 
-      calc_sampler_offsets(tail->child, instr, array_elements,
-                           indirect, b, location);
-      instr->texture_index += deref_array->base_offset * *array_elements;
+      default:
+         unreachable("Invalid sampler deref type");
+      }
 
-      if (deref_array->deref_array_type == nir_deref_array_type_indirect) {
-         nir_ssa_def *mul =
-            nir_imul(b, nir_imm_int(b, *array_elements),
-                     nir_ssa_for_src(b, deref_array->indirect, 1));
+      deref = parent;
+   }
 
-         nir_instr_rewrite_src(&instr->instr, &deref_array->indirect,
-                               NIR_SRC_INIT);
+   if (*index)
+      *index = nir_umin(b, *index, nir_imm_int(b, *array_elements - 1));
 
-         if (*indirect) {
-            *indirect = nir_iadd(b, *indirect, mul);
-         } else {
-            *indirect = mul;
-         }
-      }
+   /* We hit the deref_var.  This is the end of the line */
+   assert(deref->deref_type == nir_deref_type_var);
 
-      *array_elements *= glsl_get_length(tail->type);
-       break;
-   }
+   location += deref->var->data.location;
 
-   case nir_deref_type_struct: {
-      nir_deref_struct *deref_struct = nir_deref_as_struct(tail->child);
-      *location += glsl_get_record_location_offset(tail->type, deref_struct->index);
-      calc_sampler_offsets(tail->child, instr, array_elements,
-                           indirect, b, location);
-      break;
-   }
+   gl_shader_stage stage = b->shader->info.stage;
+   assert(location < shader_program->data->NumUniformStorage &&
+          shader_program->data->UniformStorage[location].opaque[stage].active);
 
-   default:
-      unreachable("Invalid deref type");
-      break;
-   }
+   *base_index +=
+      shader_program->data->UniformStorage[location].opaque[stage].index;
 }
 
 static bool
 lower_sampler(nir_builder *b, nir_tex_instr *instr,
               const struct gl_shader_program *shader_program)
 {
-   if (instr->texture == NULL)
-      return false;
+   int texture_idx =
+      nir_tex_instr_src_index(instr, nir_tex_src_texture_deref);
+   int sampler_idx =
+      nir_tex_instr_src_index(instr, nir_tex_src_sampler_deref);
 
-   /* In GLSL, we only fill out the texture field.  The sampler is inferred */
-   assert(instr->sampler == NULL || shader_program->data->spirv);
+   if (texture_idx < 0)
+      return false;
 
-   instr->texture_index = 0;
-   unsigned location = instr->texture->var->data.location;
-   unsigned array_elements = 1;
-   nir_ssa_def *indirect = NULL;
+   assert(texture_idx >= 0 && sampler_idx >= 0);
+   assert(instr->src[texture_idx].src.is_ssa);
+   assert(instr->src[sampler_idx].src.is_ssa);
+   assert(instr->src[texture_idx].src.ssa == instr->src[sampler_idx].src.ssa);
 
    b->cursor = nir_before_instr(&instr->instr);
-   calc_sampler_offsets(&instr->texture->deref, instr, &array_elements,
-                        &indirect, b, &location);
 
-   if (indirect) {
-      assert(array_elements >= 1);
-      indirect = nir_umin(b, indirect, nir_imm_int(b, array_elements - 1));
+   unsigned base_offset, array_elements;
+   nir_ssa_def *indirect;
+   calc_sampler_offsets(b, instr->src[texture_idx].src.ssa, shader_program,
+                        &base_offset, &indirect, &array_elements);
 
-      nir_tex_instr_add_src(instr, nir_tex_src_texture_offset,
+   instr->texture_index = base_offset;
+   instr->sampler_index = base_offset;
+   if (indirect) {
+      nir_instr_rewrite_src(&instr->instr, &instr->src[texture_idx].src,
                             nir_src_for_ssa(indirect));
-      nir_tex_instr_add_src(instr, nir_tex_src_sampler_offset,
+      instr->src[texture_idx].src_type = nir_tex_src_texture_offset;
+      nir_instr_rewrite_src(&instr->instr, &instr->src[sampler_idx].src,
                             nir_src_for_ssa(indirect));
+      instr->src[sampler_idx].src_type = nir_tex_src_sampler_offset;
 
       instr->texture_array_size = array_elements;
+   } else {
+      nir_tex_instr_remove_src(instr, texture_idx);
+      /* The sampler index may have changed */
+      sampler_idx = nir_tex_instr_src_index(instr, nir_tex_src_sampler_deref);
+      nir_tex_instr_remove_src(instr, sampler_idx);
    }
 
-   gl_shader_stage stage = b->shader->info.stage;
-   assert(location < shader_program->data->NumUniformStorage &&
-          shader_program->data->UniformStorage[location].opaque[stage].active);
-
-   instr->texture_index +=
-      shader_program->data->UniformStorage[location].opaque[stage].index;
-
-   instr->sampler_index = instr->texture_index;
-
-   instr->texture = NULL;
-   nir_instr_rewrite_deref(&instr->instr, &instr->sampler, NULL);
-
    return true;
 }
 
@@ -156,7 +173,7 @@ gl_nir_lower_samplers(nir_shader *shader,
 {
    bool progress = false;
 
-   nir_assert_lowered_derefs(shader, nir_lower_texture_derefs);
+   nir_assert_unlowered_derefs(shader, nir_lower_texture_derefs);
 
    nir_foreach_function(function, shader) {
       if (function->impl)
diff --git a/src/compiler/glsl/glsl_to_nir.cpp b/src/compiler/glsl/glsl_to_nir.cpp
index b2dbaf89d44..d3a3fb9b085 100644
--- a/src/compiler/glsl/glsl_to_nir.cpp
+++ b/src/compiler/glsl/glsl_to_nir.cpp
@@ -2010,7 +2010,7 @@ nir_visitor::visit(ir_texture *ir)
       num_srcs++;
 
    /* Add one for the texture deref */
-   num_srcs += 1;
+   num_srcs += 2;
 
    nir_tex_instr *instr = nir_tex_instr_create(this->shader, num_srcs);
 
@@ -2039,8 +2039,10 @@ nir_visitor::visit(ir_texture *ir)
    nir_deref_instr *sampler_deref = evaluate_deref(ir->sampler);
    instr->src[0].src = nir_src_for_ssa(&sampler_deref->dest.ssa);
    instr->src[0].src_type = nir_tex_src_texture_deref;
+   instr->src[1].src = nir_src_for_ssa(&sampler_deref->dest.ssa);
+   instr->src[1].src_type = nir_tex_src_sampler_deref;
 
-   unsigned src_number = 1;
+   unsigned src_number = 2;
 
    if (ir->coordinate != NULL) {
       instr->coord_components = ir->coordinate->type->vector_elements;