panvk: Move NIR lower logic out of shader_create

Signed-off-by: Mary Guillemard <mary.guillemard@collabora.com>
Reviewed-by: Boris Brezillon <boris.brezillon@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/29161>

src/panfrost/vulkan/bifrost/panvk_vX_nir_lower_descriptors.c

@@ -1029,9 +1029,10 @@ upload_shader_desc_info(struct panvk_device *dev, struct panvk_shader *shader,
 }
 
 bool
-panvk_per_arch(nir_lower_descriptors)(nir_shader *nir, struct panvk_device *dev,
+panvk_per_arch(nir_lower_descriptors)(
-                                      const struct vk_pipeline_layout *layout,
+   nir_shader *nir, struct panvk_device *dev, uint32_t set_layout_count,
-                                      struct panvk_shader *shader)
+   struct vk_descriptor_set_layout *const *set_layouts,
+   struct panvk_shader *shader)
 {
    struct lower_desc_ctx ctx = {
       .ubo_addr_format = nir_address_format_32bit_index_offset,
@@ -1045,9 +1046,8 @@ panvk_per_arch(nir_lower_descriptors)(nir_shader *nir, struct panvk_device *dev,
 
    _mesa_hash_table_set_deleted_key(ctx.ht, DELETED_KEY);
 
-   for (uint32_t i = 0; i < layout->set_count; i++) {
+   for (uint32_t i = 0; i < set_layout_count; i++) {
-      ctx.set_layouts[i] =
+      ctx.set_layouts[i] = to_panvk_descriptor_set_layout(set_layouts[i]);
-         to_panvk_descriptor_set_layout(layout->set_layouts[i]);
    }
 
    bool progress = nir_shader_instructions_pass(nir, collect_instr_desc_access,

src/panfrost/vulkan/panvk_shader.h

@@ -161,7 +161,8 @@ panvk_shader_link_cleanup(struct panvk_pool *desc_pool,
 }
 
 bool panvk_per_arch(nir_lower_descriptors)(
-   nir_shader *nir, struct panvk_device *dev,
+   nir_shader *nir, struct panvk_device *dev, uint32_t set_layout_count,
-   const struct vk_pipeline_layout *layout, struct panvk_shader *shader);
+   struct vk_descriptor_set_layout *const *set_layouts,
+   struct panvk_shader *shader);
 
 #endif

src/panfrost/vulkan/panvk_vX_shader.c

@@ -137,6 +137,105 @@ shared_type_info(const struct glsl_type *type, unsigned *size, unsigned *align)
    *size = comp_size * length, *align = comp_size * (length == 3 ? 4 : length);
 }
 
+static inline nir_address_format
+panvk_buffer_ubo_addr_format(VkPipelineRobustnessBufferBehaviorEXT robustness)
+{
+   switch (robustness) {
+   case VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_DISABLED_EXT:
+   case VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_ROBUST_BUFFER_ACCESS_EXT:
+   case VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_ROBUST_BUFFER_ACCESS_2_EXT:
+      return nir_address_format_32bit_index_offset;
+   default:
+      unreachable("Invalid robust buffer access behavior");
+   }
+}
+
+static inline nir_address_format
+panvk_buffer_ssbo_addr_format(VkPipelineRobustnessBufferBehaviorEXT robustness)
+{
+   switch (robustness) {
+   case VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_DISABLED_EXT:
+      return nir_address_format_64bit_global_32bit_offset;
+   case VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_ROBUST_BUFFER_ACCESS_EXT:
+   case VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_ROBUST_BUFFER_ACCESS_2_EXT:
+      return nir_address_format_64bit_bounded_global;
+   default:
+      unreachable("Invalid robust buffer access behavior");
+   }
+}
+
+static void
+panvk_lower_nir(struct panvk_device *dev, nir_shader *nir,
+                uint32_t set_layout_count,
+                struct vk_descriptor_set_layout *const *set_layouts,
+                const struct vk_pipeline_robustness_state *rs,
+                const struct panfrost_compile_inputs *compile_input,
+                struct panvk_shader *shader)
+{
+   struct panvk_instance *instance =
+      to_panvk_instance(dev->vk.physical->instance);
+   gl_shader_stage stage = nir->info.stage;
+
+   NIR_PASS_V(nir, panvk_per_arch(nir_lower_descriptors), dev, set_layout_count,
+              set_layouts, shader);
+
+   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_ubo,
+              panvk_buffer_ubo_addr_format(rs->uniform_buffers));
+   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_ssbo,
+              panvk_buffer_ssbo_addr_format(rs->storage_buffers));
+   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_push_const,
+              nir_address_format_32bit_offset);
+   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_global,
+              nir_address_format_64bit_global);
+
+   if (gl_shader_stage_uses_workgroup(stage)) {
+      if (!nir->info.shared_memory_explicit_layout) {
+         NIR_PASS_V(nir, nir_lower_vars_to_explicit_types, nir_var_mem_shared,
+                    shared_type_info);
+      }
+
+      NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_shared,
+                 nir_address_format_32bit_offset);
+   }
+
+   if (stage == MESA_SHADER_VERTEX) {
+      /* We need the driver_location to match the vertex attribute location,
+       * so we can use the attribute layout described by
+       * vk_vertex_input_state where there are holes in the attribute locations.
+       */
+      nir_foreach_shader_in_variable(var, nir) {
+         assert(var->data.location >= VERT_ATTRIB_GENERIC0 &&
+                var->data.location <= VERT_ATTRIB_GENERIC15);
+         var->data.driver_location = var->data.location - VERT_ATTRIB_GENERIC0;
+      }
+   } else {
+      nir_assign_io_var_locations(nir, nir_var_shader_in, &nir->num_inputs,
+                                  stage);
+   }
+
+   nir_assign_io_var_locations(nir, nir_var_shader_out, &nir->num_outputs,
+                               stage);
+
+   /* Needed to turn shader_temp into function_temp since the backend only
+    * handles the latter for now.
+    */
+   NIR_PASS_V(nir, nir_lower_global_vars_to_local);
+
+   nir_shader_gather_info(nir, nir_shader_get_entrypoint(nir));
+   if (unlikely(instance->debug_flags & PANVK_DEBUG_NIR)) {
+      fprintf(stderr, "translated nir:\n");
+      nir_print_shader(nir, stderr);
+   }
+
+   pan_shader_preprocess(nir, compile_input->gpu_id);
+
+   if (stage == MESA_SHADER_VERTEX)
+      NIR_PASS_V(nir, pan_lower_image_index, MAX_VS_ATTRIBS);
+
+   NIR_PASS_V(nir, nir_shader_instructions_pass, panvk_lower_sysvals,
+              nir_metadata_block_index | nir_metadata_dominance, NULL);
+}
+
 static VkResult
 panvk_compile_nir(struct panvk_device *dev, nir_shader *nir,
                   VkShaderCreateFlagsEXT shader_flags,
@@ -245,8 +344,6 @@ panvk_per_arch(shader_create)(struct panvk_device *dev,
 {
    struct panvk_physical_device *phys_dev =
       to_panvk_physical_device(dev->vk.physical);
-   struct panvk_instance *instance =
-      to_panvk_instance(dev->vk.physical->instance);
    gl_shader_stage stage = vk_to_mesa_shader_stage(stage_info->stage);
    struct panvk_shader *shader;
 
@@ -276,12 +373,6 @@ panvk_per_arch(shader_create)(struct panvk_device *dev,
    NIR_PASS_V(nir, nir_lower_io_to_temporaries, nir_shader_get_entrypoint(nir),
               true, true);
 
-   struct panfrost_compile_inputs inputs = {
-      .gpu_id = phys_dev->kmod.props.gpu_prod_id,
-      .no_ubo_to_push = true,
-      .no_idvs = true, /* TODO */
-   };
-
    NIR_PASS_V(nir, nir_lower_indirect_derefs,
               nir_var_shader_in | nir_var_shader_out, UINT32_MAX);
 
@@ -330,27 +421,6 @@ panvk_per_arch(shader_create)(struct panvk_device *dev,
    };
    NIR_PASS_V(nir, nir_lower_tex, &lower_tex_options);
 
-   NIR_PASS_V(nir, panvk_per_arch(nir_lower_descriptors), dev, layout, shader);
-
-   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_ubo,
-              nir_address_format_32bit_index_offset);
-   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_ssbo,
-              spirv_options.ssbo_addr_format);
-   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_push_const,
-              nir_address_format_32bit_offset);
-   NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_global,
-              nir_address_format_64bit_global);
-
-   if (gl_shader_stage_uses_workgroup(stage)) {
-      if (!nir->info.shared_memory_explicit_layout) {
-         NIR_PASS_V(nir, nir_lower_vars_to_explicit_types, nir_var_mem_shared,
-                    shared_type_info);
-      }
-
-      NIR_PASS_V(nir, nir_lower_explicit_io, nir_var_mem_shared,
-                 nir_address_format_32bit_offset);
-   }
-
    NIR_PASS_V(nir, nir_lower_system_values);
 
    nir_lower_compute_system_values_options options = {
@@ -365,42 +435,24 @@ panvk_per_arch(shader_create)(struct panvk_device *dev,
    NIR_PASS_V(nir, nir_split_var_copies);
    NIR_PASS_V(nir, nir_lower_var_copies);
 
-   if (stage == MESA_SHADER_VERTEX) {
+   struct vk_pipeline_robustness_state rs = {
-      /* We need the driver_location to match the vertex attribute location,
+      .storage_buffers =
-       * so we can use the attribute layout described by
+         dev->vk.enabled_features.robustBufferAccess
-       * vk_vertex_input_state where there are holes in the attribute locations.
+            ? VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_ROBUST_BUFFER_ACCESS_EXT
-       */
+            : VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_DISABLED_EXT,
-      nir_foreach_shader_in_variable(var, nir) {
+      .uniform_buffers = VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_DISABLED_EXT,
-         assert(var->data.location >= VERT_ATTRIB_GENERIC0 &&
+      .vertex_inputs = VK_PIPELINE_ROBUSTNESS_BUFFER_BEHAVIOR_DISABLED_EXT,
-                var->data.location <= VERT_ATTRIB_GENERIC15);
+      .images = VK_PIPELINE_ROBUSTNESS_IMAGE_BEHAVIOR_DISABLED_EXT,
-         var->data.driver_location = var->data.location - VERT_ATTRIB_GENERIC0;
+   };
-      }
-   } else {
-      nir_assign_io_var_locations(nir, nir_var_shader_in, &nir->num_inputs,
-                                  stage);
-   }
 
-   nir_assign_io_var_locations(nir, nir_var_shader_out, &nir->num_outputs,
+   struct panfrost_compile_inputs inputs = {
-                               stage);
+      .gpu_id = phys_dev->kmod.props.gpu_prod_id,
+      .no_ubo_to_push = true,
+      .no_idvs = true, /* TODO */
+   };
 
-   /* Needed to turn shader_temp into function_temp since the backend only
+   panvk_lower_nir(dev, nir, layout->set_count, layout->set_layouts, &rs,
-    * handles the latter for now.
+                   &inputs, shader);
-    */
-   NIR_PASS_V(nir, nir_lower_global_vars_to_local);
-
-   nir_shader_gather_info(nir, nir_shader_get_entrypoint(nir));
-   if (unlikely(instance->debug_flags & PANVK_DEBUG_NIR)) {
-      fprintf(stderr, "translated nir:\n");
-      nir_print_shader(nir, stderr);
-   }
-
-   pan_shader_preprocess(nir, inputs.gpu_id);
-
-   if (stage == MESA_SHADER_VERTEX)
-      NIR_PASS_V(nir, pan_lower_image_index, MAX_VS_ATTRIBS);
-
-   NIR_PASS_V(nir, nir_shader_instructions_pass, panvk_lower_sysvals,
-              nir_metadata_block_index | nir_metadata_dominance, NULL);
 
    result = panvk_compile_nir(dev, nir, 0, &inputs, shader);