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/slang-check-expr.cpp | 66 ++++++++++++++++----------------------- 1 file changed, 27 insertions(+), 39 deletions(-) (limited to 'source/slang/slang-check-expr.cpp') diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp index b7f99c4e7..a787af211 100644 --- a/source/slang/slang-check-expr.cpp +++ b/source/slang/slang-check-expr.cpp @@ -1526,48 +1526,42 @@ namespace Slang return expr; } - // This function proceses primal params (i.e params of the inner function that is being - // differentiated) that need to be carried over to the function signature for the JVP - // function. (eg. out types can be discarded) - // - Type* primalToInputType(ASTBuilder*, Type* primalType) - { - if (auto primalOutType = as(primalType)) - return nullptr; - else if (auto primalInOutType = as(primalType)) - return primalInOutType->getValueType(); - - return primalType; - } - Type* primalToJVPParamType(ASTBuilder* builder, Type* primalType) + Type* SemanticsVisitor::_toDifferentialParamType(ASTBuilder* builder, Type* primalType) { - // Only float and vector types can be differentiated for now. - - if (primalType->equals(builder->getFloatType())) - return primalType; - else if (auto primalVectorType = as(primalType)) - { - if (auto jvpElementType = primalToJVPParamType(builder, primalVectorType->elementType)) - return builder->getVectorType(jvpElementType, primalVectorType->elementCount); - } - else if (auto primalOutType = as(primalType)) + // Check for type modifiers like 'out' and 'inout'. We need to differentiate the + // nested type. + // + if (auto primalOutType = as(primalType)) { - return builder->getOutType(primalToJVPParamType(builder, primalOutType->getValueType())); + return builder->getOutType(_toDifferentialParamType(builder, primalOutType->getValueType())); } else if (auto primalInOutType = as(primalType)) { - return builder->getInOutType(primalToJVPParamType(builder, primalInOutType->getValueType())); + return builder->getInOutType(_toDifferentialParamType(builder, primalInOutType->getValueType())); } - return nullptr; + + // Get a reference to the builtin 'IDifferentiable' interface + auto differentiableInterface = builder->getDifferentiableInterface(); + + // Check if the provided type inherits from IDifferentiable. + // If not, return the original type. + if (auto conformanceWitness = as(tryGetInterfaceConformanceWitness(primalType, differentiableInterface))) + return builder->getDifferentialPairType(primalType, conformanceWitness); + else + return primalType; + } - Type* primalToJVPReturnType(ASTBuilder* builder, Type* primalType) + Type* SemanticsVisitor::_toJVPReturnType(ASTBuilder* builder, Type* primalType) { - if(auto jvpType = primalToJVPParamType(builder, primalType)) - return jvpType; + if (auto conformanceWitness = + as(tryGetInterfaceConformanceWitness( + primalType, + builder->getDifferentiableInterface()))) + return builder->getDifferentialPairType(primalType, conformanceWitness); else - return builder->getVoidType(); + return primalType; } Expr* SemanticsExprVisitor::visitJVPDifferentiateExpr(JVPDifferentiateExpr* expr) @@ -1588,7 +1582,7 @@ namespace Slang // The JVP return type is float if primal return type is float // void otherwise. // - jvpType->resultType = primalToJVPReturnType(astBuilder, primalType->getResultType()); + jvpType->resultType = _toJVPReturnType(astBuilder, primalType->getResultType()); // No support for differentiating function that throw errors, for now. SLANG_ASSERT(primalType->errorType->equals(astBuilder->getBottomType())); @@ -1596,13 +1590,7 @@ namespace Slang for (UInt i = 0; i < primalType->getParamCount(); i++) { - if(auto primalInputType = primalToInputType(astBuilder, primalType->getParamType(i))) - jvpType->paramTypes.add(primalInputType); - } - - for (UInt i = 0; i < primalType->getParamCount(); i++) - { - if(auto jvpParamType = primalToJVPParamType(astBuilder, primalType->getParamType(i))) + if(auto jvpParamType = _toDifferentialParamType(astBuilder, primalType->getParamType(i))) jvpType->paramTypes.add(jvpParamType); } -- cgit v1.2.3