glsl: Parse and propagate derivative_group to shader_info

NV_compute_shader_derivatives allow selecting between two possible
arrangements (quads and linear) when calculating derivatives and
certain subgroup operations in case of Vulkan.  So parse and propagate
those up to shader_info.h.

v2: Do not fail when ARB_compute_variable_group_size is being used,
    since we are still clarifying what is the right thing to do here.

Reviewed-by: Ian Romanick <[email protected]>

6 files changed, 140 insertions, 4 deletions

diff --git a/src/compiler/glsl/ast.h b/src/compiler/glsl/ast.h
index 94bcf644a26..90a35654aeb 100644
--- a/src/compiler/glsl/ast.h
+++ b/src/compiler/glsl/ast.h
@@ -654,6 +654,11 @@ struct ast_type_qualifier {
          /** \{ */
          unsigned non_coherent:1;
          /** \} */
+
+         /** \name Layout qualifiers for NV_compute_shader_derivatives */
+         /** \{ */
+         unsigned derivative_group:1;
+         /** \} */
       }
       /** \brief Set of flags, accessed by name. */
       q;
@@ -767,6 +772,12 @@ struct ast_type_qualifier {
    GLenum image_format;
 
    /**
+    * Arrangement of invocations used to calculate derivatives in a compute
+    * shader.  From NV_compute_shader_derivatives.
+    */
+   enum gl_derivative_group derivative_group;
+
+   /**
     * Base type of the data read from or written to this image.  Only
     * the following enumerants are allowed: GLSL_TYPE_UINT,
     * GLSL_TYPE_INT, GLSL_TYPE_FLOAT.
diff --git a/src/compiler/glsl/ast_type.cpp b/src/compiler/glsl/ast_type.cpp
index c2b6e6b3c2b..b499ee913e7 100644
--- a/src/compiler/glsl/ast_type.cpp
+++ b/src/compiler/glsl/ast_type.cpp
@@ -444,6 +444,11 @@ ast_type_qualifier::merge_qualifier(YYLTYPE *loc,
    if (q.flags.q.bound_image)
       this->flags.q.bound_image = true;
 
+   if (q.flags.q.derivative_group) {
+      this->flags.q.derivative_group = true;
+      this->derivative_group = q.derivative_group;
+   }
+
    this->flags.i |= q.flags.i;
 
    if (this->flags.q.in &&
@@ -645,6 +650,7 @@ ast_type_qualifier::validate_in_qualifier(YYLTYPE *loc,
    case MESA_SHADER_COMPUTE:
       valid_in_mask.flags.q.local_size = 7;
       valid_in_mask.flags.q.local_size_variable = 1;
+      valid_in_mask.flags.q.derivative_group = 1;
       break;
    default:
       r = false;
@@ -741,6 +747,19 @@ ast_type_qualifier::merge_into_in_qualifier(YYLTYPE *loc,
       r = false;
    }
 
+   if (state->in_qualifier->flags.q.derivative_group) {
+      if (state->cs_derivative_group != DERIVATIVE_GROUP_NONE) {
+         if (state->in_qualifier->derivative_group != DERIVATIVE_GROUP_NONE &&
+             state->cs_derivative_group != state->in_qualifier->derivative_group) {
+            _mesa_glsl_error(loc, state,
+                             "conflicting derivative groups.");
+            r = false;
+         }
+      } else {
+         state->cs_derivative_group = state->in_qualifier->derivative_group;
+      }
+   }
+
    /* We allow the creation of multiple cs_input_layout nodes. Coherence among
     * all existing nodes is checked later, when the AST node is transformed
     * into HIR.
diff --git a/src/compiler/glsl/glsl_parser.yy b/src/compiler/glsl/glsl_parser.yy
index ddb54f4a4d6..8d6c47fb6a3 100644
--- a/src/compiler/glsl/glsl_parser.yy
+++ b/src/compiler/glsl/glsl_parser.yy
@@ -1655,6 +1655,37 @@ layout_qualifier_id:
             $$.flags.q.non_coherent = 1;
       }
 
+      // Layout qualifiers for NV_compute_shader_derivatives.
+      if (!$$.flags.i) {
+         if (match_layout_qualifier($1, "derivative_group_quadsNV", state) == 0) {
+            $$.flags.q.derivative_group = 1;
+            $$.derivative_group = DERIVATIVE_GROUP_QUADS;
+         } else if (match_layout_qualifier($1, "derivative_group_linearNV", state) == 0) {
+            $$.flags.q.derivative_group = 1;
+            $$.derivative_group = DERIVATIVE_GROUP_LINEAR;
+         }
+
+         if ($$.flags.i) {
+            if (!state->has_compute_shader()) {
+               _mesa_glsl_error(& @1, state,
+                                "qualifier `%s' requires "
+                                "a compute shader", $1);
+            }
+
+            if (!state->NV_compute_shader_derivatives_enable) {
+               _mesa_glsl_error(& @1, state,
+                                "qualifier `%s' requires "
+                                "NV_compute_shader_derivatives", $1);
+            }
+
+            if (state->NV_compute_shader_derivatives_warn) {
+               _mesa_glsl_warning(& @1, state,
+                                  "NV_compute_shader_derivatives layout "
+                                  "qualifier `%s' used", $1);
+            }
+         }
+      }
+
       if (!$$.flags.i) {
          _mesa_glsl_error(& @1, state, "unrecognized layout identifier "
                           "`%s'", $1);
diff --git a/src/compiler/glsl/glsl_parser_extras.cpp b/src/compiler/glsl/glsl_parser_extras.cpp
index ec196299b6a..c1c3c9649fe 100644
--- a/src/compiler/glsl/glsl_parser_extras.cpp
+++ b/src/compiler/glsl/glsl_parser_extras.cpp
@@ -1717,7 +1717,8 @@ set_shader_inout_layout(struct gl_shader *shader,
 {
    /* Should have been prevented by the parser. */
    if (shader->Stage != MESA_SHADER_GEOMETRY &&
-       shader->Stage != MESA_SHADER_TESS_EVAL) {
+       shader->Stage != MESA_SHADER_TESS_EVAL &&
+       shader->Stage != MESA_SHADER_COMPUTE) {
       assert(!state->in_qualifier->flags.i);
    }
 
@@ -1725,6 +1726,7 @@ set_shader_inout_layout(struct gl_shader *shader,
       /* Should have been prevented by the parser. */
       assert(!state->cs_input_local_size_specified);
       assert(!state->cs_input_local_size_variable_specified);
+      assert(state->cs_derivative_group == DERIVATIVE_GROUP_NONE);
    }
 
    if (shader->Stage != MESA_SHADER_FRAGMENT) {
@@ -1849,6 +1851,36 @@ set_shader_inout_layout(struct gl_shader *shader,
 
       shader->info.Comp.LocalSizeVariable =
          state->cs_input_local_size_variable_specified;
+
+      shader->info.Comp.DerivativeGroup = state->cs_derivative_group;
+
+      if (state->NV_compute_shader_derivatives_enable) {
+         /* We allow multiple cs_input_layout nodes, but do not store them in
+          * a convenient place, so for now live with an empty location error.
+          */
+         YYLTYPE loc = {0};
+         if (shader->info.Comp.DerivativeGroup == DERIVATIVE_GROUP_QUADS) {
+            if (shader->info.Comp.LocalSize[0] % 2 != 0) {
+               _mesa_glsl_error(&loc, state, "derivative_group_quadsNV must be used with a "
+                                "local group size whose first dimension "
+                                "is a multiple of 2\n");
+            }
+            if (shader->info.Comp.LocalSize[1] % 2 != 0) {
+               _mesa_glsl_error(&loc, state, "derivative_group_quadsNV must be used with a "
+                                "local group size whose second dimension "
+                                "is a multiple of 2\n");
+            }
+         } else if (shader->info.Comp.DerivativeGroup == DERIVATIVE_GROUP_LINEAR) {
+            if ((shader->info.Comp.LocalSize[0] *
+                 shader->info.Comp.LocalSize[1] *
+                 shader->info.Comp.LocalSize[2]) % 4 != 0) {
+               _mesa_glsl_error(&loc, state, "derivative_group_linearNV must be used with a "
+                            "local group size whose total number of invocations "
+                            "is a multiple of 4\n");
+            }
+         }
+      }
+
       break;
 
    case MESA_SHADER_FRAGMENT:
diff --git a/src/compiler/glsl/glsl_parser_extras.h b/src/compiler/glsl/glsl_parser_extras.h
index 295612e9aed..82a2ce3424b 100644
--- a/src/compiler/glsl/glsl_parser_extras.h
+++ b/src/compiler/glsl/glsl_parser_extras.h
@@ -446,6 +446,12 @@ struct _mesa_glsl_parse_state {
    bool cs_input_local_size_variable_specified;
 
    /**
+    * Arrangement of invocations used to calculate derivatives in a compute
+    * shader.  From NV_compute_shader_derivatives.
+    */
+   enum gl_derivative_group cs_derivative_group;
+
+   /**
     * True if a shader declare bindless_sampler/bindless_image, and
     * respectively bound_sampler/bound_image at global scope as specified by
     * ARB_bindless_texture.
diff --git a/src/compiler/glsl/linker.cpp b/src/compiler/glsl/linker.cpp
index c6664e98ce3..5667fb73f1e 100644
--- a/src/compiler/glsl/linker.cpp
+++ b/src/compiler/glsl/linker.cpp
@@ -2170,9 +2170,9 @@ link_gs_inout_layout_qualifiers(struct gl_shader_program *prog,
 
 
 /**
- * Perform cross-validation of compute shader local_size_{x,y,z} layout
- * qualifiers for the attached compute shaders, and propagate them to the
- * linked CS and linked shader program.
+ * Perform cross-validation of compute shader local_size_{x,y,z} layout and
+ * derivative arrangement qualifiers for the attached compute shaders, and
+ * propagate them to the linked CS and linked shader program.
  */
 static void
 link_cs_input_layout_qualifiers(struct gl_shader_program *prog,
@@ -2191,6 +2191,8 @@ link_cs_input_layout_qualifiers(struct gl_shader_program *prog,
 
    gl_prog->info.cs.local_size_variable = false;
 
+   gl_prog->info.cs.derivative_group = DERIVATIVE_GROUP_NONE;
+
    /* From the ARB_compute_shader spec, in the section describing local size
     * declarations:
     *
@@ -2234,6 +2236,17 @@ link_cs_input_layout_qualifiers(struct gl_shader_program *prog,
          }
          gl_prog->info.cs.local_size_variable = true;
       }
+
+      enum gl_derivative_group group = shader->info.Comp.DerivativeGroup;
+      if (group != DERIVATIVE_GROUP_NONE) {
+         if (gl_prog->info.cs.derivative_group != DERIVATIVE_GROUP_NONE &&
+             gl_prog->info.cs.derivative_group != group) {
+            linker_error(prog, "compute shader defined with conflicting "
+                         "derivative groups\n");
+            return;
+         }
+         gl_prog->info.cs.derivative_group = group;
+      }
    }
 
    /* Just do the intrastage -> interstage propagation right now,
@@ -2246,6 +2259,30 @@ link_cs_input_layout_qualifiers(struct gl_shader_program *prog,
                          "local group size\n");
       return;
    }
+
+   if (gl_prog->info.cs.derivative_group == DERIVATIVE_GROUP_QUADS) {
+      if (gl_prog->info.cs.local_size[0] % 2 != 0) {
+         linker_error(prog, "derivative_group_quadsNV must be used with a "
+                      "local group size whose first dimension "
+                      "is a multiple of 2\n");
+         return;
+      }
+      if (gl_prog->info.cs.local_size[1] % 2 != 0) {
+         linker_error(prog, "derivative_group_quadsNV must be used with a local"
+                      "group size whose second dimension "
+                      "is a multiple of 2\n");
+         return;
+      }
+   } else if (gl_prog->info.cs.derivative_group == DERIVATIVE_GROUP_LINEAR) {
+      if ((gl_prog->info.cs.local_size[0] *
+           gl_prog->info.cs.local_size[1] *
+           gl_prog->info.cs.local_size[2]) % 4 != 0) {
+         linker_error(prog, "derivative_group_linearNV must be used with a "
+                      "local group size whose total number of invocations "
+                      "is a multiple of 4\n");
+         return;
+      }
+   }
 }
 
 /**