Skip to content

Commit

Permalink
Add partial derivative builtins (#2277)
Browse files Browse the repository at this point in the history 
* Add partial derivative builtins

* [dot] emit derivative control

* Fix fmt

---------

Co-authored-by: Teodor Tanasoaia <[email protected]>
  • Loading branch information
@evahop @teoxoy
evahop and teoxoy authored Mar 16, 2023
1 parent 6fcd04e commit 6db8da7
Show file tree
Hide file tree
Showing 21 changed files with 376 additions and 116 deletions.
4 changes: 2 additions & 2 deletions src/back/dot/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -505,9 +505,9 @@ 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)
(format!("d{axis:?}{ctrl:?}").into(), 8)
}
E::Relational { fun, argument } => {
edges.insert("arg", argument);
Expand Down
34 changes: 24 additions & 10 deletions src/back/glsl/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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, ")")?
}
Expand Down
46 changes: 35 additions & 11 deletions src/back/hlsl/writer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
}

Expand Down Expand Up @@ -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;
Expand Down
9 changes: 5 additions & 4 deletions src/back/msl/writer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;
Expand Down
28 changes: 21 additions & 7 deletions src/back/spv/block.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 12 additions & 7 deletions src/back/wgsl/writer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;
Expand Down
30 changes: 16 additions & 14 deletions src/front/glsl/builtins.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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!(),
},
))
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(),
Expand Down
80 changes: 71 additions & 9 deletions src/front/spv/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;
Expand All @@ -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(
Expand Down Expand Up @@ -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),
)
};
}

Expand Down Expand Up @@ -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)?;
Expand Down
4 changes: 2 additions & 2 deletions src/front/wgsl/lower/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Loading

0 comments on commit 6db8da7

Please sign in to comment.