[Mesa-dev] [PATCH 2/9] mesa/st: compute+SSBO support for glsl_to_nir

[Rob Clark]

Signed-off-by: Rob Clark <robdclark at gmail.com>
---
 src/mesa/state_tracker/st_glsl_to_nir.cpp  | 14 ++++++++++++++
 src/mesa/state_tracker/st_glsl_to_tgsi.cpp |  3 ++-
 src/mesa/state_tracker/st_program.c        | 13 +++++++++++++
 src/mesa/state_tracker/st_program.h        |  3 +++
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/mesa/state_tracker/st_glsl_to_nir.cpp b/src/mesa/state_tracker/st_glsl_to_nir.cpp
index 1d850ef..fd5eeea 100644
--- a/src/mesa/state_tracker/st_glsl_to_nir.cpp
+++ b/src/mesa/state_tracker/st_glsl_to_nir.cpp
@@ -242,6 +242,7 @@ st_glsl_to_nir(struct st_context *st, struct gl_program *prog,
    NIR_PASS_V(nir, nir_split_var_copies);
    NIR_PASS_V(nir, nir_lower_var_copies);
    NIR_PASS_V(nir, st_nir_lower_builtin);
+   NIR_PASS_V(nir, nir_lower_atomics, shader_program);
 
    /* fragment shaders may need : */
    if (stage == MESA_SHADER_FRAGMENT) {
@@ -336,6 +337,8 @@ st_finalize_nir(struct st_context *st, struct gl_program *prog, nir_shader *nir)
       nir_assign_var_locations(&nir->outputs,
                                &nir->num_outputs,
                                st_glsl_type_size);
+   } else if (nir->stage == MESA_SHADER_COMPUTE) {
+       /* TODO? */
    } else {
       unreachable("invalid shader type for tgsi bypass\n");
    }
@@ -348,11 +351,17 @@ st_finalize_nir(struct st_context *st, struct gl_program *prog, nir_shader *nir)
    case MESA_SHADER_FRAGMENT:
       shader_program = ((struct st_fragment_program *)prog)->shader_program;
       break;
+   case MESA_SHADER_COMPUTE:
+      shader_program = ((struct st_compute_program *)prog)->shader_program;
+      break;
    default:
       assert(!"should not be reached");
       return;
    }
 
+   NIR_PASS_V(nir, nir_lower_atomics_to_ssbo,
+         st->ctx->Const.Program[nir->stage].MaxAtomicBuffers);
+
    st_nir_assign_uniform_locations(prog, shader_program,
                                    &nir->uniforms, &nir->num_uniforms);
 
@@ -432,6 +441,7 @@ st_nir_get_mesa_program(struct gl_context *ctx,
 
    struct st_vertex_program *stvp;
    struct st_fragment_program *stfp;
+   struct st_compute_program *stcp;
 
    switch (shader->Stage) {
    case MESA_SHADER_VERTEX:
@@ -442,6 +452,10 @@ st_nir_get_mesa_program(struct gl_context *ctx,
       stfp = (struct st_fragment_program *)prog;
       stfp->shader_program = shader_program;
       break;
+   case MESA_SHADER_COMPUTE:
+      stcp = (struct st_compute_program *)prog;
+      stcp->shader_program = shader_program;
+      break;
    default:
       assert(!"should not be reached");
       return NULL;
diff --git a/src/mesa/state_tracker/st_glsl_to_tgsi.cpp b/src/mesa/state_tracker/st_glsl_to_tgsi.cpp
index 96c08a6..80bc01d 100644
--- a/src/mesa/state_tracker/st_glsl_to_tgsi.cpp
+++ b/src/mesa/state_tracker/st_glsl_to_tgsi.cpp
@@ -7014,10 +7014,11 @@ st_link_shader(struct gl_context *ctx, struct gl_shader_program *prog)
 
       struct gl_program *linked_prog = NULL;
       if (preferred_ir == PIPE_SHADER_IR_NIR) {
-         /* TODO only for GLSL VS/FS for now: */
+         /* TODO only for GLSL VS/FS/CS for now: */
          switch (shader->Stage) {
          case MESA_SHADER_VERTEX:
          case MESA_SHADER_FRAGMENT:
+         case MESA_SHADER_COMPUTE:
             linked_prog = st_nir_get_mesa_program(ctx, prog, shader);
          default:
             break;
diff --git a/src/mesa/state_tracker/st_program.c b/src/mesa/state_tracker/st_program.c
index 0dc3b1e..09e0d86 100644
--- a/src/mesa/state_tracker/st_program.c
+++ b/src/mesa/state_tracker/st_program.c
@@ -1798,6 +1798,19 @@ st_translate_compute_program(struct st_context *st,
    struct ureg_program *ureg;
    struct pipe_shader_state prog;
 
+   if (stcp->shader_program) {
+      nir_shader *nir = st_glsl_to_nir(st, &stcp->Base, stcp->shader_program,
+                                       MESA_SHADER_COMPUTE);
+
+      /* no compute variants: */
+      st_finalize_nir(st, &stcp->Base, nir);
+
+      stcp->tgsi.ir_type = PIPE_SHADER_IR_NIR;
+      stcp->tgsi.prog = nir;
+
+      return true;
+   }
+
    ureg = ureg_create_with_screen(PIPE_SHADER_COMPUTE, st->pipe->screen);
    if (ureg == NULL)
       return false;
diff --git a/src/mesa/state_tracker/st_program.h b/src/mesa/state_tracker/st_program.h
index 70664d1..92dfe25 100644
--- a/src/mesa/state_tracker/st_program.h
+++ b/src/mesa/state_tracker/st_program.h
@@ -312,6 +312,9 @@ struct st_compute_program
    struct glsl_to_tgsi_visitor* glsl_to_tgsi;
    uint64_t affected_states; /**< ST_NEW_* flags to mark dirty when binding */
 
+   /* used when bypassing glsl_to_tgsi: */
+   struct gl_shader_program *shader_program;
+
    struct st_basic_variant *variants;
 
    /** SHA1 hash of linked tgsi shader program, used for on-disk cache */
-- 
2.9.3