Support IDifferentiablePtrType

source/slang/core.meta.slang

@@ -285,6 +285,13 @@ interface IDifferentiable
     static Differential dmul(T, Differential);
 };
 
+__magic_type(DifferentiablePtrType)
+interface IDifferentiablePtrType
+{
+    __builtin_requirement($( (int)BuiltinRequirementKind::DifferentialPtrType) )
+    associatedtype Differential : IDifferentiablePtrType;
+};
+
 
 /// Pair type that serves to wrap the primal and
 /// differential types of an arbitrary type T.
@@ -357,6 +364,36 @@ struct DifferentialPair : IDifferentiable
     }
 };
 
+__generic<T : IDifferentiablePtrType>
+__magic_type(DifferentialPtrPairType)
+__intrinsic_type($(kIROp_DifferentialPtrPairType))
+struct DifferentialPtrPair : IDifferentiablePtrType
+{
+    typedef DifferentialPtrPair<T.Differential> Differential;
+    typedef T.Differential DifferentialElementType;
+
+    __intrinsic_op($(kIROp_MakeDifferentialPtrPair))
+    __init(T _primal, T.Differential _differential);
+
+    property p : T
+    {
+        __intrinsic_op($(kIROp_DifferentialPtrPairGetPrimal))
+        get;
+    }
+
+    property v : T
+    {
+        __intrinsic_op($(kIROp_DifferentialPtrPairGetPrimal))
+        get;
+    }
+
+    property d : T.Differential
+    {
+        __intrinsic_op($(kIROp_DifferentialPtrPairGetDifferential))
+        get;
+    }
+};
+
 
 /// A type that uses a floating-point representation
 [sealed]

source/slang/slang-ast-builder.cpp

@@ -408,18 +408,32 @@ MatrixExpressionType* ASTBuilder::getMatrixType(Type* elementType, IntVal* rowCo
 
 DifferentialPairType* ASTBuilder::getDifferentialPairType(
     Type* valueType,
-    Witness* primalIsDifferentialWitness)
+    Witness* diffTypeWitness)
 {
-    Val* args[] = { valueType, primalIsDifferentialWitness };
+    Val* args[] = { valueType, diffTypeWitness };
     return as<DifferentialPairType>(getSpecializedBuiltinType(makeArrayView(args), "DifferentialPairType"));
 }
 
+DifferentialPtrPairType* ASTBuilder::getDifferentialPtrPairType(
+    Type* valueType,
+    Witness* diffRefTypeWitness)
+{
+    Val* args[] = { valueType, diffRefTypeWitness };
+    return as<DifferentialPtrPairType>(getSpecializedBuiltinType(makeArrayView(args), "DifferentialPtrPairType"));
+}
+
 DeclRef<InterfaceDecl> ASTBuilder::getDifferentiableInterfaceDecl()
 {
     DeclRef<InterfaceDecl> declRef = DeclRef<InterfaceDecl>(getBuiltinDeclRef("DifferentiableType", nullptr));
     return declRef;
 }
 
+DeclRef<InterfaceDecl> ASTBuilder::getDifferentiableRefInterfaceDecl()
+{
+    DeclRef<InterfaceDecl> declRef = DeclRef<InterfaceDecl>(getBuiltinDeclRef("DifferentiablePtrType", nullptr));
+    return declRef;
+}
+
 bool ASTBuilder::isDifferentiableInterfaceAvailable()
 {
     return (m_sharedASTBuilder->tryFindMagicDecl("DifferentiableType") != nullptr);
@@ -459,6 +473,11 @@ Type* ASTBuilder::getDifferentiableInterfaceType()
     return DeclRefType::create(this, getDifferentiableInterfaceDecl());
 }
 
