mesa/st: Add support for SPIR-V shaders

The SPIR-V codepath uses NIR linking, so we have to preprocess after
the linking steps, which makes things slightly different than GLSL.
To make more clear when the preprocess is happening, I've ended up
inlining st_nir_get_mesa_program() into its caller.

The goal was to make both GLSL and SPIR-V to use the same preprocess
function, the exceptions are:

- SPIR-V codepath don't support NIR state slots yet;
- GLSL lowers shared memory early, so we don't do the deref lowering
  for those.

For now I didn't bother to rename other functions and files (now that
many of them apply to both GLSL and SPIR-V), but we should do this in
further patches.

Reviewed-by: Timothy Arceri <[email protected]>

2 files changed, 110 insertions, 50 deletions

diff --git a/src/mesa/state_tracker/st_glsl_to_ir.cpp b/src/mesa/state_tracker/st_glsl_to_ir.cpp
index ba241c08ad9..25e16fa058c 100644
--- a/src/mesa/state_tracker/st_glsl_to_ir.cpp
+++ b/src/mesa/state_tracker/st_glsl_to_ir.cpp
@@ -59,6 +59,12 @@ st_link_shader(struct gl_context *ctx, struct gl_shader_program *prog)
 
    assert(prog->data->LinkStatus);
 
+   /* Skip the GLSL steps when using SPIR-V. */
+   if (prog->data->spirv) {
+      assert(use_nir);
+      return st_link_nir(ctx, prog);
+   }
+
    for (unsigned i = 0; i < MESA_SHADER_STAGES; i++) {
       if (prog->_LinkedShaders[i] == NULL)
          continue;
diff --git a/src/mesa/state_tracker/st_glsl_to_nir.cpp b/src/mesa/state_tracker/st_glsl_to_nir.cpp
index 5922650de2c..fead0ccc625 100644
--- a/src/mesa/state_tracker/st_glsl_to_nir.cpp
+++ b/src/mesa/state_tracker/st_glsl_to_nir.cpp
@@ -33,6 +33,7 @@
 #include "program/ir_to_mesa.h"
 #include "main/mtypes.h"
 #include "main/errors.h"
+#include "main/glspirv.h"
 #include "main/shaderapi.h"
 #include "main/uniforms.h"
 
@@ -45,6 +46,7 @@
 #include "compiler/glsl_types.h"
 #include "compiler/glsl/glsl_to_nir.h"
 #include "compiler/glsl/gl_nir.h"
+#include "compiler/glsl/gl_nir_linker.h"
 #include "compiler/glsl/ir.h"
 #include "compiler/glsl/ir_optimization.h"
 #include "compiler/glsl/string_to_uint_map.h"
@@ -265,6 +267,18 @@ st_nir_opts(nir_shader *nir, bool scalar)
    } while (progress);
 }
 
+static void
+shared_type_info(const struct glsl_type *type, unsigned *size, unsigned *align)
+{
+   assert(glsl_type_is_vector_or_scalar(type));
+
+   uint32_t comp_size = glsl_type_is_boolean(type)
+      ? 4 : glsl_get_bit_size(type) / 8;
+   unsigned length = glsl_get_vector_elements(type);
+   *size = comp_size * length,
+   *align = comp_size * (length == 3 ? 4 : length);
+}
+
 /* First third of converting glsl_to_nir.. this leaves things in a pre-
  * nir_lower_io state, so that shader variants can more easily insert/
  * replace variables, etc.
@@ -333,6 +347,15 @@ st_nir_preprocess(struct st_context *st, struct gl_program *prog,
    NIR_PASS_V(nir, gl_nir_lower_bindless_images);
    st_nir_opts(nir, is_scalar);
 
+   /* TODO: Change GLSL to not lower shared memory. */
+   if (prog->nir->info.stage == MESA_SHADER_COMPUTE &&
+       shader_program->data->spirv) {
+      NIR_PASS_V(prog->nir, nir_lower_vars_to_explicit_types,
+                 nir_var_mem_shared, shared_type_info);
+      NIR_PASS_V(prog->nir, nir_lower_explicit_io,
+                 nir_var_mem_shared, nir_address_format_32bit_offset);
+   }
+
    NIR_PASS_V(nir, gl_nir_lower_buffers, shader_program);
    /* Do a round of constant folding to clean up address calculations */
    NIR_PASS_V(nir, nir_opt_constant_folding);
