Small fixes to autodiff pass. (#2852)

* Remove obsolete assert in `extractPriamlFunc`.

* Fix `findUniqueStoredVal` to ignore diff uses.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

2 files changed, 13 insertions, 8 deletions

diff --git a/source/slang/slang-ir-autodiff-primal-hoist.cpp b/source/slang/slang-ir-autodiff-primal-hoist.cpp
index bef96f309..dd7472067 100644
--- a/source/slang/slang-ir-autodiff-primal-hoist.cpp
+++ b/source/slang/slang-ir-autodiff-primal-hoist.cpp
@@ -1044,13 +1044,6 @@ RefPtr<HoistedPrimalsInfo> ensurePrimalAvailability(
                 for (auto use : outOfScopeUses)
                 {
                     List<IndexTrackingInfo> useBlockIndices = indexedBlockInfo[getBlock(use->getUser())];
-                    if (isLoopCounter)
-                    {
-                        // The use site of a primal loop counter should be right before we enter the
-                        // loop, and therefore its index count should equal to defBlockIndices.getCount()
-                        // after we remove the first index from defBlockIndices.
-                        SLANG_RELEASE_ASSERT(useBlockIndices.getCount() == defBlockIndices.getCount());
-                    }
                     setInsertBeforeOrdinaryInst(&builder, getInstInBlock(use->getUser()));
                     builder.replaceOperand(use, loadIndexedValue(&builder, localVar, defBlockIndices, useBlockIndices));
                 }
diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp
index 4889558a5..62cb0b841 100644
--- a/source/slang/slang-ir-autodiff.cpp
+++ b/source/slang/slang-ir-autodiff.cpp
@@ -1858,6 +1858,9 @@ IRUse* findUniqueStoredVal(IRVar* var)
         {
             if (const auto callInst = as<IRCall>(use->getUser()))
             {
+                // Ignore uses from differential blocks.
+                if (callInst->getParent()->findDecoration<IRDifferentialInstDecoration>())
+                    continue;
                 // Should not see more than one IRCall. If we do
                 // we'll need to pick the primal call.
                 // 
@@ -1874,6 +1877,9 @@ IRUse* findUniqueStoredVal(IRVar* var)
         {
             if (const auto storeInst = as<IRStore>(use->getUser()))
             {
+                // Ignore uses from differential blocks.
+                if (storeInst->getParent()->findDecoration<IRDifferentialInstDecoration>())
+                    continue;
                 // Should not see more than one IRStore
                 SLANG_RELEASE_ASSERT(!storeUse);
                 storeUse = use;
@@ -1890,15 +1896,18 @@ IRUse* findUniqueStoredVal(IRVar* var)
 IRUse* findLatestUniqueWriteUse(IRVar* var)
 {
     IRUse* storeUse = nullptr;
-    // If no unique store found, try to look for a call.
     for (auto use = var->firstUse; use; use = use->nextUse)
     {
         if (const auto callInst = as<IRCall>(use->getUser()))
         {
+            // Ignore uses from differential blocks.
+            if (callInst->getParent()->findDecoration<IRDifferentialInstDecoration>())
+                continue;
             SLANG_RELEASE_ASSERT(!storeUse);
             storeUse = use;
         }
     }
+    // If no unique call found, try to look for a store.
     return findUniqueStoredVal(var);
 }
 
@@ -1917,6 +1926,9 @@ IRUse* findEarliestUniqueWriteUse(IRVar* var)
     {
         if (const auto callInst = as<IRCall>(use->getUser()))
         {
+            // Ignore uses from differential blocks.
+            if (callInst->getParent()->findDecoration<IRDifferentialInstDecoration>())
+                continue;
             SLANG_RELEASE_ASSERT(!storeUse);
             storeUse = use;
         }