From ccc310fa4e8096cda8a6c127aacc1a1fa9d8503a Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Thu, 19 Sep 2024 03:10:28 -0400
Subject: Support `IDifferentiablePtrType` (#5031)

* initial diff-ref-type interface

* Initial support for `IDifferentiablePtrType`

* Fix unused vars

* More tests + fix switch case fallthrough.

* Update slang-ir-autodiff.cpp

* Update diff-ptr-type-loop.slang

* Add optimization to allow more complex pair types

* Update slang-ir-autodiff-primal-hoist.cpp

* Update diff-ptr-type-loop.slang

* Update slang-ir-autodiff-primal-hoist.cpp

* More fixes to address reviews

* Update slang-check-expr.cpp

* Optimizations + rename `differentiableRefInterfaceType` -> `differentiablePtrInterfaceType`

* Move pair logic to ir-builder, unify the type dictionaries.

---------

Co-authored-by: Yong He <yonghe@outlook.com>
---
 source/slang/slang-ir-autodiff-rev.cpp | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

(limited to 'source/slang/slang-ir-autodiff-rev.cpp')

diff --git a/source/slang/slang-ir-autodiff-rev.cpp b/source/slang/slang-ir-autodiff-rev.cpp
index 2fb73c4ac..169dd31ee 100644
--- a/source/slang/slang-ir-autodiff-rev.cpp
+++ b/source/slang/slang-ir-autodiff-rev.cpp
@@ -152,6 +152,16 @@ namespace Slang
         builder->emitBlock();
         params = _defineFuncParams(builder, as<IRFunc>(existingPrimalFunc));
         params.removeLast();
+
+        // Unwrap any ref pairs. We need this special case for trivial funcs.
+        for (Int i = 0; i < params.getCount(); i++)
+        {
+            if (as<IRDifferentialPtrPairType>(params[i]->getDataType()))
+            {
+                params[i] = builder->emitDifferentialPtrPairGetPrimal(params[i]);
+            }
+        }
+
         IRInst* originalFuncRefFromPrimalFunc = originalFunc;
         if (originalGeneric)
             originalFuncRefFromPrimalFunc = maybeSpecializeWithGeneric(*builder, originalGeneric, existingPriamlFuncGeneric);
@@ -266,7 +276,20 @@ namespace Slang
         if (auto primalNoDiffType = _getPrimalTypeFromNoDiffType(this, builder, paramType))
             return primalNoDiffType;
 
-        return (IRType*)findOrTranscribePrimalInst(builder, paramType);
+        auto primalType = (IRType*)findOrTranscribePrimalInst(builder, paramType);
+        
+        // Differentiable pointer types are treated as primal pairs, since they aren't involved in the transposition
+        // process.
+        //
+        if (differentiableTypeConformanceContext.isDifferentiablePtrType(primalType))
+        {
+            auto diffPairType = tryGetDiffPairType(builder, primalType);
+            SLANG_ASSERT(diffPairType);
+
+            return diffPairType;
+        }
+
+        return primalType;
     }
 
     IRType* BackwardDiffTranscriberBase::transcribeParamTypeForPropagateFunc(IRBuilder* builder, IRType* paramType)
@@ -292,7 +315,7 @@ namespace Slang
         auto diffPairType = tryGetDiffPairType(builder, paramType);
         if (diffPairType)
         {
-            if (!as<IRPtrTypeBase>(diffPairType))
+            if (!as<IRPtrTypeBase>(diffPairType) && !as<IRDifferentialPtrPairType>(diffPairType))
                 return builder->getInOutType(diffPairType);
             return diffPairType;
         }
@@ -961,7 +984,7 @@ namespace Slang
                     // Initialize the var with input diff param at start.
                     // Note that we insert the store in the primal block so it won't get transposed.
                     auto storeInst = nextBlockBuilder.emitStore(tempVar, diffParam);
-                    nextBlockBuilder.markInstAsDifferential(storeInst, diffPairType);
+                    nextBlockBuilder.markInstAsDifferential(storeInst, primalType);
                     // Since this store inst is specific to propagate function, we track it in a
                     // set so we can remove it when we generate the primal func.
                     result.propagateFuncSpecificPrimalInsts.add(storeInst);
-- 
cgit v1.2.3