SPIRV Legalization fixes. (#3479)

* Fix CFG legalization for SPIRV backend.

* Emit DepthReplacing execution mode.

* Fix do-while lowering.

---------

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

8 files changed, 197 insertions, 24 deletions

diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 46ce1d0b5..a2a0bff58 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -2557,6 +2557,99 @@ struct SPIRVEmitContext
         }
     }
 
+
+    SpvExecutionMode isDepthOutput(IRInst* builtinVar)
+    {
+        SpvExecutionMode result = SpvExecutionModeMax;
+        bool isDepthVar = false;
+        if (auto layout = getVarLayout(builtinVar))
+        {
+            if (auto systemValueAttr = layout->findAttr<IRSystemValueSemanticAttr>())
+            {
+                String semanticName = systemValueAttr->getName();
+                semanticName = semanticName.toLower();
+                if (semanticName == "sv_position")
+                {
+                    auto importDecor = builtinVar->findDecoration<IRImportDecoration>();
+                    if (importDecor->getMangledName() == "gl_FragCoord")
+                    {
+                        isDepthVar = true;
+                        result = SpvExecutionModeDepthReplacing;
+                    }
+                }
+                else if (semanticName == "sv_depth")
+                {
+                    isDepthVar = true;
+                    result = SpvExecutionModeDepthReplacing;
+                }
+                else if (semanticName == "sv_depthgreaterequal")
+                {
+                    isDepthVar = true;
+                    result = SpvExecutionModeDepthGreater;
+                }
+                else if (semanticName == "sv_depthlessequal")
+                {
+                    isDepthVar = true;
+                    result = SpvExecutionModeDepthLess;
+                }
+            }
+        }
+        if (!isDepthVar)
+            return result;
+        for (auto use = builtinVar->firstUse; use; use = use->nextUse)
+        {
+            auto user = use->getUser();
+            if (user->getOp() == kIROp_Load)
+                continue;
+            if (as<IRDecoration>(user))
+                continue;
+            switch (user->getOp())
+            {
+            case kIROp_SwizzledStore:
+            case kIROp_Store:
+                return result;
+            }
+        }
+        return result;
+    }
+
+    void maybeEmitEntryPointDepthReplacingExecutionMode(
+        IRFunc* entryPoint,
+        const List<IRInst*>& referencedBuiltinIRVars)
+    {
+        // Check if the entrypoint uses any depth output builtin variables,
+        // if so, we need to emit a DepthReplacing execution mode for the
+        // fragment entrypoint.
+        bool needDepthReplacingMode = false;
+        SpvExecutionMode mode = SpvExecutionModeMax;
+        for (auto globalInst : referencedBuiltinIRVars)
+        {
+            if (auto thisMode = isDepthOutput(globalInst))
+            {
+                needDepthReplacingMode = true;
+                if (mode == SpvExecutionModeMax)
+                    mode = thisMode;
+                else if (mode != thisMode)
+                    mode = SpvExecutionModeDepthReplacing;
+                break;
+            }
+        }
+        if (!needDepthReplacingMode)
+            return;
+        emitOpExecutionMode(getSection(SpvLogicalSectionID::ExecutionModes),
+            nullptr,
+            entryPoint,
+            SpvExecutionModeDepthReplacing);
+        if (mode != SpvExecutionModeDepthReplacing &&
+            mode != SpvExecutionModeMax)
+        {
+            emitOpExecutionMode(getSection(SpvLogicalSectionID::ExecutionModes),
+                nullptr,
+                entryPoint,
+                mode);
+        }
+    }
+
         /// Emit an appropriate SPIR-V decoration for the given IR `decoration`, if necessary and possible.
         ///
         /// The given `dstID` should be the `<id>` of the SPIR-V instruction being decorated,
@@ -2624,6 +2717,7 @@ struct SPIRVEmitContext
                 auto name = entryPointDecor->getName()->getStringSlice();
                 List<SpvInst*> params;
                 HashSet<SpvInst*> paramsSet;
