nvk: Wire up nir_opt_large_constants

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

src/nouveau/vulkan/nvk_cmd_buffer.c

@@ -763,6 +763,7 @@ nvk_cmd_buffer_flush_push_descriptors(struct nvk_cmd_buffer *cmd,
 bool
 nvk_cmd_buffer_get_cbuf_descriptor(struct nvk_cmd_buffer *cmd,
                                    const struct nvk_descriptor_state *desc,
+                                   const struct nvk_shader *shader,
                                    const struct nvk_cbuf *cbuf,
                                    struct nvk_buffer_address *desc_out)
 {
@@ -775,6 +776,13 @@ nvk_cmd_buffer_get_cbuf_descriptor(struct nvk_cmd_buffer *cmd,
       unreachable("The caller should handle root descriptors");
       return false;
 
+   case NVK_CBUF_TYPE_SHADER_DATA:
+      *desc_out = (struct nvk_buffer_address) {
+         .base_addr = nvk_shader_data_address(shader),
+         .size = shader->data_size,
+      };
+      return true;
+
    case NVK_CBUF_TYPE_DESC_SET:
       *desc_out = (struct nvk_buffer_address) {
          .base_addr = desc->root.sets[cbuf->desc_set],

src/nouveau/vulkan/nvk_cmd_buffer.h

@@ -23,6 +23,7 @@ struct nvk_cmd_bo;
 struct nvk_cmd_pool;
 struct nvk_image_view;
 struct nvk_push_descriptor_set;
+struct nvk_shader;
 
 struct nvk_sample_location {
    uint8_t x_u4:4;
@@ -256,6 +257,7 @@ nvk_cmd_buffer_flush_push_descriptors(struct nvk_cmd_buffer *cmd,
 bool
 nvk_cmd_buffer_get_cbuf_descriptor(struct nvk_cmd_buffer *cmd,
                                    const struct nvk_descriptor_state *desc,
+                                   const struct nvk_shader *shader,
                                    const struct nvk_cbuf *cbuf,
                                    struct nvk_buffer_address *desc_out);
 uint64_t

src/nouveau/vulkan/nvk_cmd_dispatch.c

@@ -229,7 +229,7 @@ nvk_flush_compute_state(struct nvk_cmd_buffer *cmd,
          };
       } else {
          ASSERTED bool direct_descriptor =
-            nvk_cmd_buffer_get_cbuf_descriptor(cmd, desc, cbuf, &ba);
+            nvk_cmd_buffer_get_cbuf_descriptor(cmd, desc, shader, cbuf, &ba);
          assert(direct_descriptor);
       }
 

src/nouveau/vulkan/nvk_cmd_draw.c

@@ -1994,7 +1994,7 @@ nvk_flush_descriptors(struct nvk_cmd_buffer *cmd)
          }
 
          struct nvk_buffer_address ba;
-         if (nvk_cmd_buffer_get_cbuf_descriptor(cmd, desc, cbuf, &ba)) {
+         if (nvk_cmd_buffer_get_cbuf_descriptor(cmd, desc, shader, cbuf, &ba)) {
             assert(ba.base_addr % min_cbuf_alignment == 0);
             ba.size = align(ba.size, min_cbuf_alignment);
             ba.size = MIN2(ba.size, NVK_MAX_CBUF_SIZE);

src/nouveau/vulkan/nvk_nir_lower_descriptors.c

@@ -419,6 +419,12 @@ build_cbuf_map(nir_shader *nir, struct lower_descriptors_ctx *ctx)
       .type = NVK_CBUF_TYPE_ROOT_DESC,
    };
 
+   if (nir->constant_data_size > 0) {
+      ctx->cbuf_map->cbufs[mapped_cbuf_count++] = (struct nvk_cbuf) {
+         .type = NVK_CBUF_TYPE_SHADER_DATA,
+      };
+   }
+
    uint8_t max_cbuf_bindings;
    if (nir->info.stage == MESA_SHADER_COMPUTE ||
        nir->info.stage == MESA_SHADER_KERNEL) {
@@ -506,6 +512,35 @@ lower_load_ubo_intrin(nir_builder *b, nir_intrinsic_instr *load, void *_ctx)
    return true;
 }
 
+static bool
+lower_load_constant(nir_builder *b, nir_intrinsic_instr *load,
+                    const struct lower_descriptors_ctx *ctx)
+{
+   assert(load->intrinsic == nir_intrinsic_load_constant);
+
+   const struct nvk_cbuf cbuf_key = {
+      .type = NVK_CBUF_TYPE_SHADER_DATA,
+   };
+   int cbuf_idx = get_mapped_cbuf_idx(&cbuf_key, ctx);
+   assert(cbuf_idx >= 0);
+
+   uint32_t base = nir_intrinsic_base(load);
+   uint32_t range = nir_intrinsic_range(load);
+
+   b->cursor = nir_before_instr(&load->instr);
+
+   nir_def *offset = nir_iadd_imm(b, load->src[0].ssa, base);
+   nir_def *data = nir_load_ubo(b, load->def.num_components, load->def.bit_size,
+                                nir_imm_int(b, cbuf_idx), offset,
+                                .align_mul = nir_intrinsic_align_mul(load),
+                                .align_offset = nir_intrinsic_align_offset(load),
+                                .range_base = base, .range = range);
+
+   nir_def_rewrite_uses(&load->def, data);
+
+   return true;
+}
+
 static nir_def *
 load_descriptor_set_addr(nir_builder *b, uint32_t set,
                          UNUSED const struct lower_descriptors_ctx *ctx)
@@ -835,6 +870,9 @@ try_lower_intrin(nir_builder *b, nir_intrinsic_instr *intrin,
                  const struct lower_descriptors_ctx *ctx)
 {
    switch (intrin->intrinsic) {
+   case nir_intrinsic_load_constant:
+      return lower_load_constant(b, intrin, ctx);
+
    case nir_intrinsic_load_vulkan_descriptor:
       return try_lower_load_vulkan_descriptor(b, intrin, ctx);
 

src/nouveau/vulkan/nvk_shader.c

@@ -359,6 +359,9 @@ nvk_lower_nir(struct nvk_device *dev, nir_shader *nir,
    if (use_nak(pdev, nir->info.stage) &&
        !(pdev->debug_flags & NVK_DEBUG_NO_CBUF)) {
       cbuf_map = cbuf_map_out;
+
+      /* Large constant support assumes cbufs */
+      NIR_PASS(_, nir, nir_opt_large_constants, NULL, 32);
    } else {
       /* Codegen sometimes puts stuff in cbuf 1 and adds 1 to our cbuf indices
        * so we can't really rely on it for lowering to cbufs and instead place
@@ -477,13 +480,36 @@ nvk_compile_nir(struct nvk_device *dev, nir_shader *nir,
                 struct nvk_shader *shader)
 {
    struct nvk_physical_device *pdev = nvk_device_physical(dev);
+   VkResult result;
 
    if (use_nak(pdev, nir->info.stage)) {
-      return nvk_compile_nir_with_nak(pdev, nir, pipeline_flags, rs,
-                                      fs_key, shader);
+      result = nvk_compile_nir_with_nak(pdev, nir, pipeline_flags, rs,
+                                       fs_key, shader);
    } else {
-      return nvk_cg_compile_nir(pdev, nir, fs_key, shader);
+      result = nvk_cg_compile_nir(pdev, nir, fs_key, shader);
    }
+   if (result != VK_SUCCESS)
+      return result;
+
+   if (nir->constant_data_size > 0) {
+      uint32_t data_align = nvk_min_cbuf_alignment(&dev->pdev->info);
+      uint32_t data_size = align(nir->constant_data_size, data_align);
+
+      void *data = malloc(data_size);
+      if (data == NULL)
+         return vk_error(dev, VK_ERROR_OUT_OF_HOST_MEMORY);
+
+      memcpy(data, nir->constant_data, nir->constant_data_size);
+
+      assert(nir->constant_data_size <= data_size);
+      memset(data + nir->constant_data_size, 0,
+             data_size - nir->constant_data_size);
+
+      shader->data_ptr = data;
+      shader->data_size = data_size;
+   }
+
+   return VK_SUCCESS;
 }
 
 VkResult
@@ -519,6 +545,13 @@ nvk_shader_upload(struct nvk_device *dev, struct nvk_shader *shader)
    assert(code_offset % alignment == 0);
    total_size += shader->code_size;
 
+   uint32_t data_offset = 0;
+   if (shader->data_size > 0) {
+      total_size = align(total_size, nvk_min_cbuf_alignment(&dev->pdev->info));
+      data_offset = total_size;
+      total_size += shader->data_size;
+   }
+
    char *data = malloc(total_size);
    if (data == NULL)
       return vk_error(dev, VK_ERROR_OUT_OF_HOST_MEMORY);
@@ -526,6 +559,8 @@ nvk_shader_upload(struct nvk_device *dev, struct nvk_shader *shader)
    assert(hdr_size <= sizeof(shader->info.hdr));
    memcpy(data + hdr_offset, shader->info.hdr, hdr_size);
    memcpy(data + code_offset, shader->code_ptr, shader->code_size);
+   if (shader->data_size > 0)
+      memcpy(data + data_offset, shader->data_ptr, shader->data_size);
 
 #ifndef NDEBUG
    if (debug_get_bool_option("NV50_PROG_DEBUG", false))
@@ -537,6 +572,7 @@ nvk_shader_upload(struct nvk_device *dev, struct nvk_shader *shader)
    if (result == VK_SUCCESS) {
       shader->upload_size = total_size;
       shader->hdr_offset = hdr_offset;
+      shader->data_offset = data_offset;
    }
    free(data);
 
@@ -562,6 +598,8 @@ nvk_shader_finish(struct nvk_device *dev, struct nvk_shader *shader)
       free((void *)shader->code_ptr);
    }
 
+   free((void *)shader->data_ptr);
+
    vk_free(&dev->vk.alloc, shader);
 }
 
@@ -640,6 +678,8 @@ nvk_shader_serialize(struct vk_pipeline_cache_object *object,
    blob_write_bytes(blob, &shader->cbuf_map, sizeof(shader->cbuf_map));
    blob_write_uint32(blob, shader->code_size);
    blob_write_bytes(blob, shader->code_ptr, shader->code_size);
+   blob_write_uint32(blob, shader->data_size);
+   blob_write_bytes(blob, shader->data_ptr, shader->data_size);
 
    return true;
 }
@@ -669,6 +709,14 @@ nvk_shader_deserialize(struct vk_pipeline_cache *cache,
    blob_copy_bytes(blob, code_ptr, shader->code_size);
    shader->code_ptr = code_ptr;
 
+   shader->data_size = blob_read_uint32(blob);
+   void *data_ptr = malloc(shader->data_size);
+   if (!data_ptr)
+      goto fail;
+
+   blob_copy_bytes(blob, data_ptr, shader->data_size);
+   shader->data_ptr = data_ptr;
+
    return &shader->base;
 
 fail:

src/nouveau/vulkan/nvk_shader.h

@@ -37,6 +37,7 @@ nvk_cbuf_binding_for_stage(gl_shader_stage stage)
 enum ENUM_PACKED nvk_cbuf_type {
    NVK_CBUF_TYPE_INVALID = 0,
    NVK_CBUF_TYPE_ROOT_DESC,
+   NVK_CBUF_TYPE_SHADER_DATA,
    NVK_CBUF_TYPE_DESC_SET,
    NVK_CBUF_TYPE_DYNAMIC_UBO,
    NVK_CBUF_TYPE_UBO_DESC,
@@ -64,9 +65,13 @@ struct nvk_shader {
    const void *code_ptr;
    uint32_t code_size;
 
+   const void *data_ptr;
+   uint32_t data_size;
+
    uint32_t upload_size;
    uint64_t upload_addr;
    uint32_t hdr_offset;
+   uint32_t data_offset;
 };
 
 static inline uint64_t
@@ -75,6 +80,12 @@ nvk_shader_address(const struct nvk_shader *shader)
    return shader->upload_addr + shader->hdr_offset;
 }
 
+static inline uint64_t
+nvk_shader_data_address(const struct nvk_shader *shader)
+{
+   return shader->upload_addr + shader->data_offset;
+}
+
 static inline bool
 nvk_shader_is_enabled(const struct nvk_shader *shader)
 {