anv: fix dynamic buffes again

There were 2 issues in the commit being fixed :

  1. loading from the wrong surface state
  2. not being able to have the optimization passes cleanup the
     nir_vector_extract()

We fix the first issue by reusing the nir_vector_extract() pattern in
the broken places.

We fix the second issue by reworking the internal vec4 format we use
for passing around descriptor information. In particular we put the
set in its own component so that it can be easily optimized and the
vector extraction constant folded.

Fixes: e94cb92cb0 ("anv: use internal surface state on Gfx12.5+ to access descriptor buffers")
Reviewed-by: Ivan Briano <ivan.briano@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/40076>

src/intel/vulkan/anv_nir_apply_pipeline_layout.c

@@ -68,6 +68,7 @@ struct apply_pipeline_layout_state {
    uint8_t constants_offset;
 
    nir_def *set_idx_to_bti;
+   nir_def *set_idx_to_offset;
 
    struct {
       bool desc_buffer_used;
@@ -625,6 +626,61 @@ build_desc_address64(nir_builder *b, nir_def *set_idx, unsigned set_idx_imm,
          BRW_SHADER_RELOC_DESCRIPTORS_ADDR_HIGH));
 }
 
+static nir_def *
+build_descriptor_set_bti(nir_builder *b,
+                         uint32_t set,
+                         const struct apply_pipeline_layout_state *state)
+{
+   if (state->pdevice->info.has_lsc) {
+      nir_def *surface_handle =
+         (state->bind_map->layout_type == ANV_PIPELINE_DESCRIPTOR_SET_LAYOUT_TYPE_INDIRECT ||
+          state->bind_map->layout_type == ANV_PIPELINE_DESCRIPTOR_SET_LAYOUT_TYPE_UNKNOWN) ?
+         nir_imm_int(b, 0xdeaddead) :
+         nir_load_reloc_const_intel(
+            b,
+            state->bind_map->layout_type == ANV_PIPELINE_DESCRIPTOR_SET_LAYOUT_TYPE_BUFFER ?
+            BRW_SHADER_RELOC_DESCRIPTORS_BUFFERS_VIEW_HANDLE :
+            BRW_SHADER_RELOC_DESCRIPTORS_VIEW_HANDLE);
+
+      return nir_resource_intel(
+         b,
+         nir_imm_int(b, set),
+         surface_handle,
+         nir_imm_int(b, 0) /* array_index */,
+         nir_imm_int(b, 0) /* bindless_base_offset */,
+         .desc_set = set,
+         .binding = -1,
+         .resource_block_intel = state->set[set].push_block,
+         .resource_access_intel = nir_resource_intel_pushable |
+                                  nir_resource_intel_internal);
+   } else {
+      return nir_resource_intel(
+         b,
+         nir_imm_int(b, set),
+         nir_imm_int(b, state->set[set].desc_offset),
+         nir_imm_int(b, 0) /* array_index */,
+         nir_imm_int(b, 0) /* bindless_base_offset */,
+         .desc_set = set,
+         .binding = -1,
+         .resource_block_intel = state->set[set].desc_offset,
+         .resource_access_intel = nir_resource_intel_pushable);
+   }
+}
+
+static nir_def *
+build_descriptor_set_offset(nir_builder *b,
+                            uint32_t set,
+                            const struct apply_pipeline_layout_state *state)
+{
+   if (state->pdevice->info.has_lsc) {
+      return nir_iand_imm(b,
+                          anv_load_driver_uniform(b,  1, desc_surface_offsets[set]),
+                          ANV_DESCRIPTOR_SET_OFFSET_MASK /* array_index */);
+   } else {
+      return nir_imm_int(b, 0);
+   }
+}
+
 /** Build a 32bit_index_offset address for a descriptor set */
 static nir_def *
 build_desc_address32(nir_builder *b,
@@ -636,7 +692,11 @@ build_desc_address32(nir_builder *b,
                    nir_vector_extract(
                       b, state->set_idx_to_bti,
                       set < MAX_SETS ? nir_imm_int(b, set) : set_idx),
-                   offset);
+                   nir_iadd(b,
+                            nir_vector_extract(
+                               b, state->set_idx_to_offset,
+                               set < MAX_SETS ? nir_imm_int(b, set) : set_idx),
+                            offset));
 }
 
 /** Build a Vulkan resource index
@@ -668,8 +728,6 @@ build_res_index(nir_builder *b,
    const struct anv_descriptor_set_binding_layout *bind_layout =
       &state->set_layouts[set]->binding[binding];
 
-   uint32_t array_size = bind_layout->array_size;
-
    assert(bind_layout->dynamic_offset_index < MAX_DYNAMIC_BUFFERS);
       nir_def *dynamic_offset_index;
       if (bind_layout->dynamic_offset_index >= 0) {
@@ -689,7 +747,6 @@ build_res_index(nir_builder *b,
          dynamic_offset_index = nir_imm_int(b, 0xff); /* No dynamic offset */
       }
 
