1 files changed, 106 insertions, 27 deletions
diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang
index e604140ae..26fec224c 100644
--- a/source/slang/diff.meta.slang
+++ b/source/slang/diff.meta.slang
@@ -8,18 +8,118 @@ syntax __differentiate_jvp : JVPDerivativeModifier;
 __attributeTarget(FuncDecl)
 attribute_syntax [__custom_jvp(function)]   : CustomJVPAttribute;
 
-//@ public:
-
-    /// Interface to denote types as differentiable.
-    /// Allows for user-specified differential types as
-    /// well as automatic generation, for when the associated type
-    /// hasn't been declared explicitly.
+/// Interface to denote types as differentiable.
+/// Allows for user-specified differential types as
+/// well as automatic generation, for when the associated type
+/// hasn't been declared explicitly.
+/// Note that the requirements must currently be defined in this exact order
+/// since the auto-diff pass relies on the order to grab the struct keys.
+/// 
 __magic_type(DifferentiableType)
 interface IDifferentiable
 {
     associatedtype Differential;
+    
+    static Differential zero();
+
+    static Differential dadd(Differential, Differential);
+
+    static Differential dmul(This, Differential);
 };
 
+// Add extensions for the standard types
+extension float : IDifferentiable
+{
+    typedef float Differential;
+    
+    [__unsafeForceInlineEarly]
+    static Differential zero()
+    {
+        return 0.f;
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dadd(Differential a, Differential b)
+    {
+        return a + b;
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dmul(This a, Differential b)
+    {
+        return a * b;
+    }
+}
+
+extension vector<float, 3> : IDifferentiable
+{
+    typedef vector<float, 3> Differential;
+
+    [__unsafeForceInlineEarly]
+    static Differential zero()
+    {
+        return vector<float, 3>(0.f);
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dadd(Differential a, Differential b)
+    {
+        return a + b;
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dmul(This a, Differential b)
+    {
+        return a * b;
+    }
+}
+
+extension vector<float, 2> : IDifferentiable
+{
+    typedef vector<float, 2> Differential;
+
+    [__unsafeForceInlineEarly]
+    static Differential zero()
+    {
+        return vector<float, 2>(0.f);
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dadd(Differential a, Differential b)
+    {
+        return a + b;
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dmul(This a, Differential b)
+    {
+        return a * b;
+    }
+}
+
+extension vector<float, 4> : IDifferentiable
+{
+    typedef vector<float, 4> Differential;
+
+    [__unsafeForceInlineEarly]
+    static Differential zero()
+    {
+        return vector<float, 4>(0.f);
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dadd(Differential a, Differential b)
+    {
+        return a + b;
+    }
+
+    [__unsafeForceInlineEarly]
+    static Differential dmul(This a, Differential b)
+    {
+        return a * b;
+    }
+}
+
     /// Pair type that serves to wrap the primal and
     /// differential types of an arbitrary type T.
 __generic<T : IDifferentiable>
@@ -47,24 +147,3 @@ struct __DifferentialPair
         return p();
     }
 };
-
-// Add extensions for the standard types
-extension float : IDifferentiable
-{
-    typedef float Differential;
-}
-
-extension vector<float, 3> : IDifferentiable
-{
-    typedef vector<float, 3> Differential;
-}
-
-extension vector<float, 2> : IDifferentiable
-{
-    typedef vector<float, 2> Differential;
-}
-
-extension vector<float, 4> : IDifferentiable
-{
-    typedef vector<float, 4> Differential;
-}