From 2db8c15c04f2aade49636e42f0adee636afb3b73 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Fri, 5 Aug 2022 13:19:20 -0400
Subject: Added a new differential type system and various improvements (#2343)

* Merge slang-ir-diff-jvp.cpp

* Added support and tests for other float vector types

* Added swizzle test and code to handle it (tests failing currently)

* Fixed one test, the other is still pending

* Fixed instruction cloning logic to avoid modifying original function

* Fixed an issue with custom 'pow_jvp' and added support for vector contructor

* Minor update to comments

* Fixed support for division

* Fixed an issue with uninitialized diagnostic sink

* Moved derivative processing to after mandatory inlining.

Skip instructions that don't have side-effects and aren't used by anything.

* WIP: Handling unconditional control flow and multi-block functions

* Support for unconditional multi-block functions

* Added a dead code elimination step to the derivative pass

* Changed name of 'hasNoSideEffects()'

* Refactored variable names

* Added initial IR defs for new type system

* Added necessary logic for semantic checking

* Overhauled type system to use builtin pair types and conform to the IDifferentiable interface

* Automatically replace IRDifferentiablePairType to a custom IRStructType

* Added generics handling by expanding the conformance context functionality and allowing for type parameters

* Minor fix: early return in processPairTypes()

* Minor fixes to differentiable resolution on generic types

* Added new instructions for differential pairs. Basic tests work now.

Looking into generic types.

* Adjusted most tests to the new type system. OutType and InOutType are still not properly working.

* Updated __jvp to produce both primal and differential output

* Moved autodiff related declarations to diff.meta.slang

* Refactored variable names

* Added initial IR defs for new type system

* Added necessary logic for semantic checking

* Overhauled type system to use builtin pair types and conform to the IDifferentiable interface

* Automatically replace IRDifferentiablePairType to a custom IRStructType

* Added generics handling by expanding the conformance context functionality and allowing for type parameters

* Minor fix: early return in processPairTypes()

* Minor fixes to differentiable resolution on generic types

* Added new instructions for differential pairs. Basic tests work now.

Looking into generic types.

* Adjusted most tests to the new type system. OutType and InOutType are still not properly working.

* Updated __jvp to produce both primal and differential output

* Moved autodiff related declarations to diff.meta.slang

* Removed external changes

* Cleanup the transcription logic: each case returns a pair of insts for the primal and differential computation.
---
 source/slang/diff.meta.slang | 70 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100644 source/slang/diff.meta.slang

(limited to 'source/slang/diff.meta.slang')

diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang
new file mode 100644
index 000000000..e604140ae
--- /dev/null
+++ b/source/slang/diff.meta.slang
@@ -0,0 +1,70 @@
+
+/// Modifer to mark a function for forward-mode differentiation.
+/// i.e. the compiler will automatically generate a new function
+/// that computes the jacobian-vector product of the original.
+syntax __differentiate_jvp : JVPDerivativeModifier;
+
+// Custom JVP Function reference
+__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.
+__magic_type(DifferentiableType)
+interface IDifferentiable
+{
+    associatedtype Differential;
+};
+
+    /// Pair type that serves to wrap the primal and
+    /// differential types of an arbitrary type T.
+__generic<T : IDifferentiable>
+__magic_type(DifferentialPairType)
+__intrinsic_type($(kIROp_DifferentialPairType))
+struct __DifferentialPair
+{
+
+    __intrinsic_op($(kIROp_MakeDifferentialPair))
+    __init(T _primal, T.Differential _differential);
+
+    __intrinsic_op($(kIROp_DifferentialPairGetDifferential))
+    T.Differential d();
+
+    T.Differential getDifferential()
+    {
+        return d();
+    }
+
+    __intrinsic_op($(kIROp_DifferentialPairGetPrimal))
+    T p();
+
+    T getPrimal()
+    {
+        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;
+}
-- 
cgit v1.2.3