-   const uint32_t desc_bti = state->set[set].binding[binding].surface_offset;
    /* We don't care about the stride field for inline uniforms (see
     * build_desc_addr_for_res_index), but for anything else we should be
     * aligned to 8 bytes because we store a multiple of 8 in the packed info
@@ -704,19 +761,18 @@ build_res_index(nir_builder *b,
    nir_def *packed =
       nir_ior_imm(b,
                   dynamic_offset_index,
-                  (desc_stride << 24) |
-                  (desc_bti << 16)    |
-                  (set << 8));
+                  (desc_stride << 8));
 
-   return nir_vec4(b, packed,
-                      nir_imm_int(b, bind_layout->descriptor_surface_offset),
-                      nir_imm_int(b, array_size - 1),
-                      array_index);
+   return nir_vec4(b,
+                   nir_imm_int(b, set),
+                   packed,
+                   nir_imm_int(b, bind_layout->descriptor_surface_offset),
+                   array_index);
 }
 
 struct res_index_defs {
    nir_def *bti_idx;
-   nir_def *set_idx;
+   nir_def *set;
    nir_def *dyn_offset_base;
    nir_def *desc_offset_base;
    nir_def *array_index;
@@ -726,19 +782,15 @@ struct res_index_defs {
 static struct res_index_defs
 unpack_res_index(nir_builder *b, nir_def *index)
 {
-   struct res_index_defs defs;
+   nir_def *packed = nir_channel(b, index, 1);
 
-   nir_def *packed = nir_channel(b, index, 0);
-   defs.desc_stride =
-      nir_imul_imm(b, nir_extract_u8(b, packed, nir_imm_int(b, 3)), 8);
-   defs.bti_idx = nir_extract_u8(b, packed, nir_imm_int(b, 2));
-   defs.set_idx = nir_extract_u8(b, packed, nir_imm_int(b, 1));
-   defs.dyn_offset_base = nir_extract_u8(b, packed, nir_imm_int(b, 0));
-
-   defs.desc_offset_base = nir_channel(b, index, 1);
-   defs.array_index = nir_channel(b, index, 3);
-
-   return defs;
+   return (struct res_index_defs) {
+      .set              = nir_channel(b, index, 0),
+      .desc_stride      = nir_imul_imm(b, nir_extract_u8_imm(b, packed, 1), 8),
+      .dyn_offset_base  = nir_extract_u8_imm(b, packed, 0),
+      .desc_offset_base = nir_channel(b, index, 2),
+      .array_index      = nir_channel(b, index, 3),
+   };
 }
 
 /** Whether a surface is accessed through the bindless surface state heap */