@@ -410,7 +433,12 @@ st_glsl_to_nir_post_opts(struct st_context *st, struct gl_program *prog,
 
    st_set_prog_affected_state_flags(prog);
 
-   NIR_PASS_V(nir, st_nir_lower_builtin);
+   /* None of the builtins being lowered here can be produced by SPIR-V.  See
+    * _mesa_builtin_uniform_desc.
+    */
+   if (!shader_program->data->spirv)
+      NIR_PASS_V(nir, st_nir_lower_builtin);
+
    NIR_PASS_V(nir, gl_nir_lower_atomics, shader_program, true);
    NIR_PASS_V(nir, nir_opt_intrinsics);
 
@@ -470,54 +498,6 @@ set_st_program(struct gl_program *prog,
 }
 
 static void
-st_nir_get_mesa_program(struct gl_context *ctx,
-                        struct gl_shader_program *shader_program,
-                        struct gl_linked_shader *shader)
-{
-   struct st_context *st = st_context(ctx);
-   struct pipe_screen *pscreen = ctx->st->pipe->screen;
-   struct gl_program *prog;
-
-   validate_ir_tree(shader->ir);
-
-   prog = shader->Program;
-
-   prog->Parameters = _mesa_new_parameter_list();
-
-   _mesa_copy_linked_program_data(shader_program, shader);
-   _mesa_generate_parameters_list_for_uniforms(ctx, shader_program, shader,
-                                               prog->Parameters);
-
-   /* Remove reads from output registers. */
-   if (!pscreen->get_param(pscreen, PIPE_CAP_TGSI_CAN_READ_OUTPUTS))
-      lower_output_reads(shader->Stage, shader->ir);
-
-   if (ctx->_Shader->Flags & GLSL_DUMP) {
-      _mesa_log("\n");
-      _mesa_log("GLSL IR for linked %s program %d:\n",
-             _mesa_shader_stage_to_string(shader->Stage),
-             shader_program->Name);
-      _mesa_print_ir(_mesa_get_log_file(), shader->ir, NULL);
-      _mesa_log("\n\n");
-   }
-
-   prog->ExternalSamplersUsed = gl_external_samplers(prog);
-   _mesa_update_shader_textures_used(shader_program, prog);
-
-   if (!prog->nir) {
-      const nir_shader_compiler_options *options =
-         st->ctx->Const.ShaderCompilerOptions[prog->info.stage].NirOptions;
-      assert(options);
-
-      prog->nir = glsl_to_nir(st->ctx, shader_program,
-                              prog->info.stage, options);
-      st_nir_preprocess(st, prog, shader_program, prog->info.stage);
-   }
-
-   set_st_program(prog, shader_program, prog->nir);
-}
-
-static void
 st_nir_vectorize_io(nir_shader *producer, nir_shader *consumer)
 {
    NIR_PASS_V(producer, nir_lower_io_to_vector, nir_var_shader_out);
@@ -650,7 +630,53 @@ st_link_nir(struct gl_context *ctx,
       is_scalar[i] = screen->get_shader_param(screen, type,
                                               PIPE_SHADER_CAP_SCALAR_ISA);
 
-      st_nir_get_mesa_program(ctx, shader_program, shader);
+      struct gl_program *prog = shader->Program;
+      _mesa_copy_linked_program_data(shader_program, shader);
+
+      if (shader_program->data->spirv) {
+         const nir_shader_compiler_options *options =
+            st->ctx->Const.ShaderCompilerOptions[shader->Stage].NirOptions;
+
+         prog->Parameters = _mesa_new_parameter_list();
+         /* Parameters will be filled during NIR linking. */
+
+         /* TODO: Properly handle or dismiss `if (prog->nir)` case. */
+         prog->nir = _mesa_spirv_to_nir(ctx, shader_program, shader->Stage, options);
+         set_st_program(prog, shader_program, prog->nir);
+      } else {
+         validate_ir_tree(shader->ir);
+
+         prog->Parameters = _mesa_new_parameter_list();
+         _mesa_generate_parameters_list_for_uniforms(ctx, shader_program, shader,
+                                                     prog->Parameters);
+
+         /* Remove reads from output registers. */
+         if (!screen->get_param(screen, PIPE_CAP_TGSI_CAN_READ_OUTPUTS))
+            lower_output_reads(shader->Stage, shader->ir);
+
+         if (ctx->_Shader->Flags & GLSL_DUMP) {
+            _mesa_log("\n");
+            _mesa_log("GLSL IR for linked %s program %d:\n",
+                      _mesa_shader_stage_to_string(shader->Stage),
+                      shader_program->Name);
+            _mesa_print_ir(_mesa_get_log_file(), shader->ir, NULL);
+            _mesa_log("\n\n");
+         }
+
+         prog->ExternalSamplersUsed = gl_external_samplers(prog);
+         _mesa_update_shader_textures_used(shader_program, prog);
+
+         const nir_shader_compiler_options *options =
+            st->ctx->Const.ShaderCompilerOptions[prog->info.stage].NirOptions;
+         assert(options);
+
+         if (!prog->nir) {
+            prog->nir = glsl_to_nir(st->ctx, shader_program, shader->Stage, options);
+            set_st_program(prog, shader_program, prog->nir);
+            st_nir_preprocess(st, prog, shader_program, shader->Stage);
+         }
+      }
+
       last_stage = i;
 
       if (is_scalar[i]) {
@@ -658,6 +684,34 @@ st_link_nir(struct gl_context *ctx,
       }
    }
 
+   /* For SPIR-V, we have to perform the NIR linking before applying
+    * st_nir_preprocess.
+    */
+   if (shader_program->data->spirv) {
+      if (!gl_nir_link_uniform_blocks(ctx, shader_program))
+         return GL_FALSE;
+
+      if (!gl_nir_link_uniforms(ctx, shader_program, /* fill_parameters */ true))
+         return GL_FALSE;
+
+      gl_nir_link_assign_atomic_counter_resources(ctx, shader_program);
+      gl_nir_link_assign_xfb_resources(ctx, shader_program);
+
+      nir_build_program_resource_list(ctx, shader_program);
+
+      for (unsigned i = 0; i < MESA_SHADER_STAGES; i++) {
+         struct gl_linked_shader *shader = shader_program->_LinkedShaders[i];
+         if (shader == NULL)
+            continue;
+
+         struct gl_program *prog = shader->Program;
+         prog->ExternalSamplersUsed = gl_external_samplers(prog);
+         _mesa_update_shader_textures_used(shader_program, prog);
+
+         st_nir_preprocess(st, prog, shader_program, shader->Stage);
+      }
+   }
+
    /* Linking the stages in the opposite order (from fragment to vertex)
     * ensures that inter-shader outputs written to in an earlier stage
     * are eliminated if they are (transitively) not used in a later