nir: add derivative intrinsics

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Jesse Natalie <jenatali@microsoft.com>
Reviewed-by: Georg Lehmann <dadschoorse@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30565>

src/compiler/nir/nir.h

@@ -4278,6 +4278,16 @@ typedef struct nir_shader_compiler_options {
     */
    bool discard_is_demote;
 
+   /**
+    * Whether the new-style derivative intrinsics are supported. If false,
+    * legacy ALU derivative ops will be emitted. This transitional option will
+    * be removed once all drivers are converted to derivative intrinsics.
+    */
+   bool has_ddx_intrinsics;
+
+   /** Whether derivative intrinsics must be scalarized. */
+   bool scalarize_ddx;
+
    /** Options determining lowering and behavior of inputs and outputs. */
    nir_io_options io_options;
 

src/compiler/nir/nir_builder.h

@@ -1955,6 +1955,68 @@ nir_tex_src_for_ssa(nir_tex_src_type src_type, nir_def *def)
    return src;
 }
 
+#undef nir_ddx
+#undef nir_ddx_fine
+#undef nir_ddx_coarse
+#undef nir_ddy
+#undef nir_ddy_fine
+#undef nir_ddy_coarse
+
+static inline nir_def *
+nir_build_deriv(nir_builder *b, nir_def *x, nir_op alu, nir_intrinsic_op intrin)
+{
+   /* For derivatives in compute shaders, GLSL_NV_compute_shader_derivatives
+    * states:
+    *
+    *    If neither layout qualifier is specified, derivatives in compute
+    *    shaders return zero, which is consistent with the handling of built-in
+    *    texture functions like texture() in GLSL 4.50 compute shaders.
+    *
+    * We handle that here so the rest of the stack doesn't have to worry about
+    * it and for consistency with previous behaviour. In the future, we might
+    * move this to glsl-to-nir.
+    */
+   if (b->shader->info.stage == MESA_SHADER_COMPUTE &&
+       b->shader->info.cs.derivative_group == DERIVATIVE_GROUP_NONE) {
+
+      return nir_imm_zero(b, x->num_components, x->bit_size);
+   }
+
+   /* Otherwise, build the derivative instruction: either intrinsic or ALU. */
+   if (b->shader->options->has_ddx_intrinsics) {
+      if (b->shader->options->scalarize_ddx && x->num_components > 1) {
+         nir_def *res[NIR_MAX_VEC_COMPONENTS] = { NULL };
+
+         for (unsigned i = 0; i < x->num_components; ++i) {
+            res[i] = _nir_build_ddx(b, x->bit_size, nir_channel(b, x, i));
+            nir_instr_as_intrinsic(res[i]->parent_instr)->intrinsic = intrin;
+         }
+
+         return nir_vec(b, res, x->num_components);
+      } else {
+         nir_def *res = _nir_build_ddx(b, x->bit_size, x);
+         nir_instr_as_intrinsic(res->parent_instr)->intrinsic = intrin;
+         return res;
+      }
+   } else {
+      return nir_build_alu1(b, alu, x);
+   }
+}
+
+#define DEF_DERIV(op)                                                        \
+   static inline nir_def *                                                   \
+      nir_##op(nir_builder *build, nir_def *src0)                            \
+   {                                                                         \
+      return nir_build_deriv(build, src0, nir_op_f##op, nir_intrinsic_##op); \
+   }
+
+DEF_DERIV(ddx)
+DEF_DERIV(ddx_fine)
+DEF_DERIV(ddx_coarse)
+DEF_DERIV(ddy)
+DEF_DERIV(ddy_fine)
+DEF_DERIV(ddy_coarse)
+
 /*
  * Find a texture source, remove it, and return its nir_def. If the texture
  * source does not exist, return NULL. This is useful for texture lowering pass

src/compiler/nir/nir_intrinsics.py

@@ -414,6 +414,13 @@ intrinsic("sparse_residency_code_and", dest_comp=1, src_comp=[1, 1], bit_sizes=[
 intrinsic("is_sparse_resident_zink", dest_comp=1, src_comp=[0], bit_sizes=[1],
           flags=[CAN_ELIMINATE, CAN_REORDER])
 
+# The following intrinsics calculate screen-space partial derivatives. These are
+# not CAN_REORDER as they cannot be moved across discards.
+for suffix in ["", "_fine", "_coarse"]:
+    for axis in ["x", "y"]:
+        intrinsic(f"dd{axis}{suffix}", dest_comp=0, src_comp=[0],
+                  bit_sizes=[16, 32], flags=[CAN_ELIMINATE])
+
 # a barrier is an intrinsic with no inputs/outputs but which can't be moved
 # around/optimized in general
 def barrier(name):