Add partial derivative builtins

src/back/dot/mod.rs

@@ -505,7 +505,11 @@ fn write_function_expressions(
  edges.insert("reject", reject);
  ("Select".into(), 3)
  }
- E::Derivative { axis, expr } => {
+ E::Derivative {
+ axis,
+ ctrl: _,
+ expr,
+ } => {
  edges.insert("", expr);
  (format!("d{axis:?}").into(), 8)
  }

src/back/glsl/mod.rs

@@ -185,6 +185,10 @@ impl Version {
  fn supports_integer_functions(&self) -> bool {
  *self >= Version::Desktop(400) || *self >= Version::new_gles(310)
  }
+
+ fn supports_derivative_control(&self) -> bool {
+ *self >= Version::Desktop(450)
+ }
 }
 
 impl PartialOrd for Version {
@@ -2812,18 +2816,28 @@ impl<'a, W: Write> Writer<'a, W> {
  write!(self.out, ")")?
  }
  // `Derivative` is a function call to a glsl provided function
- Expression::Derivative { axis, expr } => {
- use crate::DerivativeAxis as Da;
-
- write!(
- self.out,
- "{}(",
+ Expression::Derivative { axis, ctrl, expr } => {
+ use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl};
+ let fun_name = if self.options.version.supports_derivative_control() {
+ match (axis, ctrl) {
+ (Axis::X, Ctrl::Coarse) => "dFdxCoarse",
+ (Axis::X, Ctrl::Fine) => "dFdxFine",
+ (Axis::X, Ctrl::None) => "dFdx",
+ (Axis::Y, Ctrl::Coarse) => "dFdyCoarse",
+ (Axis::Y, Ctrl::Fine) => "dFdyFine",
+ (Axis::Y, Ctrl::None) => "dFdy",
+ (Axis::Width, Ctrl::Coarse) => "fwidthCoarse",
+ (Axis::Width, Ctrl::Fine) => "fwidthFine",
+ (Axis::Width, Ctrl::None) => "fwidth",
+ }
+ } else {
  match axis {
- Da::X => "dFdx",
- Da::Y => "dFdy",
- Da::Width => "fwidth",
+ Axis::X => "dFdx",
+ Axis::Y => "dFdy",
+ Axis::Width => "fwidth",
  }
- )?;
+ };
+ write!(self.out, "{fun_name}(")?;
  self.write_expr(expr, ctx)?;
  write!(self.out, ")")?
  }

src/back/hlsl/writer.rs

@@ -133,6 +133,13 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
  _ => {}
  }
  }
+
+ if let Expression::Derivative { axis, ctrl, expr } = *expr {
+ use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl};
+ if axis == Axis::Width && (ctrl == Ctrl::Coarse || ctrl == Ctrl::Fine) {
+ self.need_bake_expressions.insert(expr);
+ }
+ }
  }
  }
 
@@ -2801,17 +2808,34 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
  let var_name = &self.names[&NameKey::GlobalVariable(var_handle)];
  write!(self.out, "({var_name}) - {offset}) / {stride})")?
  }
- Expression::Derivative { axis, expr } => {
- use crate::DerivativeAxis as Da;
-
- let fun_str = match axis {
- Da::X => "ddx",
- Da::Y => "ddy",
- Da::Width => "fwidth",
- };
- write!(self.out, "{fun_str}(")?;
- self.write_expr(module, expr, func_ctx)?;
- write!(self.out, ")")?
+ Expression::Derivative { axis, ctrl, expr } => {
+ use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl};
+ if axis == Axis::Width && (ctrl == Ctrl::Coarse || ctrl == Ctrl::Fine) {
+ let tail = match ctrl {
+ Ctrl::Coarse => "coarse",
+ Ctrl::Fine => "fine",
+ Ctrl::None => unreachable!(),
+ };
+ write!(self.out, "abs(ddx_{tail}(")?;
+ self.write_expr(module, expr, func_ctx)?;
+ write!(self.out, ")) + abs(ddy_{tail}(")?;
+ self.write_expr(module, expr, func_ctx)?;
+ write!(self.out, "))")?
+ } else {
+ let fun_str = match (axis, ctrl) {
+ (Axis::X, Ctrl::Coarse) => "ddx_coarse",
+ (Axis::X, Ctrl::Fine) => "ddx_fine",
+ (Axis::X, Ctrl::None) => "ddx",
+ (Axis::Y, Ctrl::Coarse) => "ddy_coarse",
+ (Axis::Y, Ctrl::Fine) => "ddy_fine",
+ (Axis::Y, Ctrl::None) => "ddy",
+ (Axis::Width, Ctrl::Coarse | Ctrl::Fine) => unreachable!(),
+ (Axis::Width, Ctrl::None) => "fwidth",
+ };
+ write!(self.out, "{fun_str}(")?;
+ self.write_expr(module, expr, func_ctx)?;
+ write!(self.out, ")")?
+ }
  }
  Expression::Relational { fun, argument } => {
  use crate::RelationalFunction as Rf;

src/back/msl/writer.rs

@@ -1579,11 +1579,12 @@ impl<W: Write> Writer<W> {
  }
  _ => return Err(Error::Validation),
  },
