Emit errors for missing returns on unsupported targets (#6633)

* initial wip

* more WIP

* preserve old lower behavior

* remove unnecessary includes

* add test

* add no target case in test

* fix broken test

---------

Co-authored-by: Ellie Hermaszewska <ellieh@nvidia.com>

9 files changed, 125 insertions, 12 deletions

diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index d4e05bc48..18192678a 100644
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1968,6 +1968,7 @@ bool isMetalTarget(TargetRequest* targetReq);
 
 /// Are we generating code for a Khronos API (OpenGL or Vulkan)?
 bool isKhronosTarget(TargetRequest* targetReq);
+bool isKhronosTarget(CodeGenTarget target);
 
 /// Are we generating code for a CUDA API (CUDA / OptiX)?
 bool isCUDATarget(TargetRequest* targetReq);
@@ -1977,6 +1978,7 @@ bool isCPUTarget(TargetRequest* targetReq);
 
 /// Are we generating code for the WebGPU API?
 bool isWGPUTarget(TargetRequest* targetReq);
+bool isWGPUTarget(CodeGenTarget target);
 
 /// A request to generate output in some target format.
 class TargetRequest : public RefObject
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index e9fa72bce..b7469deb0 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -2158,6 +2158,11 @@ DIAGNOSTIC(
 DIAGNOSTIC(41000, Warning, unreachableCode, "unreachable code detected")
 DIAGNOSTIC(41001, Error, recursiveType, "type '$0' contains cyclic reference to itself.")
 
+DIAGNOSTIC(
+    41009,
+    Error,
+    missingReturnError,
+    "non-void function must return in all cases for target '$0'")
 DIAGNOSTIC(41010, Warning, missingReturn, "non-void function does not return in all cases")
 DIAGNOSTIC(
     41011,
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 845712c2d..841f44a80 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -79,6 +79,7 @@
 #include "slang-ir-lower-tuple-types.h"
 #include "slang-ir-metadata.h"
 #include "slang-ir-metal-legalize.h"
+#include "slang-ir-missing-return.h"
 #include "slang-ir-optix-entry-point-uniforms.h"
 #include "slang-ir-pytorch-cpp-binding.h"
 #include "slang-ir-redundancy-removal.h"
@@ -326,6 +327,7 @@ struct RequiredLoweringPassSet
     bool dynamicResourceHeap;
     bool resolveVaryingInputRef;
     bool specializeStageSwitch;
+    bool missingReturn;
 };
 
 // Scan the IR module and determine which lowering/legalization passes are needed based
@@ -451,6 +453,9 @@ void calcRequiredLoweringPassSet(
     case kIROp_GetCurrentStage:
         result.specializeStageSwitch = true;
         break;
+    case kIROp_MissingReturn:
+        result.missingReturn = true;
+        break;
     }
     if (!result.generics || !result.existentialTypeLayout)
     {
@@ -1082,6 +1087,9 @@ Result linkAndOptimizeIR(
         checkForRecursiveTypes(irModule, sink);
         checkForRecursiveFunctions(codeGenContext->getTargetReq(), irModule, sink);
 
+        if (requiredLoweringPassSet.missingReturn)
+            checkForMissingReturns(irModule, sink, target, false);
+
         // For some targets, we are more restrictive about what types are allowed
         // to be used as shader parameters in ConstantBuffer/ParameterBlock.
         // We will check for these restrictions here.
diff --git a/source/slang/slang-ir-missing-return.cpp b/source/slang/slang-ir-missing-return.cpp
index 9dcda37ca..82d1171e0 100644
--- a/source/slang/slang-ir-missing-return.cpp
+++ b/source/slang/slang-ir-missing-return.cpp
@@ -1,6 +1,8 @@
 // ir-missing-return.cpp
 #include "slang-ir-missing-return.h"
 
+#include "core/slang-type-text-util.h"
+#include "slang-compiler.h"
 #include "slang-ir-insts.h"
 #include "slang-ir.h"
 
@@ -10,7 +12,45 @@ namespace Slang
 class DiagnosticSink;
 struct IRModule;
 
-void checkForMissingReturnsRec(IRInst* inst, DiagnosticSink* sink)
+// Returns false if compilation target does not allow and errors out(i.e. during downstream
+// compilation) on missing returns.
+static bool doesTargetAllowMissingReturns(CodeGenTarget target)
+{
+    if (isKhronosTarget(target) || isWGPUTarget(target))
+    {
+        return false;
+    }
+
+    return true;
+}
+
+static void diagnoseMissingReturnForTarget(
+    IRMissingReturn* missingReturn,
+    DiagnosticSink* sink,
+    CodeGenTarget target,
+    bool diagnoseWarning)
+{
+    if (doesTargetAllowMissingReturns(target))
+    {
+        if (diagnoseWarning)
+        {
+            sink->diagnose(missingReturn, Diagnostics::missingReturn);
+        }
+    }
+    else
+    {
+        sink->diagnose(
+            missingReturn,
+            Diagnostics::missingReturnError,
+            TypeTextUtil::getCompileTargetName(SlangCompileTarget(target)));
+    }
+}
+
+void checkForMissingReturnsRec(
+    IRInst* inst,
+    DiagnosticSink* sink,
+    CodeGenTarget target,
+    bool diagnoseWarning)
 {
     if (auto code = as<IRGlobalValueWithCode>(inst))
     {
@@ -20,21 +60,25 @@ void checkForMissingReturnsRec(IRInst* inst, DiagnosticSink* sink)
 
             if (auto missingReturn = as<IRMissingReturn>(terminator))
             {
-                sink->diagnose(missingReturn, Diagnostics::missingReturn);
+                diagnoseMissingReturnForTarget(missingReturn, sink, target, diagnoseWarning);
             }
         }
     }
 
     for (auto childInst : inst->getDecorationsAndChildren())
     {
-        checkForMissingReturnsRec(childInst, sink);
+        checkForMissingReturnsRec(childInst, sink, target, diagnoseWarning);
     }
 }
 
-void checkForMissingReturns(IRModule* module, DiagnosticSink* sink)
+void checkForMissingReturns(
+    IRModule* module,
+    DiagnosticSink* sink,
+    CodeGenTarget target,
+    bool diagnoseWarning)
 {
     // Look for any `missingReturn` instructions
-    checkForMissingReturnsRec(module->getModuleInst(), sink);
+    checkForMissingReturnsRec(module->getModuleInst(), sink, target, diagnoseWarning);
 }
 
 } // namespace Slang
diff --git a/source/slang/slang-ir-missing-return.h b/source/slang/slang-ir-missing-return.h
index 0be420883..68c652ebe 100644
--- a/source/slang/slang-ir-missing-return.h
+++ b/source/slang/slang-ir-missing-return.h
@@ -5,6 +5,21 @@ namespace Slang
 {
 class DiagnosticSink;
 struct IRModule;
+enum class CodeGenTarget;
+
+/// This IR check pass is performed twice, once during IR lowering(with no code gen target
+/// specified) and once during linking(with code gen target specified).
+///
+/// Some code gen targets allow missing returns while some do not. The first pass, where no target
+/// is specified, emit warnings on missing returns, and the second pass may additionally emit errors
+/// when compiling for targets that do not support missing returns.
+///
+/// On the second pass, `diagnoseWarning` is set to false to suppress warnings, ensuring that only
+/// errors are emitted. This prevents duplicate warnings from appearing in both passes.
+void checkForMissingReturns(
+    IRModule* module,
+    DiagnosticSink* sink,
+    CodeGenTarget target,
+    bool diagnoseWarning);
 
-void checkForMissingReturns(IRModule* module, DiagnosticSink* sink);
 } // namespace Slang
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index decfe4a91..3395224ec 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -11811,7 +11811,7 @@ RefPtr<IRModule> generateIRForTranslationUnit(
         // TODO: give error messages if any `undefined` or
         // instructions remain.
 
-        checkForMissingReturns(module, compileRequest->getSink());
+        checkForMissingReturns(module, compileRequest->getSink(), CodeGenTarget::None, true);
         // Check for invalid differentiable function body.
         checkAutoDiffUsages(module, compileRequest->getSink());
 
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index e6729ca85..63bcd1ba2 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -2508,9 +2508,9 @@ bool isMetalTarget(TargetRequest* targetReq)
     }
 }
 
-bool isKhronosTarget(TargetRequest* targetReq)
+bool isKhronosTarget(CodeGenTarget target)
 {
-    switch (targetReq->getTarget())
+    switch (target)
     {
     default:
         return false;
@@ -2522,6 +2522,11 @@ bool isKhronosTarget(TargetRequest* targetReq)
     }
 }
 
+bool isKhronosTarget(TargetRequest* targetReq)
+{
+    return isKhronosTarget(targetReq->getTarget());
+}
+
 bool isCPUTarget(TargetRequest* targetReq)
 {
     return ArtifactDescUtil::isCpuLikeTarget(
@@ -2541,9 +2546,9 @@ bool isCUDATarget(TargetRequest* targetReq)
     }
 }
 
-bool isWGPUTarget(TargetRequest* targetReq)
+bool isWGPUTarget(CodeGenTarget target)
 {
-    switch (targetReq->getTarget())
+    switch (target)
     {
     default:
         return false;
@@ -2555,6 +2560,11 @@ bool isWGPUTarget(TargetRequest* targetReq)
     }
 }
 
+bool isWGPUTarget(TargetRequest* targetReq)
+{
+    return isWGPUTarget(targetReq->getTarget());
+}
+
 SourceLanguage getIntermediateSourceLanguageForTarget(TargetProgram* targetProgram)
 {
     // If we are emitting directly, there is no intermediate source language
diff --git a/tests/diagnostics/missing-return-target.slang b/tests/diagnostics/missing-return-target.slang
new file mode 100644
index 000000000..8fc65efd1
--- /dev/null
+++ b/tests/diagnostics/missing-return-target.slang
@@ -0,0 +1,29 @@
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_NOT_SUPP): -entry computeMain -stage compute -target spirv
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_NOT_SUPP): -entry computeMain -stage compute -target glsl
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_NOT_SUPP): -entry computeMain -stage compute -target wgsl
+
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_SUPP):
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_SUPP): -entry computeMain -stage compute -target hlsl
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_SUPP): -entry computeMain -stage compute -target metal
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_SUPP): -entry computeMain -stage compute -target cpp
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK_SUPP): -entry computeMain -stage compute -target cuda
+
+// Some compilation targets allow missing returns while some do not.
+// This test ensures that either errors and warnings are emitted appropriately.
+
+RWStructuredBuffer<uint> out;
+
+// CHECK_NOT_SUPP: warning 41010: non-void function
+// CHECK_NOT_SUPP: error 41009: non-void function
+
+// CHECK_SUPP: warning 41010: non-void function
+
+uint func()
+{
+}
+
+[shader("compute")]
+void computeMain()
+{
+    out[0] = func();
+}
diff --git a/tests/hlsl/tbuffer.slang b/tests/hlsl/tbuffer.slang
index a1bebf2e1..cb014bf2b 100644
--- a/tests/hlsl/tbuffer.slang
+++ b/tests/hlsl/tbuffer.slang
@@ -35,6 +35,6 @@ tbuffer tbuf2 : register(t1)
 
 RWStructuredBuffer<float4> outputBuffer;
 [numthreads(1,1,1)]
-float4 computeMain() {
+void computeMain() {
     outputBuffer[0] = tb_val1 + texture2D[0] + tb_val2;
 }