@@ -768,10 +820,11 @@ is_binding_bindless(unsigned set, unsigned binding, bool sampler,
 static nir_def *
 build_res_reindex(nir_builder *b, nir_def *orig, nir_def *delta)
 {
-   return nir_vec4(b, nir_channel(b, orig, 0),
-                      nir_channel(b, orig, 1),
-                      nir_channel(b, orig, 2),
-                      nir_iadd(b, nir_channel(b, orig, 3), delta));
+   return nir_vec4(b,
+                   nir_channel(b, orig, 0),
+                   nir_channel(b, orig, 1),
+                   nir_channel(b, orig, 2),
+                   nir_iadd(b, nir_channel(b, orig, 3), delta));
 }
 
 /** Get the address for a descriptor given its resource index
@@ -808,7 +861,7 @@ build_desc_addr_for_res_index(nir_builder *b,
    case nir_address_format_64bit_bounded_global: {
       switch (state->desc_addr_format) {
       case nir_address_format_64bit_global_32bit_offset: {
-         nir_def *base_addr = build_desc_address64(b, res.set_idx, 0, state);
+         nir_def *base_addr = build_desc_address64(b, res.set, UINT32_MAX, state);
          return nir_vec4(b, nir_unpack_64_2x32_split_x(b, base_addr),
                             nir_unpack_64_2x32_split_y(b, base_addr),
                             nir_imm_int(b, UINT32_MAX),
@@ -816,7 +869,7 @@ build_desc_addr_for_res_index(nir_builder *b,
       }
 
       case nir_address_format_32bit_index_offset:
-         return build_desc_address32(b, res.set_idx, UINT32_MAX, desc_offset, state);
+         return build_desc_address32(b, res.set, UINT32_MAX, desc_offset, state);
 
       default:
          UNREACHABLE("Unhandled address format");
@@ -826,7 +879,7 @@ build_desc_addr_for_res_index(nir_builder *b,
    case nir_address_format_32bit_index_offset:
       assert(desc_type == VK_DESCRIPTOR_TYPE_INLINE_UNIFORM_BLOCK);
       assert(state->desc_addr_format == nir_address_format_32bit_index_offset);
-      return build_desc_address32(b, res.set_idx, UINT32_MAX, desc_offset, state);
+      return build_desc_address32(b, res.set, UINT32_MAX, desc_offset, state);
 
    default:
       UNREACHABLE("Unhandled address format");
@@ -1110,8 +1163,8 @@ build_indirect_buffer_addr_for_res_index(nir_builder *b,
       return build_desc_addr_for_res_index(b, desc_type, res_index,
                                            addr_format, state);
    } else if (addr_format == nir_address_format_32bit_index_offset) {
-      return nir_vec2(b, nir_iadd(b, res.bti_idx, res.array_index),
-                         nir_imm_int(b, 0));
+      return build_desc_address32(b, res.bti_idx, UINT32_MAX,
+                                  nir_imm_int(b, 0), state);
    }
 
    nir_def *desc_addr =
@@ -1171,9 +1224,10 @@ build_direct_buffer_addr_for_res_index(nir_builder *b,
    } else if (addr_format == nir_address_format_32bit_index_offset) {
       struct res_index_defs res = unpack_res_index(b, res_index);
 
-      return nir_vec2(b, nir_iadd(b, res.desc_offset_base,
-                                  nir_imul(b, res.array_index, res.desc_stride)),
-                      nir_imm_int(b, 0));
+      return build_desc_address32(b, res.set, UINT32_MAX,
+                                  nir_iadd(b, res.desc_offset_base,
+                                           nir_imul(b, res.array_index, res.desc_stride)),
+                                  state);
    }
 
    nir_def *desc_addr =
@@ -2439,46 +2493,6 @@ build_packed_binding_table(struct apply_pipeline_layout_state *state,
    }
 }
 
-static nir_def *
-build_descriptor_set_bti(nir_builder *b,
-                         uint32_t set,
-                         const struct apply_pipeline_layout_state *state)
-{
-   if (state->pdevice->info.has_lsc) {
-      nir_def *surface_handle =
-         nir_load_reloc_const_intel(
-            b,
-            state->bind_map->layout_type == ANV_PIPELINE_DESCRIPTOR_SET_LAYOUT_TYPE_BUFFER ?
-            BRW_SHADER_RELOC_DESCRIPTORS_BUFFERS_VIEW_HANDLE :
-            BRW_SHADER_RELOC_DESCRIPTORS_VIEW_HANDLE);
-
-      return nir_resource_intel(
-         b,
-         nir_imm_int(b, set),
-         surface_handle,
-         nir_iand_imm(b,
-                      anv_load_driver_uniform(b,  1, desc_surface_offsets[set]),
-                      ANV_DESCRIPTOR_SET_OFFSET_MASK) /* array_index */,
-         nir_imm_int(b, 0) /* bindless_base_offset */,
-         .desc_set = set,
-         .binding = -1,
-         .resource_block_intel = state->set[set].push_block,
-         .resource_access_intel = nir_resource_intel_pushable |
-                                  nir_resource_intel_internal);
-   } else {
-      return nir_resource_intel(
-         b,
-         nir_imm_int(b, set),
-         nir_imm_int(b, state->set[set].desc_offset),
-         nir_imm_int(b, 0) /* array_index */,
-         nir_imm_int(b, 0) /* bindless_base_offset */,
-         .desc_set = set,
-         .binding = -1,
-         .resource_block_intel = state->set[set].desc_offset,
-         .resource_access_intel = nir_resource_intel_pushable);
-   }
-}
-
 static nir_def *
 build_descriptor_sets_bti_vec(nir_builder *b,
                               const struct apply_pipeline_layout_state *state)