- crate::Expression::Derivative { axis, expr } => {
+ crate::Expression::Derivative { axis, expr, .. } => {
+ use crate::DerivativeAxis as Axis;
  let op = match axis {
- crate::DerivativeAxis::X => "dfdx",
- crate::DerivativeAxis::Y => "dfdy",
- crate::DerivativeAxis::Width => "fwidth",
+ Axis::X => "dfdx",
+ Axis::Y => "dfdy",
+ Axis::Width => "fwidth",
  };
  write!(self.out, "{NAMESPACE}::{op}")?;
  self.put_call_parameters(iter::once(expr), context)?;

src/back/spv/block.rs

@@ -1334,15 +1334,29 @@ impl<'w> BlockContext<'w> {
  block.body.push(instruction);
  id
  }
- crate::Expression::Derivative { axis, expr } => {
- use crate::DerivativeAxis as Da;
-
+ crate::Expression::Derivative { axis, ctrl, expr } => {
+ use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl};
+ match ctrl {
+ Ctrl::Coarse | Ctrl::Fine => {
+ self.writer.require_any(
+ "DerivativeControl",
+ &[spirv::Capability::DerivativeControl],
+ )?;
+ }
+ Ctrl::None => {}
+ }
  let id = self.gen_id();
  let expr_id = self.cached[expr];
- let op = match axis {
- Da::X => spirv::Op::DPdx,
- Da::Y => spirv::Op::DPdy,
- Da::Width => spirv::Op::Fwidth,
+ let op = match (axis, ctrl) {
+ (Axis::X, Ctrl::Coarse) => spirv::Op::DPdxCoarse,
+ (Axis::X, Ctrl::Fine) => spirv::Op::DPdxFine,
+ (Axis::X, Ctrl::None) => spirv::Op::DPdx,
+ (Axis::Y, Ctrl::Coarse) => spirv::Op::DPdyCoarse,
+ (Axis::Y, Ctrl::Fine) => spirv::Op::DPdyFine,
+ (Axis::Y, Ctrl::None) => spirv::Op::DPdy,
+ (Axis::Width, Ctrl::Coarse) => spirv::Op::FwidthCoarse,
+ (Axis::Width, Ctrl::Fine) => spirv::Op::FwidthFine,
+ (Axis::Width, Ctrl::None) => spirv::Op::Fwidth,
  };
  block
  .body

src/back/wgsl/writer.rs

@@ -1590,13 +1590,18 @@ impl<W: Write> Writer<W> {
  self.write_expr(module, condition, func_ctx)?;
  write!(self.out, ")")?
  }
- Expression::Derivative { axis, expr } => {
- use crate::DerivativeAxis as Da;
-
- let op = match axis {
- Da::X => "dpdx",
- Da::Y => "dpdy",
- Da::Width => "fwidth",
+ Expression::Derivative { axis, ctrl, expr } => {
+ use crate::{DerivativeAxis as Axis, DerivativeControl as Ctrl};
+ let op = match (axis, ctrl) {
+ (Axis::X, Ctrl::Coarse) => "dpdxCoarse",
+ (Axis::X, Ctrl::Fine) => "dpdxFine",
+ (Axis::X, Ctrl::None) => "dpdx",
+ (Axis::Y, Ctrl::Coarse) => "dpdyCoarse",
+ (Axis::Y, Ctrl::Fine) => "dpdyFine",
+ (Axis::Y, Ctrl::None) => "dpdy",
+ (Axis::Width, Ctrl::Coarse) => "fwidthCoarse",
+ (Axis::Width, Ctrl::Fine) => "fwidthFine",
+ (Axis::Width, Ctrl::None) => "fwidth",
  };
  write!(self.out, "{op}(")?;
  self.write_expr(module, expr, func_ctx)?;

src/front/glsl/builtins.rs

@@ -7,9 +7,10 @@ use super::{
  Error, ErrorKind, Frontend, Result,
 };
 use crate::{
- BinaryOperator, Block, Constant, DerivativeAxis, Expression, Handle, ImageClass,
- ImageDimension as Dim, ImageQuery, MathFunction, Module, RelationalFunction, SampleLevel,
- ScalarKind as Sk, Span, Type, TypeInner, UnaryOperator, VectorSize,
+ BinaryOperator, Block, Constant, DerivativeAxis as Axis, DerivativeControl as Ctrl, Expression,
+ Handle, ImageClass, ImageDimension as Dim, ImageQuery, MathFunction, Module,
+ RelationalFunction, SampleLevel, ScalarKind as Sk, Span, Type, TypeInner, UnaryOperator,
+ VectorSize,
 };
 
 impl crate::ScalarKind {
@@ -571,15 +572,15 @@ fn inject_standard_builtins(
  "degrees" => MacroCall::MathFunction(MathFunction::Degrees),
  "floatBitsToInt" => MacroCall::BitCast(Sk::Sint),
  "floatBitsToUint" => MacroCall::BitCast(Sk::Uint),
- "dFdx" | "dFdxFine" | "dFdxCoarse" => {
-  MacroCall::Derivate(DerivativeAxis::X)
- }
- "dFdy" | "dFdyFine" | "dFdyCoarse" => {
-  MacroCall::Derivate(DerivativeAxis::Y)
- }
- "fwidth" | "fwidthFine" | "fwidthCoarse" => {
-  MacroCall::Derivate(DerivativeAxis::Width)
- }
+ "dFdxCoarse" => MacroCall::Derivate(Axis::X, Ctrl::Coarse),
+ "dFdyCoarse" => MacroCall::Derivate(Axis::Y, Ctrl::Coarse),
+ "fwidthCoarse" => MacroCall::Derivate(Axis::Width, Ctrl::Coarse),
+ "dFdxFine" => MacroCall::Derivate(Axis::X, Ctrl::Fine),
+ "dFdyFine" => MacroCall::Derivate(Axis::Y, Ctrl::Fine),
+ "fwidthFine" => MacroCall::Derivate(Axis::Width, Ctrl::Fine),
+ "dFdx" => MacroCall::Derivate(Axis::X, Ctrl::None),
+ "dFdy" => MacroCall::Derivate(Axis::Y, Ctrl::None),
+ "fwidth" => MacroCall::Derivate(Axis::Width, Ctrl::None),
  _ => unreachable!(),
  },
  ))
@@ -1661,7 +1662,7 @@ pub enum MacroCall {
  MixBoolean,
  Clamp(Option<VectorSize>),
  BitCast(Sk),
- Derivate(DerivativeAxis),
+ Derivate(Axis, Ctrl),
  Barrier,
  /// SmoothStep needs a separate variant because it might need it's inputs
  /// to be splatted depending on the overload
@@ -2143,9 +2144,10 @@ impl MacroCall {
  Span::default(),
  body,
  ),
- MacroCall::Derivate(axis) => ctx.add_expression(
+ MacroCall::Derivate(axis, ctrl) => ctx.add_expression(
  Expression::Derivative {
  axis,
+ ctrl,
  expr: args[0],
  },
  Span::default(),

src/front/spv/mod.rs

@@ -1129,7 +1129,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
  block: &mut crate::Block,
  block_id: spirv::Word,
  body_idx: usize,
- axis: crate::DerivativeAxis,
+ (axis, ctrl): (crate::DerivativeAxis, crate::DerivativeControl),
  ) -> Result<(), Error> {
  let start = self.data_offset;
  let result_type_id = self.next()?;
@@ -1141,6 +1141,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
 
  let expr = crate::Expression::Derivative {
  axis,
+ ctrl,
  expr: arg_handle,
  };
  self.lookup_expression.insert(
@@ -1294,8 +1295,15 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
  ($op:expr, UNARY) => {
  self.parse_expr_unary_op(ctx, &mut emitter, &mut block, block_id, body_idx, $op)
  };
- ($axis:expr, DERIVATIVE) => {
- self.parse_expr_derivative(ctx, &mut emitter, &mut block, block_id, body_idx, $axis)
+ ($axis:expr, $ctrl:expr, DERIVATIVE) => {
+ self.parse_expr_derivative(
+ ctx,
+ &mut emitter,
+ &mut block,
+ block_id,
+ body_idx,
+ ($axis, $ctrl),
+ )
  };
  }
 
@@ -3346,14 +3354,68 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
  });
  body_idx = loop_body_idx;
  }
- Op::DPdx | Op::DPdxFine | Op::DPdxCoarse => {
- parse_expr_op!(crate::DerivativeAxis::X, DERIVATIVE)?;
+ Op::DPdxCoarse => {
+ parse_expr_op!(
+ crate::DerivativeAxis::X,
+ crate::DerivativeControl::Coarse,
+ DERIVATIVE
+ )?;
+ }
+ Op::DPdyCoarse => {
+ parse_expr_op!(
+ crate::DerivativeAxis::Y,
+ crate::DerivativeControl::Coarse,
+ DERIVATIVE
+ )?;
+ }
+ Op::FwidthCoarse => {
+ parse_expr_op!(
+ crate::DerivativeAxis::Width,
+ crate::DerivativeControl::Coarse,
+ DERIVATIVE
+ )?;
+ }
+ Op::DPdxFine => {
+ parse_expr_op!(
+ crate::DerivativeAxis::X,
+ crate::DerivativeControl::Fine,
+ DERIVATIVE
+ )?;
+ }
+ Op::DPdyFine => {
+ parse_expr_op!(
+ crate::DerivativeAxis::Y,
+ crate::DerivativeControl::Fine,
+ DERIVATIVE
+ )?;
+ }
+ Op::FwidthFine => {
+ parse_expr_op!(
+ crate::DerivativeAxis::Width,
+ crate::DerivativeControl::Fine,
+ DERIVATIVE
+ )?;
+ }
+ Op::DPdx => {
+ parse_expr_op!(
+ crate::DerivativeAxis::X,
+ crate::DerivativeControl::None,
+ DERIVATIVE
+ )?;
  }
- Op::DPdy | Op::DPdyFine | Op::DPdyCoarse => {
- parse_expr_op!(crate::DerivativeAxis::Y, DERIVATIVE)?;
+ Op::DPdy => {
+ parse_expr_op!(
+ crate::DerivativeAxis::Y,
+ crate::DerivativeControl::None,
+ DERIVATIVE
+ )?;
  }
- Op::Fwidth | Op::FwidthFine | Op::FwidthCoarse => {
- parse_expr_op!(crate::DerivativeAxis::Width, DERIVATIVE)?;
+ Op::Fwidth => {
+ parse_expr_op!(
+ crate::DerivativeAxis::Width,
+ crate::DerivativeControl::None,
+ DERIVATIVE
+ )?;
  }
  Op::ArrayLength => {
  inst.expect(5)?;

src/front/wgsl/lower/mod.rs

@@ -1572,12 +1572,12 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
  args.finish()?;
 
  crate::Expression::Relational { fun, argument }
- } else if let Some(axis) = conv::map_derivative_axis(function.name) {
+ } else if let Some((axis, ctrl)) = conv::map_derivative(function.name) {
  let mut args = ctx.prepare_args(arguments, 1, span);
  let expr = self.expression(args.next()?, ctx.reborrow())?;
  args.finish()?;
 
- crate::Expression::Derivative { axis, expr }
+ crate::Expression::Derivative { axis, ctrl, expr }
  } else if let Some(fun) = conv::map_standard_fun(function.name) {
  let expected = fun.argument_count() as _;
  let mut args = ctx.prepare_args(arguments, expected, span);