nvk: Implement push constants

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

src/nouveau/vulkan/nvk_cmd_buffer.c

@@ -421,3 +421,27 @@ nvk_CmdBindDescriptorSets(VkCommandBuffer commandBuffer,
    }
    assert(next_dyn_offset <= dynamicOffsetCount);
 }
+
+VKAPI_ATTR void VKAPI_CALL
+nvk_CmdPushConstants(VkCommandBuffer commandBuffer,
+                     VkPipelineLayout layout,
+                     VkShaderStageFlags stageFlags,
+                     uint32_t offset,
+                     uint32_t size,
+                     const void *pValues)
+{
+   VK_FROM_HANDLE(nvk_cmd_buffer, cmd, commandBuffer);
+   if (stageFlags & VK_SHADER_STAGE_COMPUTE_BIT) {
+      struct nvk_descriptor_state *desc =
+         nvk_get_descriptors_state(cmd, VK_PIPELINE_BIND_POINT_COMPUTE);
+
+      memcpy(desc->root.push + offset, pValues, size);
+   }
+
+//   if (stageFlags & VK_SHADER_STAGE_ALL_GRAPHICS) {
+//      struct nvk_descriptor_state *desc =
+//         nvk_get_descriptors_state(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS);
+//
+//      memcpy(desc->root.push + offset, pValues, size);
+//   }
+}

src/nouveau/vulkan/nvk_cmd_buffer.h

@@ -33,7 +33,7 @@ struct nvk_root_descriptor_table {
    };
 
    /* Client push constants */
-   uint8_t push[128];
+   uint8_t push[NVK_MAX_PUSH_SIZE];
 
    /* Descriptor set base addresses */
    uint64_t sets[NVK_MAX_SETS];

src/nouveau/vulkan/nvk_nir_lower_descriptors.c

@@ -95,6 +95,32 @@ lower_load_vulkan_descriptor(nir_builder *b, nir_intrinsic_instr *intrin,
    return true;
 }
 
+static bool
+lower_load_push_constant(nir_builder *b, nir_intrinsic_instr *load,
+                         const struct lower_descriptors_ctx *ctx)
+{
+   const uint32_t push_region_offset =
+      offsetof(struct nvk_root_descriptor_table, push);
+   const uint32_t base = nir_intrinsic_base(load);
+
+   b->cursor = nir_before_instr(&load->instr);
+
+   nir_ssa_def *offset = nir_iadd_imm(b, load->src[0].ssa,
+                                         push_region_offset + base);
+
+   nir_ssa_def *val =
+      nir_load_ubo(b, load->dest.ssa.num_components, load->dest.ssa.bit_size,
+                   nir_imm_int(b, 0), offset,
+                   .align_mul = load->dest.ssa.bit_size / 8,
+                   .align_offset = 0,
+                   .range = push_region_offset + base +
+                            nir_intrinsic_range(load));
+
+   nir_ssa_def_rewrite_uses(&load->dest.ssa, val);
+
+   return true;
+}
+
 static void
 get_resource_deref_binding(nir_builder *b, nir_deref_instr *deref,
                            uint32_t *set, uint32_t *binding,
@@ -146,6 +172,9 @@ lower_intrin(nir_builder *b, nir_intrinsic_instr *intrin,
    case nir_intrinsic_load_vulkan_descriptor:
       return lower_load_vulkan_descriptor(b, intrin, ctx);
 
+   case nir_intrinsic_load_push_constant:
+      return lower_load_push_constant(b, intrin, ctx);
+
    case nir_intrinsic_image_deref_load:
    case nir_intrinsic_image_deref_store:
    case nir_intrinsic_image_deref_atomic:

src/nouveau/vulkan/nvk_physical_device.c

@@ -98,6 +98,7 @@ nvk_GetPhysicalDeviceProperties2(VkPhysicalDevice physicalDevice,
          .maxImageArrayLayers = 2048,
          .maxImageDimension1D = pdevice->dev->chipset >= 0x130 ? 0x8000 : 0x4000,
          .maxImageDimension2D = pdevice->dev->chipset >= 0x130 ? 0x8000 : 0x4000,
+         .maxPushConstantsSize = NVK_MAX_PUSH_SIZE,
          .maxSamplerAllocationCount = 4096,
          .maxBoundDescriptorSets = NVK_MAX_SETS,
          .maxPerStageDescriptorSamplers = UINT32_MAX,

src/nouveau/vulkan/nvk_private.h

@@ -15,6 +15,7 @@
 #include <xf86drm.h>
 
 #define NVK_MAX_SETS 8
+#define NVK_MAX_PUSH_SIZE 128
 #define NVK_MAX_DYNAMIC_BUFFERS 64
 #define NVK_MIN_UBO_ALIGNMENT 16
 

src/nouveau/vulkan/nvk_shader.c

@@ -118,6 +118,10 @@ nvk_shader_compile_to_nir(struct nvk_device *device,
    /* Vulkan uses the separate-shader linking model */
    nir->info.separate_shader = true;
 
+   /* Lower push constants before lower_descriptors */
+   NIR_PASS(_, nir, nir_lower_explicit_io, nir_var_mem_push_const,
+            nir_address_format_32bit_offset);
+
    NIR_PASS(_, nir, nvk_nir_lower_descriptors, layout, true);
    NIR_PASS(_, nir, nir_lower_explicit_io, nir_var_mem_ssbo,
             spirv_options.ssbo_addr_format);