panvk: Add support for push constants

Push constants are stored in a separate UBO.

Signed-off-by: Boris Brezillon <boris.brezillon@collabora.com>
Reviewed-by: Jason Ekstrand <jason.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/15248>

src/panfrost/vulkan/panvk_cmd_buffer.c

@@ -121,7 +121,25 @@ panvk_CmdPushConstants(VkCommandBuffer commandBuffer,
                        uint32_t size,
                        const void *pValues)
 {
-   panvk_stub();
+   VK_FROM_HANDLE(panvk_cmd_buffer, cmdbuf, commandBuffer);
+
+   memcpy(cmdbuf->push_constants + offset, pValues, size);
+
+   if (stageFlags & VK_SHADER_STAGE_ALL_GRAPHICS) {
+      struct panvk_descriptor_state *desc_state =
+         panvk_cmd_get_desc_state(cmdbuf, GRAPHICS);
+
+      desc_state->ubos = 0;
+      desc_state->push_constants = 0;
+   }
+
+   if (stageFlags & VK_SHADER_STAGE_COMPUTE_BIT) {
+      struct panvk_descriptor_state *desc_state =
+         panvk_cmd_get_desc_state(cmdbuf, COMPUTE);
+
+      desc_state->ubos = 0;
+      desc_state->push_constants = 0;
+   }
 }
 
 void

src/panfrost/vulkan/panvk_descriptor_set.c

@@ -307,6 +307,16 @@ panvk_CreatePipelineLayout(VkDevice _device,
       }
    }
 
+   for (unsigned range = 0; range < pCreateInfo->pushConstantRangeCount; range++) {
+      layout->push_constants.size =
+         MAX2(pCreateInfo->pPushConstantRanges[range].offset +
+              pCreateInfo->pPushConstantRanges[range].size,
+              layout->push_constants.size);
+   }
+
+   if (layout->push_constants.size)
+      layout->push_constants.ubo_idx = ubo_idx++;
+
    layout->num_samplers = sampler_idx;
    layout->num_textures = tex_idx;
    layout->num_ubos = ubo_idx;

src/panfrost/vulkan/panvk_private.h

@@ -419,6 +419,11 @@ struct panvk_pipeline_layout {
    unsigned num_dynoffsets;
    uint32_t num_sets;
 
+   struct {
+      uint32_t size;
+      unsigned ubo_idx;
+   } push_constants;
+
    struct {
       struct panvk_descriptor_set_layout *layout;
       unsigned sampler_offset;
@@ -485,6 +490,7 @@ struct panvk_descriptor_state {
    mali_ptr ubos;
    mali_ptr textures;
    mali_ptr samplers;
+   mali_ptr push_constants;
 };
 
 #define INVOCATION_DESC_WORDS 2

src/panfrost/vulkan/panvk_vX_cmd_buffer.c

@@ -346,6 +346,26 @@ panvk_cmd_prepare_sysvals(struct panvk_cmd_buffer *cmdbuf,
    }
 }
 
+static void
+panvk_cmd_prepare_push_constants(struct panvk_cmd_buffer *cmdbuf,
+                                 struct panvk_cmd_bind_point_state *bind_point_state)
+{
+   struct panvk_descriptor_state *desc_state = &bind_point_state->desc_state;
+   const struct panvk_pipeline *pipeline = bind_point_state->pipeline;
+
+   if (!pipeline->layout->push_constants.size || desc_state->push_constants)
+      return;
+
+   struct panfrost_ptr push_constants =
+      pan_pool_alloc_aligned(&cmdbuf->desc_pool.base,
+                             ALIGN_POT(pipeline->layout->push_constants.size, 16),
+                             16);
+
+   memcpy(push_constants.cpu, cmdbuf->push_constants,
+          pipeline->layout->push_constants.size);
+   desc_state->push_constants = push_constants.gpu;
+}
+
 static void
 panvk_cmd_prepare_ubos(struct panvk_cmd_buffer *cmdbuf,
                        struct panvk_cmd_bind_point_state *bind_point_state)
@@ -357,6 +377,7 @@ panvk_cmd_prepare_ubos(struct panvk_cmd_buffer *cmdbuf,
       return;
 
    panvk_cmd_prepare_sysvals(cmdbuf, bind_point_state);
+   panvk_cmd_prepare_push_constants(cmdbuf, bind_point_state);
 
    struct panfrost_ptr ubos =
       pan_pool_alloc_desc_array(&cmdbuf->desc_pool.base,

src/panfrost/vulkan/panvk_vX_cs.c

@@ -379,6 +379,12 @@ panvk_per_arch(emit_ubos)(const struct panvk_pipeline *pipeline,
                                pipeline->sysvals[i].ids.sysval_count * 16,
                                &ubos[pipeline->sysvals[i].ubo_idx]);
    }
+
+   if (pipeline->layout->push_constants.size) {
+      panvk_per_arch(emit_ubo)(state->push_constants,
+                               ALIGN_POT(pipeline->layout->push_constants.size, 16),
+                               &ubos[pipeline->layout->push_constants.ubo_idx]);
+   }
 }
 
 void

src/panfrost/vulkan/panvk_vX_shader.c

@@ -432,6 +432,33 @@ panvk_lower_blend(struct panfrost_device *pdev,
    }
 }
 
+static bool
+panvk_lower_load_push_constant(nir_builder *b, nir_instr *instr, void *data)
+{
+   if (instr->type != nir_instr_type_intrinsic)
+      return false;
+
+   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
+   if (intr->intrinsic != nir_intrinsic_load_push_constant)
+      return false;
+
+   const struct panvk_pipeline_layout *layout = data;
+
+   b->cursor = nir_before_instr(instr);
+   nir_ssa_def *ubo_load =
+      nir_load_ubo(b, nir_dest_num_components(intr->dest),
+                   nir_dest_bit_size(intr->dest),
+                   nir_imm_int(b, layout->push_constants.ubo_idx),
+                   intr->src[0].ssa,
+                   .align_mul = nir_dest_bit_size(intr->dest) / 8,
+                   .align_offset = 0,
+                   .range_base = nir_intrinsic_base(intr),
+                   .range = nir_intrinsic_range(intr));
+   nir_ssa_def_rewrite_uses(&intr->dest.ssa, ubo_load);
+   nir_instr_remove(instr);
+   return true;
+}
+
 struct panvk_shader *
 panvk_per_arch(shader_create)(struct panvk_device *dev,
                               gl_shader_stage stage,
@@ -513,6 +540,14 @@ panvk_per_arch(shader_create)(struct panvk_device *dev,
    NIR_PASS_V(nir, nir_lower_explicit_io,
               nir_var_mem_ubo | nir_var_mem_ssbo,
               nir_address_format_32bit_index_offset);
+   NIR_PASS_V(nir, nir_lower_explicit_io,
+              nir_var_mem_push_const,
+              nir_address_format_32bit_offset);
+   NIR_PASS_V(nir, nir_shader_instructions_pass,
+              panvk_lower_load_push_constant,
+              nir_metadata_block_index |
+              nir_metadata_dominance,
+              (void *)layout);
 
    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);