Module: Mesa
Branch: main
Commit: bdec097bb8b392af3e721e4df6c17caad7660a44
URL:    
http://cgit.freedesktop.org/mesa/mesa/commit/?id=bdec097bb8b392af3e721e4df6c17caad7660a44

Author: Faith Ekstrand <faith.ekstr...@collabora.com>
Date:   Fri Dec  8 16:44:07 2023 -0600

nvk/nir: Lower UBO loads to load_ubo when we have a cbuf

This will make it go through the cbuf path in NAK.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/26615>

---

 src/nouveau/vulkan/nvk_nir_lower_descriptors.c | 125 ++++++++++++++++++++-----
 1 file changed, 101 insertions(+), 24 deletions(-)

diff --git a/src/nouveau/vulkan/nvk_nir_lower_descriptors.c 
b/src/nouveau/vulkan/nvk_nir_lower_descriptors.c
index 9ed3d4d9a1d..e08a823e3d2 100644
--- a/src/nouveau/vulkan/nvk_nir_lower_descriptors.c
+++ b/src/nouveau/vulkan/nvk_nir_lower_descriptors.c
@@ -211,12 +211,8 @@ ubo_deref_to_cbuf(nir_deref_instr *deref,
    uint64_t offset = 0;
    uint64_t range = glsl_get_explicit_size(deref->type, false);
    bool offset_valid = true;
-   while (true) {
+   while (deref->deref_type != nir_deref_type_cast) {
       nir_deref_instr *parent = nir_deref_instr_parent(deref);
-      if (parent == NULL) {
-         assert(deref->deref_type == nir_deref_type_cast);
-         break;
-      }
 
       switch (deref->deref_type) {
       case nir_deref_type_var:
@@ -225,9 +221,6 @@ ubo_deref_to_cbuf(nir_deref_instr *deref,
       case nir_deref_type_array:
       case nir_deref_type_array_wildcard: {
          uint32_t stride = nir_deref_instr_array_stride(deref);
-         if (range > stride)
-            offset_valid = false;
-
          if (deref->deref_type == nir_deref_type_array &&
              nir_src_is_const(deref->arr.index)) {
             offset += nir_src_as_uint(deref->arr.index) * stride;
@@ -242,17 +235,14 @@ ubo_deref_to_cbuf(nir_deref_instr *deref,
           * anyway, even with variable pointers.
           */
          offset_valid = false;
+         unreachable("Variable pointers aren't allowed on UBOs");
          break;
 
-      case nir_deref_type_struct:
+      case nir_deref_type_struct: {
          offset += glsl_get_struct_field_offset(parent->type,
                                                 deref->strct.index);
          break;
-
-      case nir_deref_type_cast:
-         /* nir_explicit_io_address_from_deref() can't handle casts */
-         offset_valid = false;
-         break;
+      }
 
       default:
          unreachable("Unknown deref type");
@@ -453,6 +443,67 @@ build_cbuf_map(nir_shader *nir, struct 
lower_descriptors_ctx *ctx)
    ctx->cbufs = NULL;
 }
 
+static int
+get_mapped_cbuf_idx(const struct nvk_cbuf *key,
+                    const struct lower_descriptors_ctx *ctx)
+{
+   if (ctx->cbuf_map == NULL)
+      return -1;
+
+   for (uint32_t c = 0; c < ctx->cbuf_map->cbuf_count; c++) {
+      if (cbufs_equal(&ctx->cbuf_map->cbufs[c], key)) {
+         return c;
+      }
+   }
+
+   return -1;
+}
+
+static bool
+lower_load_ubo_intrin(nir_builder *b, nir_intrinsic_instr *load, void *_ctx)
+{
+   const struct lower_descriptors_ctx *ctx = _ctx;
+
+   if (load->intrinsic != nir_intrinsic_load_deref)
+      return false;
+
+   nir_deref_instr *deref = nir_src_as_deref(load->src[0]);
+   if (!nir_deref_mode_is(deref, nir_var_mem_ubo))
+      return false;
+
+   uint64_t offset, end;
+   UNUSED uint64_t start;
+   UNUSED nir_intrinsic_instr *res_index;
+   struct nvk_cbuf cbuf_key =
+      ubo_deref_to_cbuf(deref, &res_index, &offset, &start, &end, ctx);
+
+   if (cbuf_key.type == NVK_CBUF_TYPE_INVALID)
+      return false;
+
+   if (end > NVK_MAX_CBUF_SIZE)
+      return false;
+
+   int cbuf_idx = get_mapped_cbuf_idx(&cbuf_key, ctx);
+   if (cbuf_idx < 0)
+      return false;
+
+   b->cursor = nir_before_instr(&load->instr);
+
+   nir_deref_path path;
+   nir_deref_path_init(&path, deref, NULL);
+
+   nir_def *addr = nir_imm_ivec2(b, cbuf_idx, offset);
+   nir_address_format addr_format = nir_address_format_32bit_index_offset;
+   for (nir_deref_instr **p = &path.path[1]; *p != NULL; p++)
+      addr = nir_explicit_io_address_from_deref(b, *p, addr, addr_format);
+
+   nir_deref_path_finish(&path);
+
+   nir_lower_explicit_io_instr(b, load, addr, addr_format);
+
+   return true;
+}
+
 static nir_def *
 load_descriptor_set_addr(nir_builder *b, uint32_t set,
                          UNUSED const struct lower_descriptors_ctx *ctx)
@@ -529,12 +580,27 @@ load_descriptor(nir_builder *b, unsigned num_components, 
unsigned bit_size,
       unsigned desc_align_offset = binding_layout->offset + offset_B;
       desc_align_offset %= desc_align_mul;
 
-      nir_def *set_addr = load_descriptor_set_addr(b, set, ctx);
-      nir_def *desc =
-         nir_load_global_constant_offset(b, num_components, bit_size,
-                                         set_addr, desc_ubo_offset,
-                                         .align_mul = desc_align_mul,
-                                         .align_offset = desc_align_offset);
+      const struct nvk_cbuf cbuf_key = {
+         .type = NVK_CBUF_TYPE_DESC_SET,
+         .desc_set = set,
+      };
+      int cbuf_idx = get_mapped_cbuf_idx(&cbuf_key, ctx);
+
+      nir_def *desc;
+      if (cbuf_idx >= 0) {
+         desc = nir_load_ubo(b, num_components, bit_size,
+                             nir_imm_int(b, cbuf_idx),
+                             desc_ubo_offset,
+                             .align_mul = desc_align_mul,
+                             .align_offset = desc_align_offset,
+                             .range = ~0);
+      } else {
+         nir_def *set_addr = load_descriptor_set_addr(b, set, ctx);
+         desc = nir_load_global_constant_offset(b, num_components, bit_size,
+                                                set_addr, desc_ubo_offset,
+                                                .align_mul = desc_align_mul,
+                                                .align_offset = 
desc_align_offset);
+      }
       if (binding_layout->type == VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER ||
           binding_layout->type == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER) {
          /* We know a priori that the the .w compnent (offset) is zero */
@@ -1083,21 +1149,32 @@ nvk_nir_lower_descriptors(nir_shader *nir,
       .ubo_addr_format = nvk_buffer_addr_format(rs->uniform_buffers),
    };
 
-   /* We run in three passes:
+   /* We run in four passes:
     *
     *  1. Find ranges of UBOs that we can promote to bound UBOs.  Nothing is
     *     actually lowered in this pass.  It's just analysis.
     *
-    *  2. Attempt to lower everything with direct descriptors.  This may fail
+    *  2. Try to lower UBO loads to cbufs based on the map we just created.
+    *     We need to do this before the main lowering pass because it relies
+    *     on the original descriptor load intrinsics.
+    *
+    *  3. Attempt to lower everything with direct descriptors.  This may fail
     *     to lower some SSBO intrinsics when variable pointers are used.
     *
-    *  3. Clean up any SSBO intrinsics which are left and lower them to
+    *  4. Clean up any SSBO intrinsics which are left and lower them to
     *     slightly less efficient but variable- pointers-correct versions.
     */
 
+   bool pass_lower_ubo = false;
    if (cbuf_map_out != NULL) {
       ctx.cbuf_map = cbuf_map_out;
       build_cbuf_map(nir, &ctx);
+
+      pass_lower_ubo =
+         nir_shader_intrinsics_pass(nir, lower_load_ubo_intrin,
+                                    nir_metadata_block_index |
+                                    nir_metadata_dominance,
+                                    (void *)&ctx);
    }
 
    bool pass_lower_descriptors =
@@ -1110,5 +1187,5 @@ nvk_nir_lower_descriptors(nir_shader *nir,
                                    nir_metadata_block_index |
                                    nir_metadata_dominance,
                                    (void *)&ctx);
-   return pass_lower_descriptors || pass_lower_ssbo;
+   return pass_lower_ubo || pass_lower_descriptors || pass_lower_ssbo;
 }

Reply via email to