+                List<IRInst*> referencedBuiltinIRVars;
                 // `interface` part: reference all global variables that are used by this entrypoint.
                 for (auto globalInst : m_irModule->getModuleInst()->getChildren())
                 {
@@ -2641,6 +2735,7 @@ struct SPIRVEmitContext
                             {
                                 paramsSet.add(spvGlobalInst);
                                 params.add(spvGlobalInst);
+                                referencedBuiltinIRVars.add(globalInst);
                             }
                         }
                         break;
@@ -2648,8 +2743,11 @@ struct SPIRVEmitContext
                     default:
                         break;
                     }
+                    if (entryPointDecor->getProfile().getStage() == Stage::Fragment)
+                    {
+                        maybeEmitEntryPointDepthReplacingExecutionMode(entryPoint, referencedBuiltinIRVars);
+                    }
                 }
-
                 // Add remaining builtin variables that does not have a corresponding IR global var/param.
                 // These variables could be added from SPIRV ASM blocks.
                 for (auto builtinVar : m_builtinGlobalVars)
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index aa436a796..7ad4ad5d4 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -933,7 +933,7 @@ Result linkAndOptimizeIR(
 
     if (isKhronosTarget(targetRequest) && targetRequest->shouldEmitSPIRVDirectly())
     {
-        performIntrinsicFunctionFunctionInlining(irModule);
+        performIntrinsicFunctionInlining(irModule);
         eliminateDeadCode(irModule);
     }
     eliminateMultiLevelBreak(irModule);
diff --git a/source/slang/slang-ir-inline.cpp b/source/slang/slang-ir-inline.cpp
index eea47533d..a5538425f 100644
--- a/source/slang/slang-ir-inline.cpp
+++ b/source/slang/slang-ir-inline.cpp
@@ -988,7 +988,7 @@ struct IntrinsicFunctionInliningPass : InliningPassBase
     }
 };
 
