From 499b0253c224e68ceed6e5b6b1ee9cd7d65aad0f Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 30 Jan 2023 19:24:09 -0800
Subject: Make ArrayExpressionType a DeclRefType and define its autodiff
 extension in stdlib. (#2615)

* Allow array parameters in forward diff.

* Use type canonicalization instead of coersion.

* Reimplement array type.

* Fix.

* Update test case.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 .../slang/slang-ir-autodiff-transcriber-base.cpp   | 35 ++++++++++++++--------
 1 file changed, 22 insertions(+), 13 deletions(-)

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

diff --git a/source/slang/slang-ir-autodiff-transcriber-base.cpp b/source/slang/slang-ir-autodiff-transcriber-base.cpp
index 05a5f8f56..91374e006 100644
--- a/source/slang/slang-ir-autodiff-transcriber-base.cpp
+++ b/source/slang/slang-ir-autodiff-transcriber-base.cpp
@@ -194,27 +194,19 @@ IRWitnessTable* AutoDiffTranscriberBase::getDifferentialPairWitness(IRBuilder* b
     return table;
 }
 
-IRType* AutoDiffTranscriberBase::getOrCreateDiffPairType(IRBuilder* builder, IRInst* primalType, IRInst* witness)
+IRInst* AutoDiffTranscriberBase::tryGetDifferentiableWitness(IRBuilder* builder, IRInst* originalType)
 {
-    return builder->getDifferentialPairType(
-        (IRType*)primalType,
-        witness);
-}
-
-IRType* AutoDiffTranscriberBase::getOrCreateDiffPairType(IRBuilder* builder, IRInst* originalType)
-{
-    auto primalType = lookupPrimalInst(builder, originalType, nullptr);
-    SLANG_RELEASE_ASSERT(primalType);
-
-    IRInst* witness = 
+    IRInst* witness =
         differentiableTypeConformanceContext.lookUpConformanceForType((IRType*)originalType);
     if (witness)
     {
         witness = lookupPrimalInst(builder, witness, nullptr);
-        SLANG_RELEASE_ASSERT(witness);
+        SLANG_RELEASE_ASSERT(witness || as<IRArrayType>(originalType));
     }
     if (!witness)
     {
+        auto primalType = lookupPrimalInst(builder, originalType, nullptr);
+        SLANG_RELEASE_ASSERT(primalType);
         if (auto primalPairType = as<IRDifferentialPairType>(primalType))
         {
             witness = getDifferentialPairWitness(builder, originalType, primalPairType);
@@ -224,6 +216,23 @@ IRType* AutoDiffTranscriberBase::getOrCreateDiffPairType(IRBuilder* builder, IRI
             differentiateExtractExistentialType(builder, extractExistential, witness);
         }
     }
+    return witness;
+}
+
+IRType* AutoDiffTranscriberBase::getOrCreateDiffPairType(IRBuilder* builder, IRInst* primalType, IRInst* witness)
+{
+    return builder->getDifferentialPairType(
+        (IRType*)primalType,
+        witness);
+}
+
+IRType* AutoDiffTranscriberBase::getOrCreateDiffPairType(IRBuilder* builder, IRInst* originalType)
+{
+    auto primalType = lookupPrimalInst(builder, originalType, nullptr);
+    SLANG_RELEASE_ASSERT(primalType);
+
+    IRInst* witness = tryGetDifferentiableWitness(builder, originalType);
+    SLANG_RELEASE_ASSERT(witness);
 
     return builder->getDifferentialPairType(
         (IRType*)primalType,
-- 
cgit v1.2.3