Added basic auto-diff capabilities for local load/store and simple arithmetic. Also added type-checking during the semantic stage. (shader-slang#2303)

* Added JVPTranscriber to handle differentiation of load, store, var, param and return instructions, as well as conversion of data and function types

* Changed class names to be more in line with convention. Added correct type checking for __jvp() and verified that simple calls with only loads and stores are processed correctly

* Added logic to differentiate basic arithmetic and literals inside IRConstruct and fixed the way parameters are differentiated

Co-authored-by: Yong He <yonghe@outlook.com>

Author: saipraveenb25

source/slang/slang-ast-expr.h

@@ -372,10 +372,10 @@ class ExtractExistentialValueExpr: public Expr
     /// An expression of the form `__jvp(fn)` to access the 
     /// forward-mode derivative version of the function `fn`
     ///
-class JVPDerivativeOfExpr: public Expr
+class JVPDifferentiateExpr: public Expr
 {
-    SLANG_AST_CLASS(JVPDerivativeOfExpr)
-    Expr* baseFn;
+    SLANG_AST_CLASS(JVPDifferentiateExpr)
+    Expr* baseFunction;
 };
 
     /// A type expression of the form `__TaggedUnion(A, ...)`.

source/slang/slang-check-expr.cpp

@@ -1509,16 +1509,46 @@ namespace Slang
         return expr;
     }
 
-    Expr* SemanticsExprVisitor::visitJVPDerivativeOfExpr(JVPDerivativeOfExpr* expr)
+    Expr* SemanticsExprVisitor::visitJVPDifferentiateExpr(JVPDifferentiateExpr* expr)
     {
         // Check/Resolve inner function declaration.
-        expr->baseFn = CheckTerm(expr->baseFn);
+        expr->baseFunction = CheckTerm(expr->baseFunction);
 
-        if(auto funcType = as<FuncType>(expr->baseFn->type))
+        if(auto primalType = as<FuncType>(expr->baseFunction->type))
         {
             // Resolve JVP type here. 
-            // Temporarily resolving to the same type as the original function. 
-            expr->type = expr->baseFn->type;
+            // Note that this type checking needs to be in sync with
+            // the auto-generation logic in slang-ir-jvp-diff.cpp
+            
+            auto astBuilder = this->getASTBuilder();
+            FuncType* jvpType = astBuilder->create<FuncType>();
+
+            // Only float types can be differentiated for now.
+
+            // The JVP return type is float if primal return type is float
+            // void otherwise.
+            //
+            if (primalType->resultType->equals(astBuilder->getFloatType()))
+                    jvpType->resultType = astBuilder->getFloatType();
+                else
+                    jvpType->resultType = astBuilder->getVoidType();
+            
+            // No support for differentiating function that throw errors, for now.
+            SLANG_ASSERT(primalType->errorType->equals(astBuilder->getBottomType()));
+            jvpType->errorType = primalType->errorType;
+
+            for (UInt i = 0; i < primalType->getParamCount(); i++)
+            {
+                jvpType->paramTypes.add(primalType->getParamType(i));
+            }
+
+            for (UInt i = 0; i < primalType->getParamCount(); i++)
+            {
+                if(primalType->getParamType(i)->equals(astBuilder->getFloatType()))
+                    jvpType->paramTypes.add(astBuilder->getFloatType());
+            }
+
+            expr->type = jvpType;
         }
         else
         {

source/slang/slang-check-impl.h

@@ -1732,7 +1732,7 @@ namespace Slang
         Expr* visitAndTypeExpr(AndTypeExpr* expr);
         Expr* visitModifiedTypeExpr(ModifiedTypeExpr* expr);
 
-        Expr* visitJVPDerivativeOfExpr(JVPDerivativeOfExpr* expr);
+        Expr* visitJVPDifferentiateExpr(JVPDifferentiateExpr* expr);
 
             /// Perform semantic checking on a `modifier` that is being applied to the given `type`
         Val* checkTypeModifier(Modifier* modifier, Type* type);

source/slang/slang-ir-diff-call.cpp

@@ -34,10 +34,10 @@ struct DerivativeCallProcessContext
                     do 
                     {
                         auto nextChild = child->getNextInst();
-                        // Look for IRJVPDerivativeOf
-                        if (auto derivOf = as<IRJVPDerivativeOf>(child))
+                        // Look for IRJVPDifferentiate
+                        if (auto derivOf = as<IRJVPDifferentiate>(child))
                         {
-                            processDerivativeOf(derivOf);
+                            processDifferentiate(derivOf);
                         }
                         child = nextChild;
                     } 
@@ -50,7 +50,7 @@ struct DerivativeCallProcessContext
 
     // Perform forward-mode automatic differentiation on 
     // the intstructions.
-    void processDerivativeOf(IRJVPDerivativeOf* derivOfInst)
+    void processDifferentiate(IRJVPDifferentiate* derivOfInst)
     {
         IRFunc* jvpFunc = nullptr;