-void performIntrinsicFunctionFunctionInlining(IRModule* module)
+void performIntrinsicFunctionInlining(IRModule* module)
 {
     IntrinsicFunctionInliningPass pass(module);
     bool changed = true;
diff --git a/source/slang/slang-ir-inline.h b/source/slang/slang-ir-inline.h
index 539bb26c0..a15887b00 100644
--- a/source/slang/slang-ir-inline.h
+++ b/source/slang/slang-ir-inline.h
@@ -32,7 +32,7 @@ namespace Slang
     void performGLSLResourceReturnFunctionInlining(IRModule* module);
 
         /// Inline simple intrinsic functions whose definition is a single asm block.
-    void performIntrinsicFunctionFunctionInlining(IRModule* module);
+    void performIntrinsicFunctionInlining(IRModule* module);
 
         /// Inline a specific call.
     bool inlineCall(IRCall* call);
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index c4b0a471a..81bec8aa7 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -19,6 +19,7 @@
 #include "slang-ir-simplify-cfg.h"
 #include "slang-ir-peephole.h"
 #include "slang-ir-redundancy-removal.h"
+#include "slang-ir-loop-unroll.h"
 
 namespace Slang
 {
@@ -1125,23 +1126,11 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
         //       - the back-edge block must structurally post dominate the
         //         Continue Target
 
-        // If the continue block has only a single predecessor, pretend like it
-        // is just ordinary control flow
-        //
-        // TODO: could this fail in cases like this, where it had a single
-        // predecessor, but it's still nested inside a region?
-        // do{
-        //   if(x)
-        //     continue;
-        //   unreachable
-        // } while(foo)
+        // By this point, we should have already eliminated all continue jumps and
+        // turned them into a break jump. So all loop insts should satisfy
+        // continueBlock == targetBlock.
         const auto t = loop->getTargetBlock();
         auto c = loop->getContinueBlock();
-        if(c->getPredecessors().getCount() <= 1)
-        {
-            c = t;
-            loop->continueBlock.set(c);
-        }
 
         // Our IR allows multiple back-edges to a loop header if this is also
         // the loop continue block. SPIR-V does not so replace them with a
@@ -1565,6 +1554,9 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
             case kIROp_SPIRVAsm:
                 processSPIRVAsm(as<IRSPIRVAsm>(inst));
                 break;
+            case kIROp_Func:
+                eliminateContinueBlocksInFunc(m_module, as<IRFunc>(inst));
+                [[fallthrough]];
             default:
                 for (auto child = inst->getLastChild(); child; child = child->getPrevInst())
                 {
@@ -1826,7 +1818,10 @@ void simplifyIRForSpirvLegalization(TargetRequest* target, DiagnosticSink* sink,
                 funcChanged |= applySparseConditionalConstantPropagation(func, sink);
                 funcChanged |= peepholeOptimize(target, func);
                 funcChanged |= removeRedundancyInFunc(func);
-                funcChanged |= simplifyCFG(func, CFGSimplificationOptions::getFast());
+                CFGSimplificationOptions options;
+                options.removeTrivialSingleIterationLoops = true;
+                options.removeSideEffectFreeLoops = false;
+                funcChanged |= simplifyCFG(func, options);
                 eliminateDeadCode(func);
             }
         }
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index f624b0baa..00db77511 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -5695,13 +5695,32 @@ struct StmtLoweringVisitor : StmtVisitor<StmtLoweringVisitor>
         {
             auto irCondition = getSimpleVal(context,
                 lowerRValueExpr(context, condExpr));
+            auto invCondition = builder->emitNot(irCondition->getDataType(), irCondition);
 
             // Now we want to `break` if the loop condition is false,
             // otherwise we will jump back to the head of the loop.
-            builder->emitLoopTest(
-                irCondition,
-                loopHead,
-                breakLabel);
+            // 
+            // We need to make sure not to reuse the break block of the loop as
+            // the break/merge block of the ifelse test.
+            // Therefore, we introduce a separate merge block for the loop test.
+            // 
+            // Emit the following structure:
+            // 
+            // [merge(mergeBlock)]
+            // if (cond) goto loopHead;
+            // else goto mergeBlock;
+            // 
+            // mergeBlock:
+            //   goto breakLabel;
+            auto mergeBlock = builder->emitBlock();
+            builder->emitBranch(loopHead);
+
+            builder->setInsertInto(testLabel);
+            builder->emitIfElse(
+                invCondition,
+                breakLabel,
+                mergeBlock,
+                mergeBlock);
         }
 
         // Finally we insert the label that a `break` will jump to
diff --git a/tests/spirv/depth-replacing.slang b/tests/spirv/depth-replacing.slang
new file mode 100644
index 000000000..df2ab17e1
--- /dev/null
+++ b/tests/spirv/depth-replacing.slang
@@ -0,0 +1,16 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv -stage fragment -entry main -emit-spirv-directly
+
+struct PSOut
+{
+    float4 color : SV_Target;
+    float depth : SV_Depth;
+}
+// CHECK: OpExecutionMode {{.*}} DepthReplacing
+
+PSOut main()
+{
+    PSOut psout;
+    psout.color = float4(1.0f, 0.0f, 0.0f, 1.0f);
+    psout.depth = 0.5f;
+    return psout;
+}
diff --git a/tests/spirv/while-continue.slang b/tests/spirv/while-continue.slang
new file mode 100644
index 000000000..9c7b89295
--- /dev/null
+++ b/tests/spirv/while-continue.slang
@@ -0,0 +1,45 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv -stage compute -entry main -emit-spirv-directly
+
+// Test that we can transform a continue in a while loop to valid spirv.
+
+uniform int loopCount;
+
+// CHECK: OpEntryPoint
+
+bool condition1()
+{
+    AllMemoryBarrier();
+    return output[2] < 1;
+}
+bool condition2()
+{
+    AllMemoryBarrier();
+    return output[3] < 1;
+}
+
+RWStructuredBuffer<float> output;
+
+[numthreads(1,1,1)]
+void main()
+{
+    float weight = 0.0;
+    do 
+    {
+        if (condition2()) 
+        { 
+            continue;
+        }
+        AllMemoryBarrier();
+    } while (condition1());
+
+    while(condition1())
+    {
+        if (condition2())
+        {
+            continue;
+        }
+        AllMemoryBarrier();
+    }
+
+    output[0] = weight;
+}