+Type* ASTBuilder::getDifferentiableRefInterfaceType()
+{
+    return DeclRefType::create(this, getDifferentiableRefInterfaceDecl());
+}
+
 DeclRef<Decl> ASTBuilder::getBuiltinDeclRef(const char* builtinMagicTypeName, Val* genericArg)
 {
     auto decl = m_sharedASTBuilder->findMagicDecl(builtinMagicTypeName);

source/slang/slang-ast-builder.h

@@ -489,10 +489,17 @@ class ASTBuilder : public RefObject
 
     DifferentialPairType* getDifferentialPairType(
         Type* valueType,
-        Witness* primalIsDifferentialWitness);
+        Witness* diffTypeWitness);
+
+    DifferentialPtrPairType* getDifferentialPtrPairType(
+        Type* valueType,
+        Witness* diffRefTypeWitness);
 
     DeclRef<InterfaceDecl> getDifferentiableInterfaceDecl();
+    DeclRef<InterfaceDecl> getDifferentiableRefInterfaceDecl();
+
     Type* getDifferentiableInterfaceType();
+    Type* getDifferentiableRefInterfaceType();
 
     bool isDifferentiableInterfaceAvailable();
 

source/slang/slang-ast-support-types.h

@@ -9,7 +9,7 @@
 
 #include "slang-profile.h"
 #include "slang-type-system-shared.h"
-#include "slang.h"
+#include "../../include/slang.h"
 
 #include "../core/slang-semantic-version.h"
 
@@ -1606,6 +1606,7 @@ namespace Slang
         DefaultInitializableConstructor, ///< The `IDefaultInitializable.__init()` method
 
         DifferentialType, ///< The `IDifferentiable.Differential` associated type requirement 
+        DifferentialPtrType, ///< The `IDifferentiable.DifferentialPtr` associated type requirement
         DZeroFunc, ///< The `IDifferentiable.dzero` function requirement 
         DAddFunc, ///< The `IDifferentiable.dadd` function requirement 
         DMulFunc, ///< The `IDifferentiable.dmul` function requirement 

source/slang/slang-ast-type.h

@@ -462,11 +462,22 @@ class DifferentialPairType : public ArithmeticExpressionType
     Type* getPrimalType();
 };
 
+class DifferentialPtrPairType : public ArithmeticExpressionType 
+{
+    SLANG_AST_CLASS(DifferentialPtrPairType)
+    Type* getPrimalRefType();
+};
+
 class DifferentiableType : public BuiltinType
 {
     SLANG_AST_CLASS(DifferentiableType)
 };
 
+class DifferentiablePtrType : public BuiltinType
+{
+    SLANG_AST_CLASS(DifferentiablePtrType)
+};
+
 class DefaultInitializableType : public BuiltinType
 {
     SLANG_AST_CLASS(DefaultInitializableType);

source/slang/slang-check-conformance.cpp

@@ -274,9 +274,14 @@ namespace Slang
         return isInterfaceType(type);
     }
 
-    bool SemanticsVisitor::isTypeDifferentiable(Type* type)
+    SubtypeWitness* SemanticsVisitor::isTypeDifferentiable(Type* type)
     {
-        return isSubtype(type, m_astBuilder->getDiffInterfaceType(), IsSubTypeOptions::None);
+        if (auto valueWitness = isSubtype(type, m_astBuilder->getDiffInterfaceType(), IsSubTypeOptions::None))
+            return valueWitness;
+        else if (auto ptrWitness = isSubtype(type, m_astBuilder->getDifferentiableRefInterfaceType(), IsSubTypeOptions::None))
+            return ptrWitness;
+
+        return nullptr;
     }
 
     bool SemanticsVisitor::doesTypeHaveTag(Type* type, TypeTag tag)

source/slang/slang-check-decl.cpp

@@ -10204,7 +10204,8 @@ namespace Slang
             bool isDiffParam = (!param->findModifier<NoDiffModifier>());
             if (isDiffParam)
             {
-                if (auto pairType = as<DifferentialPairType>(visitor->getDifferentialPairType(param->getType())))
+                auto diffPair = visitor->getDifferentialPairType(param->getType());
+                if (auto pairType = as<DifferentialPairType>(diffPair))
                 {
                     arg->type.type = pairType;
                     arg->type.isLeftValue = true;
@@ -10225,6 +10226,11 @@ namespace Slang
                         direction = ParameterDirection::kParameterDirection_InOut;
                     }
                 }
