radv: Handle nir_intrinsic_printf

Makes it possible to use printf statements inside glsl meta shaders.

Reviewed-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/34208>

src/amd/vulkan/meson.build

@@ -75,6 +75,7 @@ libradv_files = files(
   'nir/radv_nir_lower_intrinsics_early.c',
   'nir/radv_nir_lower_io.c',
   'nir/radv_nir_lower_primitive_shading_rate.c',
+  'nir/radv_nir_lower_printf.c',
   'nir/radv_nir_lower_ray_queries.c',
   'nir/radv_nir_lower_view_index.c',
   'nir/radv_nir_lower_viewport_to_zero.c',

src/amd/vulkan/nir/radv_nir.h

@@ -75,6 +75,8 @@ bool radv_nir_lower_draw_id_to_zero(nir_shader *shader);
 
 bool radv_nir_remap_color_attachment(nir_shader *shader, const struct radv_graphics_state_key *gfx_state);
 
+bool radv_nir_lower_printf(nir_shader *shader);
+
 typedef struct radv_nir_opt_tid_function_options {
    bool use_masked_swizzle_amd : 1;
    bool use_dpp16_shift_amd : 1;

src/amd/vulkan/nir/radv_nir_lower_printf.c

@@ -0,0 +1,52 @@
+/*
+ * Copyright © 2025 Valve Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "radv_nir.h"
+#include "radv_printf.h"
+
+#include "util/u_printf.h"
+
+#include "nir.h"
+#include "nir_builder.h"
+
+static bool
+pass(nir_builder *b, nir_intrinsic_instr *instr, void *state)
+{
+   if (instr->intrinsic != nir_intrinsic_printf)
+      return false;
+
+   u_printf_info *info = &b->shader->printf_info[nir_intrinsic_fmt_idx(instr)];
+
+   nir_def **args = malloc(info->num_args * sizeof(nir_def *));
+
+   b->cursor = nir_after_instr(&instr->instr);
+
+   nir_deref_instr *packed_args = nir_src_as_deref(instr->src[0]);
+   for (uint32_t i = 0; i < info->num_args; i++)
+      args[i] = nir_load_deref(b, nir_build_deref_struct(b, packed_args, i));
+
+   radv_build_printf_args(b, NULL, info->strings, info->num_args, args);
+
+   nir_instr_remove(&instr->instr);
+
+   free(args);
+
+   return true;
+}
+
+bool
+radv_nir_lower_printf(nir_shader *shader)
+{
+   bool progress = nir_shader_intrinsics_pass(shader, pass, nir_metadata_none, NULL);
+
+   /* cleanup */
+   if (progress) {
+      nir_split_struct_vars(shader, nir_var_function_temp);
+      nir_lower_vars_to_ssa(shader);
+   }
+
+   return progress;
+}

src/amd/vulkan/radv_printf.c

@@ -96,14 +96,20 @@ radv_printf_data_finish(struct radv_device *device)
    util_dynarray_fini(&device->printf.formats);
 }
 
