anv: Lower some SSBO operations in apply_pipeline_layout

In order to avoid the potential overhead of A64 operations on all SSBO
ops, we look for those SSBO ops where we can get to the descriptor set
from the SSBO access operation and lower those to a binding-table
approach.  When robustBufferAccess is enabled, this lets the hardware do
the bounds checking for us.  It also avoids some potentially expensive
64-bit integer calculations.

Reviewed-by: Lionel Landwerlin <[email protected]>
Reviewed-by: Caio Marcelo de Oliveira Filho <[email protected]>

1 files changed, 212 insertions, 2 deletions

diff --git a/src/intel/vulkan/anv_nir_apply_pipeline_layout.c b/src/intel/vulkan/anv_nir_apply_pipeline_layout.c
index ea02ed1be78..7abc27be103 100644
--- a/src/intel/vulkan/anv_nir_apply_pipeline_layout.c
+++ b/src/intel/vulkan/anv_nir_apply_pipeline_layout.c
@@ -25,6 +25,7 @@
 #include "program/prog_parameter.h"
 #include "nir/nir_builder.h"
 #include "compiler/brw_nir.h"
+#include "util/set.h"
 
 /* Sampler tables don't actually have a maximum size but we pick one just so
  * that we don't end up emitting too much state on-the-fly.
@@ -41,6 +42,9 @@ struct apply_pipeline_layout_state {
    struct anv_pipeline_layout *layout;
    bool add_bounds_checks;
 
+   /* Place to flag lowered instructions so we don't lower them twice */
+   struct set *lowered_instrs;
+
    bool uses_constants;
    uint8_t constants_offset;
    struct {
@@ -143,6 +147,176 @@ get_used_bindings_block(nir_block *block,
    }
 }
 
