derivative: refactor derivatives, allows f32, Vec2, Vec3, Vec4 and Vec3A to be derived

disallow f64 as spec expects only 32bit floats

Author: Firestar99

crates/spirv-std/src/arch/derivative.rs

@@ -1,102 +1,152 @@
-use crate::float::Float;
+use crate::sealed::Sealed;
+use glam::{Vec2, Vec3, Vec3A, Vec4};
+
+macro_rules! cap_deriv_control {
+    () => {
+        unsafe {
+            core::arch::asm!("OpCapability DerivativeControl");
+        }
+    };
+}
 
-#[cfg(target_arch = "spirv")]
 macro_rules! deriv_fn {
-    ($p:ident, $inst:ident) => {
+    ($inst:ident, $param:expr) => {
         unsafe {
-            let mut o = Default::default();
-            ::core::arch::asm!(
-                "%input = OpLoad _ {0}",
-                concat!("%result = ", stringify!($inst), " _ %input"),
-                "OpStore {1} %result",
-                in(reg) &$p,
-                in(reg) &mut o,
+            let mut result = Default::default();
+            core::arch::asm!(
+                "%input = OpLoad typeof*{1} {1}",
+                concat!("%result = ", stringify!($inst), " typeof*{1} %input"),
+                "OpStore {0} %result",
+                in(reg) &mut result,
+                in(reg) &$param,
             );
-            o
+            result
         }
     };
 }
 
-/// Returns the partial derivative of `component` with respect to the window's X
-/// coordinate. Returns the same result as either [`ddx_fine`] or
-/// [`ddx_coarse`], selection of which one is dependent on external factors.
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn ddx<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpDPdx)
-}
+/// Types that can be derived by partial derivatives
+pub unsafe trait Derivative: Sealed + Default {
+    /// Result is the partial derivative of `Self` with respect to the window x coordinate. Uses local differencing
+    /// based on the value of `Self`. Same result as either [`ddx_fine`] or [`ddx_coarse`] on `Self`. Selection of which
+    /// one is based on external factors.
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn ddx(self) -> Self {
+        deriv_fn!(OpDPdx, self)
+    }
 
-/// Returns the partial derivative of `component` with respect to the window's X
-/// coordinate. Uses local differencing based on the value of `component` for
-/// the current fragment and its immediate neighbor(s).
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn ddx_fine<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpDPdxFine)
-}
+    /// Result is the partial derivative of `Self` with respect to the window x coordinate. Uses local differencing
+    /// based on the value of `Self` for the current fragment and its immediate neighbor(s).
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn ddx_fine(self) -> Self {
+        cap_deriv_control!();
+        deriv_fn!(OpDPdxFine, self)
+    }
 
-/// Returns the partial derivative of `component` with respect to the window's X
-/// coordinate. Uses local differencing based on the value of `component` for
-/// the current fragment’s neighbors, and possibly, but not necessarily,
-/// includes the value of `component` for the current fragment. That is, over a
-/// given area, the implementation can compute X derivatives in fewer unique
-/// locations than would be allowed by [`ddx_fine`].
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn ddx_coarse<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpDPdxCoarse)
-}
+    /// Result is the partial derivative of `Self` with respect to the window x coordinate. Uses local differencing
+    /// based on the value of `Self` for the current fragment’s neighbors, and possibly, but not necessarily, includes
+    /// the value of `Self` for the current fragment. That is, over a given area, the implementation can compute x
+    /// derivatives in fewer unique locations than would be allowed for [`ddx_fine`].
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn ddx_coarse(self) -> Self {
+        cap_deriv_control!();
+        deriv_fn!(OpDPdxCoarse, self)
+    }
 
-/// Returns the partial derivative of `component` with respect to the window's Y
-/// coordinate. Returns the same result as either [`ddy_fine`] or
-/// [`ddy_coarse`], selection of which one is dependent on external factors.
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn ddy<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpDPdy)
-}
+    /// Result is the partial derivative of `Self` with respect to the window y coordinate. Uses local differencing
+    /// based on the value of `Self`. Same result as either [`ddy_fine`] or [`ddy_coarse`] on `Self`. Selection of which
+    /// one is based on external factors.
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn ddy(self) -> Self {
+        deriv_fn!(OpDPdy, self)
+    }
 
