Type legalization and autodiff bug fixes. (#2722)

* Bug fixes.

* Fix.

* Only perform autodiff for functions whose derivative is actually used.

* Fix loop optimize bug.

* Fix high order diff.

* Fix trivial diff func generation.

* Fixes.

* Cleanup.

---------

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

1 files changed, 20 insertions, 9 deletions

diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index 050c9bfc7..05c28d131 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -62,7 +62,7 @@ struct SpecializationContext
     // if it is in our set.
     //
     bool isInstFullySpecialized(
-        IRInst*                     inst)
+        IRInst* inst)
     {
         // A small wrinkle is that a null instruction pointer
         // sometimes appears a a type, and so should be treated
@@ -70,7 +70,7 @@ struct SpecializationContext
         //
         // TODO: It would be nice to remove this wrinkle.
         //
-        if(!inst) return true;
+        if (!inst) return true;
 
         // An interface requirement entry should always be considered
         // to be fully specialized, even if it hasn't been visited.
@@ -79,7 +79,7 @@ struct SpecializationContext
         // can't mark an interface as used until its requirements are
         // used, etc.
         //
-        if(inst->getOp() == kIROp_InterfaceRequirementEntry)
+        if (inst->getOp() == kIROp_InterfaceRequirementEntry)
             return true;
 
         // A generic parameter is never specialized.
@@ -93,6 +93,20 @@ struct SpecializationContext
                 inst->getParent()->getParent()->getOp() == kIROp_Generic)
                 return false;
         }
+
+        // A global value is always specialized.
+        if (inst->getParent() == module->getModuleInst())
+        {
+            switch (inst->getOp())
+            {
+            case kIROp_LookupWitness:
+            case kIROp_Specialize:
+                return false;
+            default:
+                return true;
+            }
+        }
+
         return fullySpecializedInsts.Contains(inst);
     }
 
@@ -504,12 +518,9 @@ struct SpecializationContext
             // be considered as fully specialized as soon
             // as all of its operands are.
             //
-            // TODO: We realistically need a more refined
-            // check here that uses an allow-list of instructions
-            // that can represent values suitable for use
-            // as generic arguments.
-            //
-            if(areAllOperandsFullySpecialized(inst))
+            // Anything defined in global scope can be viewed as fully specialized.
+            if (inst->getParent() == module->getModuleInst() ||
+                areAllOperandsFullySpecialized(inst))
             {
                 markInstAsFullySpecialized(inst);
             }