Added support for differentiating calls to basic functions, as well as arithmetic on the float3 type (#2313)

* Added support for differentiating calls to basic functions, as well as arithmetic on the float3 type

* Added test expected result

4 files changed, 132 insertions, 15 deletions

diff --git a/source/compiler-core/slang-diagnostic-sink.cpp b/source/compiler-core/slang-diagnostic-sink.cpp
index 0110b16d7..34a3c4968 100644
--- a/source/compiler-core/slang-diagnostic-sink.cpp
+++ b/source/compiler-core/slang-diagnostic-sink.cpp
@@ -642,13 +642,13 @@ void DiagnosticSink::diagnoseRaw(
     // Did the client supply a callback for us to use?
     if(writer)
     {
-        // If so, pass the error string along to them
+        // If so, pass the error string along to them.
         writer->write(message.begin(), message.getLength());
     }
     else
     {
         // If the user doesn't have a callback, then just
-        // collect our diagnostic messages into a buffer
+        // collect our diagnostic messages into a buffer.
         outputBuffer.append(message);
     }
 
diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
index df58b11ed..a3fec4802 100644
--- a/source/slang/slang-check-expr.cpp
+++ b/source/slang/slang-check-expr.cpp
@@ -1509,6 +1509,29 @@ namespace Slang
         return expr;
     }
 
+    Type* primalToJVPParamType(ASTBuilder* builder, Type* primalType)
+    {
+        // Only float and float3 types can be differentiated for now.
+        
+        if(primalType->equals(builder->getFloatType()))
+            return primalType;
+        else if(auto primalVectorType = as<VectorExpressionType>(primalType))
+        {
+            // TODO(sai): There's probably a more elegant way to check if a type is a float3?
+            if (getIntVal(primalVectorType->elementCount) == 3 && primalVectorType->elementType->equals(builder->getFloatType()))
+                return primalVectorType;
+        }
+        return nullptr;
+    }
+
+    Type* primalToJVPReturnType(ASTBuilder* builder, Type* primalType)
+    {
+        if(auto jvpType = primalToJVPParamType(builder, primalType))
+            return jvpType;
+        else
+            return builder->getVoidType();
+    }
+
     Expr* SemanticsExprVisitor::visitJVPDifferentiateExpr(JVPDifferentiateExpr* expr)
     {
         // Check/Resolve inner function declaration.
@@ -1524,18 +1547,10 @@ namespace Slang
             
             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
-            {
-                //TODO(yong): issue proper diagnostic here.
-                jvpType->resultType = astBuilder->getVoidType();
-            }
+            jvpType->resultType = primalToJVPReturnType(astBuilder, primalType->getResultType());
             
             // No support for differentiating function that throw errors, for now.
             SLANG_ASSERT(primalType->errorType->equals(astBuilder->getBottomType()));
@@ -1548,8 +1563,8 @@ namespace Slang
 
             for (UInt i = 0; i < primalType->getParamCount(); i++)
             {
-                if(primalType->getParamType(i)->equals(astBuilder->getFloatType()))
-                    jvpType->paramTypes.add(astBuilder->getFloatType());
+                if(auto jvpParamType = primalToJVPParamType(astBuilder, primalType->getParamType(i)))
+                    jvpType->paramTypes.add(jvpParamType);
             }
 
             expr->type = jvpType;
diff --git a/source/slang/slang-ir-diff-jvp.cpp b/source/slang/slang-ir-diff-jvp.cpp
index 00210daaa..5b77d483d 100644
--- a/source/slang/slang-ir-diff-jvp.cpp
+++ b/source/slang/slang-ir-diff-jvp.cpp
@@ -38,6 +38,16 @@ struct JVPTranscriber
         return instMapD[instP];
     }
 
+    IRInst* getDifferentialInst(IRInst* instP, IRInst* defaultInst)
+    {
+        return (hasDifferentialInst(instP)) ? instMapD[instP] : defaultInst;
+    }
+
+    bool hasDifferentialInst(IRInst* instP)
+    {
+        return instMapD.ContainsKey(instP);
+    }
+
     IRFuncType* differentiateFunctionType(IRBuilder* builder, IRFuncType* funcType)
     {
         List<IRType*> parameterTypesD;
@@ -74,7 +84,8 @@ struct JVPTranscriber
             case kIROp_FloatType:
             case kIROp_DoubleType:
                 return builder->getType(typeP->getOp());
-            
+            case kIROp_VectorType:
+                return as<IRVectorType>(typeP);
             default:
                 return nullptr;
         }
@@ -252,6 +263,94 @@ struct JVPTranscriber
         return nullptr;
     }
 
