From 87f00a36a123e36b415eeea82e02a8366cc5b881 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Fri, 10 Jan 2025 03:16:24 +0530
Subject: [Auto-diff] Overhaul auto-diff type tracking + Overhaul dynamic
 dispatch for differentiable functions (#5866)

* Overhauled the auto-diff system for dynamic dispatch

* More fixes

* remove intermediate dumps

* Update slang-ast-type.h

* More fixes + add a workaround for existential no-diff

* Update reverse-control-flow-3.slang

* remove dumps

* remove more dumps

* Delete working-reverse-control-flow-3.hlsl

* Cleanup comments + unused variables

* More comment cleanup

* Add support for lowering `DiffPairType(TypePack)` & `MakePair(MakeValuePack, MakeValuePack)`

* Fix array of issues in Falcor tests.

* Update slang-ir-autodiff-pairs.cpp

* More fixes for Falcor image tests

* Small fixups.

---------

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

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

diff --git a/source/slang/slang-ir-autodiff.h b/source/slang/slang-ir-autodiff.h
index 2b03f3923..433b6093f 100644
--- a/source/slang/slang-ir-autodiff.h
+++ b/source/slang/slang-ir-autodiff.h
@@ -221,6 +221,8 @@ struct DifferentiableTypeConformanceContext
     IRGlobalValueWithCode* parentFunc = nullptr;
     OrderedDictionary<IRType*, IRInst*> differentiableTypeWitnessDictionary;
 
+    Dictionary<IRInst*, List<IRDifferentiableTypeAnnotation*>> annotationCache;
+
     IRFunc* existentialDAddFunc = nullptr;
 
     DifferentiableTypeConformanceContext(AutoDiffSharedContext* shared)
@@ -235,6 +237,10 @@ struct DifferentiableTypeConformanceContext
 
     void setFunc(IRGlobalValueWithCode* func);
 
+    List<IRDifferentiableTypeAnnotation*> getAnnotations(IRGlobalValueWithCode* inst);
+
+    List<IRDifferentiableTypeAnnotation*> getAnnotations(IRModuleInst* inst);
+
     void buildGlobalWitnessDictionary();
 
     // Lookup a witness table for the concreteType. One should exist if concreteType
@@ -445,6 +451,20 @@ struct DifferentiableTypeConformanceContext
         IRBuilder* builder,
         IRExtractExistentialType* extractExistentialType,
         DiffConformanceKind target);
+
+    IRInst* emitDAddOfDiffInstType(
+        IRBuilder* builder,
+        IRType* primalType,
+        IRInst* op1,
+        IRInst* op2);
+
+    IRInst* emitDAddForExistentialType(
+        IRBuilder* builder,
+        IRType* primalType,
+        IRInst* op1,
+        IRInst* op2);
+
+    IRInst* emitDZeroOfDiffInstType(IRBuilder* builder, IRType* primalType);
 };
 
 
@@ -461,9 +481,15 @@ struct DifferentialPairTypeBuilder
 
     IRInst* emitFieldAccessor(IRBuilder* builder, IRInst* baseInst, IRStructKey* key);
 
-    IRInst* emitPrimalFieldAccess(IRBuilder* builder, IRInst* baseInst);
+    IRInst* emitPrimalFieldAccess(IRBuilder* builder, IRType* loweredPairType, IRInst* baseInst);
+
+    IRInst* emitDiffFieldAccess(IRBuilder* builder, IRType* loweredPairType, IRInst* baseInst);
 
-    IRInst* emitDiffFieldAccess(IRBuilder* builder, IRInst* baseInst);
+    IRInst* emitExistentialMakePair(
+        IRBuilder* builder,
+        IRInst* type,
+        IRInst* primalInst,
+        IRInst* diffInst);
 
     IRStructKey* _getOrCreateDiffStructKey();
 
@@ -471,17 +497,52 @@ struct DifferentialPairTypeBuilder
 
     IRInst* _createDiffPairType(IRType* origBaseType, IRType* diffType);
 
+    IRInst* _createDiffPairInterfaceRequirement(IRType* origBaseType, IRType* diffType);
+
     IRInst* lowerDiffPairType(IRBuilder* builder, IRType* originalPairType);
 
+    IRInst* getOrCreateCommonDiffPairInterface(IRBuilder* builder);
+
     struct PairStructKey
     {
         IRInst* originalType;
         IRInst* diffType;
     };
 
-    // Cache from `IRDifferentialPairType` to materialized struct type.
+    // Cache from pair types to lowered type.
     Dictionary<IRInst*, IRInst*> pairTypeCache;
 
+    // Cache from existential pair types to their lowered interface keys.
+    // We use a different cache because an interface type can have
+    // a regular pair for the pair of interface types, as well as an
+    // interface key for the associated pair types used for its implementations
+    //
+    Dictionary<IRInst*, IRInst*> existentialPairTypeCache;
+
+    // Cache for any interface requirement keys (generated for existential
+    // pair types)
+    //
+    Dictionary<IRInst*, IRStructKey*> assocPairTypeKeyMap;
+    Dictionary<IRInst*, IRStructKey*> makePairKeyMap;
+    Dictionary<IRInst*, IRStructKey*> getPrimalKeyMap;
+    Dictionary<IRInst*, IRStructKey*> getDiffKeyMap;
+
+    // More caches for easier lookups of the types associated with the
+    // keys. (avoid having to keep recomputing or performing complicated
+    // lookups)
+    //
+    Dictionary<IRInst*, IRFuncType*> makePairFuncTypeMap;
+    Dictionary<IRInst*, IRFuncType*> getPrimalFuncTypeMap;
+    Dictionary<IRInst*, IRFuncType*> getDiffFuncTypeMap;
+
+    // Even more caches for easier access to original primal/diff types
+    // (Only used for existential pair types). For regular pair types,
+    // these are easy to find right on the type itself.
+    //
+    Dictionary<IRInst*, IRType*> primalTypeMap;
+    Dictionary<IRInst*, IRType*> diffTypeMap;
+
+
     IRStructKey* globalPrimalKey = nullptr;
 
     IRStructKey* globalDiffKey = nullptr;
@@ -491,6 +552,8 @@ struct DifferentialPairTypeBuilder
     List<IRInst*> generatedTypeList;
 
     AutoDiffSharedContext* sharedContext = nullptr;
+
+    IRInterfaceType* commonDiffPairInterface = nullptr;
 };
 
 void stripAutoDiffDecorations(IRModule* module);
@@ -551,6 +614,10 @@ inline bool isRelevantDifferentialPair(IRType* type)
     return false;
 }
 
+bool isRuntimeType(IRType* type);
+
+IRInst* getExistentialBaseWitnessTable(IRBuilder* builder, IRType* type);
+
 UIndex addPhiOutputArg(
     IRBuilder* builder,
     IRBlock* block,
-- 
cgit v1.2.3