-/// Returns the partial derivative of `component` with respect to the window's Y
-/// coordinate. Uses local differencing based on the value of `component` for
-/// the current fragment and its immediate neighbor(s).
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn ddy_fine<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpDPdyFine)
-}
+    /// Result is the partial derivative of `Self` with respect to the window y coordinate. Uses local differencing
+    /// based on the value of `Self` for the current fragment and its immediate neighbor(s).
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn ddy_fine(self) -> Self {
+        cap_deriv_control!();
+        deriv_fn!(OpDPdyFine, self)
+    }
 
-/// Returns the partial derivative of `component` with respect to the window's Y
-/// coordinate. Uses local differencing based on the value of `component` for
-/// the current fragment’s neighbors, and possibly, but not necessarily,
-/// includes the value of `component` for the current fragment. That is, over a
-/// given area, the implementation can compute Y derivatives in fewer unique
-/// locations than would be allowed by [`ddy_fine`].
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn ddy_coarse<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpDPdyCoarse)
-}
+    /// Result is the partial derivative of `Self` with respect to the window y coordinate. Uses local differencing
+    /// based on the value of `Self` for the current fragment’s neighbors, and possibly, but not necessarily, includes
+    /// the value of `Self` for the current fragment. That is, over a given area, the implementation can compute y
+    /// derivatives in fewer unique locations than would be allowed for [`ddy_fine`].
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn ddy_coarse(self) -> Self {
+        cap_deriv_control!();
+        deriv_fn!(OpDPdyCoarse, self)
+    }
 
-/// Returns the sum of the absolute values of [`ddx`] and [`ddy`] as a single
-/// operation.
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn fwidth<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpFwidth)
-}
+    /// Result is the same as computing the sum of the absolute values of [`ddx`] and [`ddy`] on P.
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn fwidth(self) -> Self {
+        deriv_fn!(OpFwidth, self)
+    }
 
-/// Returns the sum of the absolute values of [`ddx_fine`] and [`ddy_fine`] as a
-/// single operation.
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn fwidth_fine<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpFwidthFine)
-}
+    /// Result is the same as computing the sum of the absolute values of [`ddx_fine`] and [`ddy_fine`] on P.
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn fwidth_fine(self) -> Self {
+        cap_deriv_control!();
+        deriv_fn!(OpFwidthFine, self)
+    }
 
-/// Returns the sum of the absolute values of [`ddx_coarse`] and [`ddy_coarse`]
-/// as a single operation.
-#[crate::macros::vectorized]
-#[crate::macros::gpu_only]
-pub fn fwidth_coarse<F: Float>(component: F) -> F {
-    deriv_fn!(component, OpFwidthCoarse)
+    /// Result is the same as computing the sum of the absolute values of [`ddx_coarse`] and [`ddy_coarse`] on P.
+    ///
+    /// An invocation will not execute a dynamic instance of this instruction (X') until all invocations in its
+    /// derivative group have executed all dynamic instances that are program-ordered before X'.
+    ///
+    /// This instruction is only valid in the Fragment Execution Model.
+    #[crate::macros::gpu_only]
+    fn fwidth_coarse(self) -> Self {
+        cap_deriv_control!();
+        deriv_fn!(OpFwidthCoarse, self)
+    }
 }
+
+unsafe impl Derivative for f32 {}
+unsafe impl Derivative for Vec2 {}
+unsafe impl Derivative for Vec3 {}
+unsafe impl Derivative for Vec4 {}
+unsafe impl Derivative for Vec3A {}

crates/spirv-std/src/sealed.rs

@@ -13,3 +13,18 @@ impl Sealed for i8 {}
 impl Sealed for i16 {}
 impl Sealed for i32 {}
 impl Sealed for i64 {}
+
+impl Sealed for glam::Vec2 {}
+impl Sealed for glam::Vec3 {}
+impl Sealed for glam::Vec4 {}
+impl Sealed for glam::DVec2 {}
+impl Sealed for glam::DVec3 {}
+impl Sealed for glam::DVec4 {}
+impl Sealed for glam::UVec2 {}
+impl Sealed for glam::UVec3 {}
+impl Sealed for glam::UVec4 {}
+impl Sealed for glam::IVec2 {}
+impl Sealed for glam::IVec3 {}
+impl Sealed for glam::IVec4 {}
+
+impl Sealed for glam::Vec3A {}