Fix bug in legalizeFuncType that leads to invalid IR. (#2902)

* Fix bug in legalizeFuncType that leads to invalid IR.

* Diagnose on functions that never returns when differentiate it.

---------

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

1 files changed, 8 insertions, 1 deletions

diff --git a/source/slang/slang-ir-autodiff-rev.cpp b/source/slang/slang-ir-autodiff-rev.cpp
index 70c43cdcb..3f0036b06 100644
--- a/source/slang/slang-ir-autodiff-rev.cpp
+++ b/source/slang/slang-ir-autodiff-rev.cpp
@@ -537,10 +537,17 @@ namespace Slang
         DifferentiableTypeConformanceContext diffTypeContext(autoDiffSharedContext);
         diffTypeContext.setFunc(func);
 
-        if (!isSingleReturnFunc(func))
+        auto returnCount = getReturnCount(func);
+        if (returnCount > 1)
         {
             convertFuncToSingleReturnForm(func->getModule(), func);
         }
+        else if (returnCount == 0)
+        {
+            // The function is ill-formed and never returns (such as having an infinite loop),
+            // we can't possibly reverse-differentiate such functions, so we will diagnose it here.
+            getSink()->diagnose(func->sourceLoc, Diagnostics::functionNeverReturnsFatal, func);
+        }
 
         eliminateContinueBlocksInFunc(func->getModule(), func);