Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add partial derivative builtins #2277

Merged
merged 3 commits into from
Mar 16, 2023
Merged

Add partial derivative builtins #2277

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5866824
Add partial derivative builtins
evahop Mar 14, 2023
b0c70d8
[dot] emit derivative control
teoxoy Mar 16, 2023
2ea62f4
Fix fmt
teoxoy Mar 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/back/dot/mod.rs
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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