+static bool
+find_descriptor_for_index_src(nir_src src,
+                              struct apply_pipeline_layout_state *state)
+{
+   nir_intrinsic_instr *intrin = nir_src_as_intrinsic(src);
+
+   while (intrin && intrin->intrinsic == nir_intrinsic_vulkan_resource_reindex)
+      intrin = nir_src_as_intrinsic(intrin->src[0]);
+
+   if (!intrin || intrin->intrinsic != nir_intrinsic_vulkan_resource_index)
+      return false;
+
+   return true;
+}
+
+static bool
+nir_deref_find_descriptor(nir_deref_instr *deref,
+                          struct apply_pipeline_layout_state *state)
+{
+   while (1) {
+      /* Nothing we will use this on has a variable */
+      assert(deref->deref_type != nir_deref_type_var);
+
+      nir_deref_instr *parent = nir_src_as_deref(deref->parent);
+      if (!parent)
+         break;
+
+      deref = parent;
+   }
+   assert(deref->deref_type == nir_deref_type_cast);
+
+   nir_intrinsic_instr *intrin = nir_src_as_intrinsic(deref->parent);
+   if (!intrin || intrin->intrinsic != nir_intrinsic_load_vulkan_descriptor)
+      return false;
+
+   return find_descriptor_for_index_src(intrin->src[0], state);
+}
+
+static nir_ssa_def *
+build_index_for_res_reindex(nir_intrinsic_instr *intrin,
+                            struct apply_pipeline_layout_state *state)
+{
+   nir_builder *b = &state->builder;
+
+   if (intrin->intrinsic == nir_intrinsic_vulkan_resource_reindex) {
+      nir_ssa_def *bti =
+         build_index_for_res_reindex(nir_src_as_intrinsic(intrin->src[0]), state);
+
+      b->cursor = nir_before_instr(&intrin->instr);
+      return nir_iadd(b, bti, nir_ssa_for_src(b, intrin->src[1], 1));
+   }
+
+   assert(intrin->intrinsic == nir_intrinsic_vulkan_resource_index);
+
+   uint32_t set = nir_intrinsic_desc_set(intrin);
+   uint32_t binding = nir_intrinsic_binding(intrin);
+
+   const struct anv_descriptor_set_binding_layout *bind_layout =
+      &state->layout->set[set].layout->binding[binding];
+
+   uint32_t surface_index = state->set[set].surface_offsets[binding];
+   uint32_t array_size = bind_layout->array_size;
+
+   b->cursor = nir_before_instr(&intrin->instr);
+
+   nir_ssa_def *array_index = nir_ssa_for_src(b, intrin->src[0], 1);
+   if (nir_src_is_const(intrin->src[0]) || state->add_bounds_checks)
+      array_index = nir_umin(b, array_index, nir_imm_int(b, array_size - 1));
+
+   return nir_iadd_imm(b, array_index, surface_index);
+}
+
+static nir_ssa_def *
+build_index_offset_for_deref(nir_deref_instr *deref,
+                             struct apply_pipeline_layout_state *state)
+{
+   nir_builder *b = &state->builder;
+
+   nir_deref_instr *parent = nir_deref_instr_parent(deref);
+   if (parent) {
+      nir_ssa_def *addr = build_index_offset_for_deref(parent, state);
+
+      b->cursor = nir_before_instr(&deref->instr);
+      return nir_explicit_io_address_from_deref(b, deref, addr,
+                                                nir_address_format_32bit_index_offset);
+   }
+
+   nir_intrinsic_instr *load_desc = nir_src_as_intrinsic(deref->parent);
+   assert(load_desc->intrinsic == nir_intrinsic_load_vulkan_descriptor);
+
+   nir_ssa_def *index =
+      build_index_for_res_reindex(nir_src_as_intrinsic(load_desc->src[0]), state);
+
+   /* Return a 0 offset which will get picked up by the recursion */
+   b->cursor = nir_before_instr(&deref->instr);
+   return nir_vec2(b, index, nir_imm_int(b, 0));
+}
+
+static bool
+try_lower_direct_buffer_intrinsic(nir_intrinsic_instr *intrin,
+                                  struct apply_pipeline_layout_state *state)
+{
+   nir_builder *b = &state->builder;
+
+   nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
+   if (deref->mode != nir_var_mem_ssbo)
+      return false;
+
+   if (!nir_deref_find_descriptor(deref, state))
+      return false;
+
+   nir_ssa_def *addr = build_index_offset_for_deref(deref, state);
+
+   b->cursor = nir_before_instr(&intrin->instr);
+   nir_lower_explicit_io_instr(b, intrin, addr,
+                               nir_address_format_32bit_index_offset);
+   return true;
+}
+
+static void
+lower_direct_buffer_access(nir_function_impl *impl,
+                           struct apply_pipeline_layout_state *state)
+{
+   nir_foreach_block(block, impl) {
+      nir_foreach_instr_safe(instr, block) {
+         if (instr->type != nir_instr_type_intrinsic)
+            continue;
+
+         nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+         switch (intrin->intrinsic) {
+         case nir_intrinsic_load_deref:
+         case nir_intrinsic_store_deref:
+         case nir_intrinsic_deref_atomic_add:
+         case nir_intrinsic_deref_atomic_imin:
+         case nir_intrinsic_deref_atomic_umin:
+         case nir_intrinsic_deref_atomic_imax:
+         case nir_intrinsic_deref_atomic_umax:
+         case nir_intrinsic_deref_atomic_and:
+         case nir_intrinsic_deref_atomic_or:
+         case nir_intrinsic_deref_atomic_xor:
+         case nir_intrinsic_deref_atomic_exchange:
+         case nir_intrinsic_deref_atomic_comp_swap:
+         case nir_intrinsic_deref_atomic_fmin:
+         case nir_intrinsic_deref_atomic_fmax:
+         case nir_intrinsic_deref_atomic_fcomp_swap:
+            try_lower_direct_buffer_intrinsic(intrin, state);
+            break;
+
+         case nir_intrinsic_get_buffer_size: {
+            /* The get_buffer_size intrinsic always just takes a
+             * index/reindex intrinsic.
+             */
+            if (!find_descriptor_for_index_src(intrin->src[0], state))
+               break;
+
+            nir_ssa_def *index =
+               build_index_for_res_reindex(nir_src_as_intrinsic(intrin->src[0]),
+                                           state);
+            nir_instr_rewrite_src(&intrin->instr, &intrin->src[0],
+                                  nir_src_for_ssa(index));
+            _mesa_set_add(state->lowered_instrs, intrin);
+         }
+
+         default:
+            break;
+         }
+      }
+   }
+}
+
 static void
 lower_res_index_intrinsic(nir_intrinsic_instr *intrin,
                           struct apply_pipeline_layout_state *state)