+    // Differentiating a call instruction here is primarily about generating
+    // an appropriate call list based on whichever parameters have differentials 
+    // in the current transcription context.
+    // Note(sai): Currently we don't look at modifiers (in, out, const etc..) in the function
+    // type, and so only support 'plain' parameters. We need to validte this somewhere to
+    // avoid weird behaviour
+    // 
+    IRInst* differentiateCall(IRBuilder* builder, IRCall* callP)
+    {   
+        if (auto calleeP = as<IRFunc>(callP->getCallee()))
+        {
+            
+            // Build the differential callee
+            IRInst* calleeD = builder->emitJVPDifferentiateInst(
+                differentiateFunctionType(builder, as<IRFuncType>(calleeP->getFullType())),
+                calleeP);
+            
+            List<IRInst*> args;
+            // Go over the parameter list and all primal arguments.
+            for (UIndex ii = 0; ii < callP->getArgCount(); ii++)
+            {
+                args.add(callP->getArg(ii));
+            }
+
+            {
+                IRParam* param = calleeP->getFirstParam();
+                // Go over the parameter list again and arguments for types that need differentials.
+                for (UIndex ii = 0; ii < callP->getArgCount(); ii++)
+                {
+                    // Look the parameter up in the callee's signature. If it requires a derivative, proceed.
+                    // Otherwise, continue.
+                    //
+                    if (differentiateType(builder, param->getDataType()))
+                    {
+                        // If the corresponding argument does not have a differential, create and place a
+                        // 0 argument.
+                        //
+                        auto argP = callP->getArg(ii);
+                        if (auto argD = getDifferentialInst(argP, nullptr))
+                            args.add(argD);
+                        else
+                            args.add(getZeroOfType(builder, argP->getDataType()));
+                    }
+
+                    param = param->getNextParam();
+                }
+            }
+            
+            return builder->emitCallInst(differentiateType(builder, callP->getFullType()),
+                                         calleeD,
+                                         args);
+        }
+        else
+        {
+            // Note that this can only happen if the callee is a result
+            // of a higher-order operation. For now, we assume that we cannot
+            // differentiate such calls safely.
+            // TODO(sai): Should probably get checked in the front-end.
+            //
+            getSink()->diagnose(callP->sourceLoc,
+                Diagnostics::internalCompilerError,
+                "attempting to differentiate unresolved callee");
+        }
+        return nullptr;
+    }
+
+    // In differential computation, the 'default' differential value is always zero.
+    // This is a consequence of differential computing being inherently linear. As a 
+    // result, it's useful to have a method to generate zero literals of any (arithmetic) type.
+    // 
+    IRInst* getZeroOfType(IRBuilder* builder, IRType* type)
+    {
+        switch (type->getOp())
+        {
+            case kIROp_FloatType:
+            case kIROp_HalfType:
+            case kIROp_DoubleType:
+                return builder->getFloatValue(type, 0.0);
+            case kIROp_IntType:
+                return builder->getIntValue(type, 0);
+            default:
+                getSink()->diagnose(type->sourceLoc,
+                    Diagnostics::internalCompilerError,
+                    "could not generate zero value for given type");
+                return nullptr;       
+        }
+    }
+
     // Logic for whether a primal instruction needs to be replicated
     // in the differential function. For puerly functional blocks with
     // no side-effects, it's safe to replicate everything except the
@@ -307,6 +406,9 @@ struct JVPTranscriber
 
         case kIROp_Construct:
             return differentiateConstruct(builder, instP);
+        
+        case kIROp_Call:
+            return differentiateCall(builder, as<IRCall>(instP));
 
         default:
             getSink()->diagnose(instP->sourceLoc,
diff --git a/source/slang/slang-ir-lower-generic-call.cpp b/source/slang/slang-ir-lower-generic-call.cpp
index 7dbe11f52..2a97bc28a 100644
--- a/source/slang/slang-ir-lower-generic-call.cpp
+++ b/source/slang/slang-ir-lower-generic-call.cpp
@@ -39,7 +39,7 @@ namespace Slang
             if (auto ptrType = as<IRPtrTypeBase>(paramType))
             {
                 paramValType = ptrType->getValueType();
-            }
+            } 
             auto argType = arg->getDataType();
             if (auto argPtrType = as<IRPtrTypeBase>(argType))
             {