@@ -2495,6 +2509,22 @@ build_descriptor_sets_bti_vec(nir_builder *b,
                    build_descriptor_set_bti(b, 7, state));
 }
 
+static nir_def *
+build_descriptor_sets_offset_vec(nir_builder *b,
+                                 const struct apply_pipeline_layout_state *state)
+{
+   STATIC_ASSERT(MAX_SETS == 8);
+   return nir_vec8(b,
+                   build_descriptor_set_offset(b, 0, state),
+                   build_descriptor_set_offset(b, 1, state),
+                   build_descriptor_set_offset(b, 2, state),
+                   build_descriptor_set_offset(b, 3, state),
+                   build_descriptor_set_offset(b, 4, state),
+                   build_descriptor_set_offset(b, 5, state),
+                   build_descriptor_set_offset(b, 6, state),
+                   build_descriptor_set_offset(b, 7, state));
+}
+
 bool
 anv_nir_apply_pipeline_layout(nir_shader *shader,
                               const struct anv_physical_device *pdevice,
@@ -2578,6 +2608,7 @@ anv_nir_apply_pipeline_layout(nir_shader *shader,
    nir_foreach_function_impl(impl, shader) {
       nir_builder _b = nir_builder_at(nir_before_impl(impl)), *b = &_b;
       state.set_idx_to_bti = build_descriptor_sets_bti_vec(b, &state);
+      state.set_idx_to_offset = build_descriptor_sets_offset_vec(b, &state);
       progress |= nir_function_instructions_pass(impl,
                                                  lower_direct_buffer_instr,
                                                  nir_metadata_control_flow,
@@ -2592,6 +2623,7 @@ anv_nir_apply_pipeline_layout(nir_shader *shader,
    nir_foreach_function_impl(impl, shader) {
       nir_builder _b = nir_builder_at(nir_before_impl(impl)), *b = &_b;
       state.set_idx_to_bti = build_descriptor_sets_bti_vec(b, &state);
+      state.set_idx_to_offset = build_descriptor_sets_offset_vec(b, &state);
       progress |= nir_function_instructions_pass(impl,
                                                  apply_pipeline_layout,
                                                  nir_metadata_control_flow,