@@ -228,6 +402,9 @@ static void
 lower_get_buffer_size(nir_intrinsic_instr *intrin,
                       struct apply_pipeline_layout_state *state)
 {
+   if (_mesa_set_search(state->lowered_instrs, intrin))
+      return;
+
    nir_builder *b = &state->builder;
 
    b->cursor = nir_before_instr(&intrin->instr);
@@ -539,15 +716,16 @@ anv_nir_apply_pipeline_layout(const struct anv_physical_device *pdevice,
                               struct brw_stage_prog_data *prog_data,
                               struct anv_pipeline_bind_map *map)
 {
+   void *mem_ctx = ralloc_context(NULL);
+
    struct apply_pipeline_layout_state state = {
       .pdevice = pdevice,
       .shader = shader,
       .layout = layout,
       .add_bounds_checks = robust_buffer_access,
+      .lowered_instrs = _mesa_pointer_set_create(mem_ctx),
    };
 
-   void *mem_ctx = ralloc_context(NULL);
-
    for (unsigned s = 0; s < layout->num_sets; s++) {
       const unsigned count = layout->set[s].layout->binding_count;
       state.set[s].use_count = rzalloc_array(mem_ctx, uint8_t, count);
@@ -735,6 +913,38 @@ anv_nir_apply_pipeline_layout(const struct anv_physical_device *pdevice,
       if (!function->impl)
          continue;
 
+      /* Before we do the normal lowering, we look for any SSBO operations
+       * that we can lower to the BTI model and lower them up-front.  The BTI
+       * model can perform better than the A64 model for a couple reasons:
+       *
+       *  1. 48-bit address calculations are potentially expensive and using
+       *     the BTI model lets us simply compute 32-bit offsets and the
+       *     hardware adds the 64-bit surface base address.
+       *
+       *  2. The BTI messages, because they use surface states, do bounds
+       *     checking for us.  With the A64 model, we have to do our own
+       *     bounds checking and this means wider pointers and extra
+       *     calculations and branching in the shader.
+       *
+       * The solution to both of these is to convert things to the BTI model
+       * opportunistically.  The reason why we need to do this as a pre-pass
+       * is for two reasons:
+       *
+       *  1. The BTI model requires nir_address_format_32bit_index_offset
+       *     pointers which are not the same type as the pointers needed for
+       *     the A64 model.  Because all our derefs are set up for the A64
+       *     model (in case we have variable pointers), we have to crawl all
+       *     the way back to the vulkan_resource_index intrinsic and build a
+       *     completely fresh index+offset calculation.
+       *
+       *  2. Because the variable-pointers-capable lowering that we do as part
+       *     of apply_pipeline_layout_block is destructive (It really has to
+       *     be to handle variable pointers properly), we've lost the deref
+       *     information by the time we get to the load/store/atomic
+       *     intrinsics in that pass.
+       */
+      lower_direct_buffer_access(function->impl, &state);
+
       nir_builder_init(&state.builder, function->impl);
       nir_foreach_block(block, function->impl)
          apply_pipeline_layout_block(block, &state);