-void
-radv_build_printf(nir_builder *b, nir_def *cond, const char *format_string, ...)
+static bool
+radv_shader_printf_enabled(nir_shader *shader)
 {
    if (!device_ht)
-      return;
+      return false;
 
-   struct radv_device *device = _mesa_hash_table_search(device_ht, b->shader)->data;
-   if (!device->printf.buffer_addr)
+   struct radv_device *device = _mesa_hash_table_search(device_ht, shader)->data;
+   return !!device->printf.buffer_addr;
+}
+
+void
+radv_build_printf_args(nir_builder *b, nir_def *cond, const char *format_string, uint32_t argc, nir_def **in_args)
+{
+   if (!radv_shader_printf_enabled(b->shader))
       return;
 
    struct radv_printf_format format = {0};
@@ -111,6 +117,7 @@ radv_build_printf(nir_builder *b, nir_def *cond, const char *format_string, ...)
    if (!format.string)
       return;
 
+   struct radv_device *device = _mesa_hash_table_search(device_ht, b->shader)->data;
    uint32_t format_index = util_dynarray_num_elements(&device->printf.formats, struct radv_printf_format);
 
    if (cond)
@@ -121,22 +128,14 @@ radv_build_printf(nir_builder *b, nir_def *cond, const char *format_string, ...)
 
    nir_def *size = nir_imm_int(b, 4);
 
-   va_list arg_list;
-   va_start(arg_list, format_string);
-
-   uint32_t num_args = 0;
-   for (uint32_t i = 0; i < strlen(format_string); i++)
-      if (format_string[i] == '%')
-         num_args++;
-
-   nir_def **args = malloc(num_args * sizeof(nir_def *));
-   nir_def **strides = malloc(num_args * sizeof(nir_def *));
+   nir_def **args = malloc(argc * sizeof(nir_def *));
+   nir_def **strides = malloc(argc * sizeof(nir_def *));
 
    nir_def *ballot = nir_ballot(b, 1, 64, nir_imm_true(b));
    nir_def *active_invocation_count = nir_bit_count(b, ballot);
 
-   for (uint32_t i = 0; i < num_args; i++) {
-      nir_def *arg = va_arg(arg_list, nir_def *);
+   for (uint32_t i = 0; i < argc; i++) {
+      nir_def *arg = in_args[i];
       bool divergent = arg->divergent;
 
       if (arg->bit_size == 1)
@@ -157,8 +156,6 @@ radv_build_printf(nir_builder *b, nir_def *cond, const char *format_string, ...)
       size = nir_iadd(b, size, strides[i]);
    }
 
-   va_end(arg_list);
-
    nir_def *offset;
    nir_def *undef;
 
@@ -187,7 +184,7 @@ radv_build_printf(nir_builder *b, nir_def *cond, const char *format_string, ...)
       nir_store_global(b, addr, 4, nir_ior_imm(b, active_invocation_count, format_index << 16), 1);
       addr = nir_iadd_imm(b, addr, 4);
 
-      for (uint32_t i = 0; i < num_args; i++) {
+      for (uint32_t i = 0; i < argc; i++) {
          nir_def *arg = args[i];
 
          if (arg->divergent) {
@@ -216,6 +213,32 @@ radv_build_printf(nir_builder *b, nir_def *cond, const char *format_string, ...)
    util_dynarray_append(&device->printf.formats, struct radv_printf_format, format);
 }
 
+void
+radv_build_printf(nir_builder *b, nir_def *cond, const char *format_string, ...)
+{
+   if (!radv_shader_printf_enabled(b->shader))
+      return;
+
+   va_list arg_list;
+   va_start(arg_list, format_string);
+
+   uint32_t num_args = 0;
+   for (uint32_t i = 0; i < strlen(format_string); i++)
+      if (format_string[i] == '%')
+         num_args++;
+
+   nir_def **args = malloc(num_args * sizeof(nir_def *));
+
+   for (uint32_t i = 0; i < num_args; i++)
+      args[i] = va_arg(arg_list, nir_def *);
+
+   va_end(arg_list);
+
+   radv_build_printf_args(b, cond, format_string, num_args, args);
+
+   free(args);
+}
+
 void
 radv_dump_printf_data(struct radv_device *device, FILE *out)
 {

src/amd/vulkan/radv_printf.h

@@ -48,6 +48,8 @@ VkResult radv_printf_data_init(struct radv_device *device);
 
 void radv_printf_data_finish(struct radv_device *device);
 
+void radv_build_printf_args(nir_builder *b, nir_def *cond, const char *format, uint32_t argc, nir_def **args);
+
 void radv_build_printf(nir_builder *b, nir_def *cond, const char *format, ...);
 
 void radv_dump_printf_data(struct radv_device *device, FILE *out);

src/amd/vulkan/radv_shader.c

@@ -394,6 +394,7 @@ radv_shader_spirv_to_nir(struct radv_device *device, const struct radv_shader_st
             },
          .emit_debug_break = !!device->trap_handler_shader,
          .debug_info = !!(instance->debug_flags & RADV_DEBUG_NIR_DEBUG_INFO),
+         .printf = !!device->printf.buffer_addr,
       };
       nir = spirv_to_nir(spirv, stage->spirv.size / 4, spec_entries, num_spec_entries, stage->stage, stage->entrypoint,
                          &spirv_options, &pdev->nir_options[stage->stage]);
@@ -405,6 +406,9 @@ radv_shader_spirv_to_nir(struct radv_device *device, const struct radv_shader_st
 
       radv_device_associate_nir(device, nir);
 
+      if (device->printf.buffer_addr)
+         NIR_PASS(_, nir, radv_nir_lower_printf);
+
       const struct nir_lower_sysvals_to_varyings_options sysvals_to_varyings = {
          .point_coord = true,
       };