+                else if (auto refPairType = as<DifferentialPtrPairType>(diffPair))
+                {
+                    // no need to change direction of ref-pairs.
+                    arg->type.type = refPairType;
+                }
                 else
                 {
                     isDiffParam = false;

source/slang/slang-check-expr.cpp

@@ -1131,15 +1131,19 @@ namespace Slang
         {
             if (auto builtinRequirement = declRefType->getDeclRef().getDecl()->findModifier<BuiltinRequirementModifier>())
             {
-                if (builtinRequirement->kind == BuiltinRequirementKind::DifferentialType)
+                if (builtinRequirement->kind == BuiltinRequirementKind::DifferentialType
+                    || builtinRequirement->kind == BuiltinRequirementKind::DifferentialPtrType)
                 {
                     // We are trying to get differential type from a differential type.
                     // The result is itself.
                     return type;
                 }
             }
             type = resolveType(type);
-            if (const auto witness = as<SubtypeWitness>(tryGetInterfaceConformanceWitness(type, builder->getDifferentiableInterfaceType())))
+            auto witness = as<SubtypeWitness>(tryGetInterfaceConformanceWitness(type, builder->getDifferentiableInterfaceType()));
+            if (!witness)
+                witness = as<SubtypeWitness>(tryGetInterfaceConformanceWitness(type, builder->getDifferentiableRefInterfaceType()));
+            if (witness)
             {
                 auto diffTypeLookupResult = lookUpMember(
                     getASTBuilder(),
@@ -1367,6 +1371,13 @@ namespace Slang
             {
                 addDifferentiableTypeToDiffTypeRegistry((DeclRefType*)type, subtypeWitness);
             }
+
+            if (auto subtypeWitness = as<SubtypeWitness>(
+                tryGetInterfaceConformanceWitness(type, getASTBuilder()->getDifferentiableRefInterfaceType())))
+            {
+                addDifferentiableTypeToDiffTypeRegistry((DeclRefType*)type, subtypeWitness);
+            }
+
             if (auto aggTypeDeclRef = declRefType->getDeclRef().as<AggTypeDecl>())
             {
                 foreachDirectOrExtensionMemberOfType<InheritanceDecl>(this, aggTypeDeclRef, [&](DeclRef<InheritanceDecl> member)
@@ -2899,18 +2910,25 @@ namespace Slang
                 return m_astBuilder->getExpandType(diffPairEachType, makeArrayViewSingle(primalType));
             }
         }
+
         // Get a reference to the builtin 'IDifferentiable' interface
         auto differentiableInterface = getASTBuilder()->getDifferentiableInterfaceType();
+        auto differentiableRefInterface = getASTBuilder()->getDifferentiableRefInterfaceType();
 
-        auto conformanceWitness = as<Witness>(isSubtype(primalType, differentiableInterface, IsSubTypeOptions::None));
         // Check if the provided type inherits from IDifferentiable.
         // If not, return the original type.
-        if (conformanceWitness)
+        if (auto conformanceWitness = isTypeDifferentiable(primalType))
         {
-            return m_astBuilder->getDifferentialPairType(primalType, conformanceWitness);
+            if (conformanceWitness->getSup() == differentiableInterface)
+            {
+                return m_astBuilder->getDifferentialPairType(primalType, conformanceWitness);
+            }
+            else if (conformanceWitness->getSup() == differentiableRefInterface)
+            {
+                return m_astBuilder->getDifferentialPtrPairType(primalType, conformanceWitness);
+            }
         }
-        else
-            return primalType;
+        return primalType;
     }
 
     Type* SemanticsVisitor::getForwardDiffFuncType(FuncType* originalType)

source/slang/slang-check-impl.h

@@ -2208,7 +2208,7 @@ namespace Slang
 
         bool isValidGenericConstraintType(Type* type);
 
-        bool isTypeDifferentiable(Type* type);
+        SubtypeWitness* isTypeDifferentiable(Type* type);
 
         bool doesTypeHaveTag(Type* type, TypeTag tag);