Repository: shader-slang/slang
Pull request #6740: Add flag to hoist instructions
Author: login: jarcherNV
Labels: name: pr: non-breaking
description: PRs without breaking changes
Created: 2025-04-04T03:43:05Z | Merged: 2025-04-11T21:51:48Z

This fixes issue #6654

Only hoist instructions that are optimized by prepareFuncForForwardDiff. Do this by adding a flag tryToHoist to IRSimplificationOptions and set this to true only if called from prepareFuncForForwardDiff, then only hoist if the flag is set. Additionally, do not hoist loops if they only have a single trivial iteration.


Diff:
diff --git a/source/slang/slang-ir-autodiff-fwd.cpp b/source/slang/slang-ir-autodiff-fwd.cpp
index 92c35a618b4..08ca56f7d4e 100644
--- a/source/slang/slang-ir-autodiff-fwd.cpp
+++ b/source/slang/slang-ir-autodiff-fwd.cpp
@@ -1877,6 +1877,7 @@ SlangResult ForwardDiffTranscriber::prepareFuncForForwardDiff(IRFunc* func)
         disableIRValidationAtInsert();
         auto simplifyOptions = IRSimplificationOptions::getDefault(nullptr);
         simplifyOptions.removeRedundancy = true;
+        simplifyOptions.hoistLoopInvariantInsts = true;
         simplifyFunc(autoDiffSharedContext->targetProgram, func, simplifyOptions);
         enableIRValidationAtInsert();
     }
diff --git a/source/slang/slang-ir-defer-buffer-load.cpp b/source/slang/slang-ir-defer-buffer-load.cpp
index fea73e70529..bd3b78f9e17 100644
--- a/source/slang/slang-ir-defer-buffer-load.cpp
+++ b/source/slang/slang-ir-defer-buffer-load.cpp
@@ -167,7 +167,7 @@ struct DeferBufferLoadContext
 
     void deferBufferLoadInFunc(IRFunc* func)
     {
-        removeRedundancyInFunc(func);
+        removeRedundancyInFunc(func, false);
 
         currentFunc = func;
         dominatorTree = func->getModule()->findOrCreateDominatorTree(func);
diff --git a/source/slang/slang-ir-redundancy-removal.cpp b/source/slang/slang-ir-redundancy-removal.cpp
index 0d9b910c4f5..a6dac723e27 100644
--- a/source/slang/slang-ir-redundancy-removal.cpp
+++ b/source/slang/slang-ir-redundancy-removal.cpp
@@ -1,6 +1,7 @@
 #include "slang-ir-redundancy-removal.h"
 
 #include "slang-ir-dominators.h"
+#include "slang-ir-simplify-cfg.h"
 #include "slang-ir-util.h"
 
 namespace Slang
@@ -20,6 +21,10 @@ struct RedundancyRemovalContext
             auto terminatorInst = parentBlock->getTerminator();
             if (auto loop = as<IRLoop>(terminatorInst))
             {
+                // Don't bother hoisting if a loop has only a single trivial iteration.
+                if (isTrivialSingleIterationLoop(dom, func, loop))
+                    continue;
+
                 // If `inst` is outside of the loop region, don't hoist it into the loop.
                 if (dom->dominates(loop->getBreakBlock(), inst))
                     continue;
@@ -62,7 +67,8 @@ struct RedundancyRemovalContext
     bool removeRedundancyInBlock(
         Dictionary<IRBlock*, DeduplicateContext>& mapBlockToDedupContext,
         IRGlobalValueWithCode* func,
-        IRBlock* block)
+        IRBlock* block,
+        bool hoistLoopInvariantInsts)
     {
         bool result = false;
         auto& deduplicateContext = mapBlockToDedupContext.getValue(block);
@@ -89,7 +95,8 @@ struct RedundancyRemovalContext
             {
                 // This inst is unique, we should consider hoisting it
                 // if it is inside a loop.
-                result |= tryHoistInstToOuterMostLoop(func, resultInst);
+                if (hoistLoopInvariantInsts)
+                    result |= tryHoistInstToOuterMostLoop(func, resultInst);
             }
         }
         for (auto child : dom->getImmediatelyDominatedBlocks(block))
@@ -101,25 +108,25 @@ struct RedundancyRemovalContext
     }
 };
 
-bool removeRedundancy(IRModule* module)
+bool removeRedundancy(IRModule* module, bool hoistLoopInvariantInsts)
 {
     bool changed = false;
     for (auto inst : module->getGlobalInsts())
     {
         if (auto genericInst = as<IRGeneric>(inst))
         {
-            removeRedundancyInFunc(genericInst);
+            removeRedundancyInFunc(genericInst, hoistLoopInvariantInsts);
             inst = findGenericReturnVal(genericInst);
         }
         if (auto func = as<IRFunc>(inst))
         {
-            changed |= removeRedundancyInFunc(func);
+            changed |= removeRedundancyInFunc(func, hoistLoopInvariantInsts);
         }
     }
     return changed;
 }
 
-bool removeRedundancyInFunc(IRGlobalValueWithCode* func)
+bool removeRedundancyInFunc(IRGlobalValueWithCode* func, bool hoistLoopInvariantInsts)
 {
     auto root = func->getFirstBlock();
     if (!root)
@@ -139,7 +146,11 @@ bool removeRedundancyInFunc(IRGlobalValueWithCode* func)
     {
         for (auto block : workList)
         {
-            result |= context.removeRedundancyInBlock(mapBlockToDeduplicateContext, func, block);
+            result |= context.removeRedundancyInBlock(
+                mapBlockToDeduplicateContext,
+                func,
+                block,
+                hoistLoopInvariantInsts);
 
             for (auto child : context.dom->getImmediatelyDominatedBlocks(block))
             {
diff --git a/source/slang/slang-ir-redundancy-removal.h b/source/slang/slang-ir-redundancy-removal.h
index 40f3b07a946..9d9fb43f012 100644
--- a/source/slang/slang-ir-redundancy-removal.h
+++ b/source/slang/slang-ir-redundancy-removal.h
@@ -7,8 +7,8 @@ namespace Slang
 struct IRModule;
 struct IRGlobalValueWithCode;
 
-bool removeRedundancy(IRModule* module);
-bool removeRedundancyInFunc(IRGlobalValueWithCode* func);
+bool removeRedundancy(IRModule* module, bool hoistLoopInvariantInsts);
+bool removeRedundancyInFunc(IRGlobalValueWithCode* func, bool hoistLoopInvariantInsts);
 
 bool eliminateRedundantLoadStore(IRGlobalValueWithCode* func);
 
diff --git a/source/slang/slang-ir-simplify-cfg.cpp b/source/slang/slang-ir-simplify-cfg.cpp
index 68d79617a83..280874f74d1 100644
--- a/source/slang/slang-ir-simplify-cfg.cpp
+++ b/source/slang/slang-ir-simplify-cfg.cpp
@@ -66,8 +66,8 @@ static IRInst* findBreakableRegionHeaderInst(IRDominatorTree* domTree, IRBlock*
 // Test if a loop is trivial: a trivial loop runs for a single iteration without any back edges, and
 // there is only one break out of the loop at the very end. The function generates `regionTree` if
 // it is needed and hasn't been generated yet.
-static bool isTrivialSingleIterationLoop(
-    CFGSimplificationContext& context,
+bool isTrivialSingleIterationLoop(
+    IRDominatorTree* domTree,
     IRGlobalValueWithCode* func,
     IRLoop* loop)
 {
@@ -91,21 +91,21 @@ static bool isTrivialSingleIterationLoop(
     //
     // We need to verify this is a trivial loop by checking if there is any multi-level breaks
     // that skips out of this loop.
-    if (!context.domTree)
-        context.domTree = computeDominatorTree(func);
+    if (!domTree)
+        domTree = computeDominatorTree(func);
     bool hasMultiLevelBreaks = false;
-    auto loopBlocks = collectBlocksInRegion(context.domTree, loop, &hasMultiLevelBreaks);
+    auto loopBlocks = collectBlocksInRegion(domTree, loop, &hasMultiLevelBreaks);
     if (hasMultiLevelBreaks)
         return false;
     for (auto block : loopBlocks)
     {
         for (auto branchTarget : block->getSuccessors())
         {
-            if (!context.domTree->dominates(loop->getParent(), branchTarget))
+            if (!domTree->dominates(loop->getParent(), branchTarget))
                 return false;
             if (branchTarget != loop->getBreakBlock())
                 continue;
-            if (findBreakableRegionHeaderInst(context.domTree, block) != loop)
+            if (findBreakableRegionHeaderInst(domTree, block) != loop)
             {
                 // If the break is initiated from a nested region, this is not trivial.
                 return false;
@@ -127,7 +127,7 @@ static bool isTrivialSingleIterationLoop(
         auto breakOriginBlock = *loop->getBreakBlock()->getPredecessors().begin();
 
         for (auto currBlock = breakOriginBlock; currBlock;
-             currBlock = context.domTree->getImmediateDominator(currBlock))
+             currBlock = domTree->getImmediateDominator(currBlock))
         {
             auto terminator = currBlock->getTerminator();
             if (terminator == loop)
@@ -139,11 +139,11 @@ static bool isTrivialSingleIterationLoop(
             switch (terminator->getOp())
             {
             case kIROp_loop:
-                if (isBlockInRegion(context.domTree, as<IRLoop>(terminator), breakOriginBlock))
+                if (isBlockInRegion(domTree, as<IRLoop>(terminator), breakOriginBlock))
                     return false;
                 break;
             case kIROp_Switch:
-                if (isBlockInRegion(context.domTree, as<IRSwitch>(terminator), breakOriginBlock))
+                if (isBlockInRegion(domTree, as<IRSwitch>(terminator), breakOriginBlock))
                     return false;
                 break;
             default:
@@ -853,8 +853,10 @@ static bool processFunc(IRGlobalValueWithCode* func, CFGSimplificationOptions op
                     // break at the end of the loop, we can remove the header and turn it into
                     // a normal branch.
                     auto targetBlock = loop->getTargetBlock();
+                    if (!simplificationContext.domTree)
+                        simplificationContext.domTree = computeDominatorTree(func);
                     if (options.removeTrivialSingleIterationLoops &&
-                        isTrivialSingleIterationLoop(simplificationContext, func, loop))
+                        isTrivialSingleIterationLoop(simplificationContext.domTree, func, loop))
                     {
                         builder.setInsertBefore(loop);
                         List<IRInst*> args;
diff --git a/source/slang/slang-ir-simplify-cfg.h b/source/slang/slang-ir-simplify-cfg.h
index 8adde3a4240..3822b0fa830 100644
--- a/source/slang/slang-ir-simplify-cfg.h
+++ b/source/slang/slang-ir-simplify-cfg.h
@@ -5,6 +5,8 @@ namespace Slang
 {
 struct IRModule;
 struct IRGlobalValueWithCode;
+struct IRLoop;
+struct IRDominatorTree;
 
 struct CFGSimplificationOptions
 {
@@ -14,6 +16,11 @@ struct CFGSimplificationOptions
     static CFGSimplificationOptions getFast() { return CFGSimplificationOptions{false, false}; }
 };
 
+bool isTrivialSingleIterationLoop(
+    IRDominatorTree* domTree,
+    IRGlobalValueWithCode* func,
+    IRLoop* loop);
+
 /// Simplifies control flow graph by merging basic blocks that
 /// forms a simple linear chain.
 /// Returns true if changed.
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 2e2ba358fa9..82b424daafd 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -2183,7 +2183,7 @@ void simplifyIRForSpirvLegalization(TargetProgram* target, DiagnosticSink* sink,
                 funcChanged = false;
                 funcChanged |= applySparseConditionalConstantPropagation(func, sink);
                 funcChanged |= peepholeOptimize(target, func);
-                funcChanged |= removeRedundancyInFunc(func);
+                funcChanged |= removeRedundancyInFunc(func, false);
                 CFGSimplificationOptions options;
                 options.removeTrivialSingleIterationLoops = true;
                 options.removeSideEffectFreeLoops = false;
diff --git a/source/slang/slang-ir-ssa-simplification.cpp b/source/slang/slang-ir-ssa-simplification.cpp
index b5081dab7a5..41da868b658 100644
--- a/source/slang/slang-ir-ssa-simplification.cpp
+++ b/source/slang/slang-ir-ssa-simplification.cpp
@@ -90,7 +90,7 @@ void simplifyIR(
                 funcChanged |= applySparseConditionalConstantPropagation(func, sink);
                 funcChanged |= peepholeOptimize(target, func);
                 if (options.removeRedundancy)
-                    funcChanged |= removeRedundancyInFunc(func);
+                    funcChanged |= removeRedundancyInFunc(func, options.hoistLoopInvariantInsts);
                 funcChanged |= simplifyCFG(func, options.cfgOptions);
                 // Note: we disregard the `changed` state from dead code elimination pass since
                 // SCCP pass could be generating temporarily evaluated constant values and never
@@ -122,7 +122,7 @@ void simplifyNonSSAIR(TargetProgram* target, IRModule* module, IRSimplificationO
         changed |= peepholeOptimize(target, module, options.peepholeOptions);
 
         if (!options.minimalOptimization)
-            changed |= removeRedundancy(module);
+            changed |= removeRedundancy(module, options.hoistLoopInvariantInsts);
         changed |= simplifyCFG(module, options.cfgOptions);
 
         // Note: we disregard the `changed` state from dead code elimination pass since
@@ -153,7 +153,7 @@ void simplifyFunc(
         changed |= applySparseConditionalConstantPropagation(func, sink);
         changed |= peepholeOptimize(target, func);
         if (!options.minimalOptimization)
-            changed |= removeRedundancyInFunc(func);
+            changed |= removeRedundancyInFunc(func, options.hoistLoopInvariantInsts);
         changed |= simplifyCFG(func, options.cfgOptions);
 
         // Note: we disregard the `changed` state from dead code elimination pass since
diff --git a/source/slang/slang-ir-ssa-simplification.h b/source/slang/slang-ir-ssa-simplification.h
index fbff2fddadb..e3db559d737 100644
--- a/source/slang/slang-ir-ssa-simplification.h
+++ b/source/slang/slang-ir-ssa-simplification.h
@@ -20,6 +20,7 @@ struct IRSimplificationOptions
 
     bool minimalOptimization = false;
     bool removeRedundancy = false;
+    bool hoistLoopInvariantInsts = false;
 
     static IRSimplificationOptions getDefault(TargetProgram* targetProgram);
 
diff --git a/tests/spirv/forceinline-nohoist.slang b/tests/spirv/forceinline-nohoist.slang
new file mode 100644
index 00000000000..54db2838f40
--- /dev/null
+++ b/tests/spirv/forceinline-nohoist.slang
@@ -0,0 +1,32 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv -stage compute -entry main
+
+// Verify that the call to dot is after the conditional branch.
+
+// CHECK: OpBranchConditional
+// CHECK: OpDot
+
+[ForceInline]
+float test(bool x, float3 a, float3 b) {
+    float result = 0;
+    if(x) {
+        result = dot(a, b);
+    }
+    return result;
+}
+
+float caller(uniform bool x, uniform float3 a, uniform float3 b) {
+    return test(x, a, b);
+}
+
+RWStructuredBuffer<float> output;
+
+uniform bool branchCheck;
+uniform float3 uniformA;
+uniform float3 uniformB;
+
+[numthreads(1,1,1)]
+[shader("compute")]
+void main()
+{
+    output[0] = caller(branchCheck, uniformA, uniformB);
+}


Reviews:
user: login: csyonghe
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T21:33:02Z
user: login: jarcherNV
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T02:39:25Z
user: login: jarcherNV
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T02:39:32Z
user: login: jarcherNV
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-08T02:41:07Z
user: login: csyonghe
body: Looks good!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-08T20:02:24Z

Review comments:
diff_hunk: @@ -14,6 +25,11 @@ struct CFGSimplificationOptions
     static CFGSimplificationOptions getFast() { return CFGSimplificationOptions{false, false}; }
 };
 
+bool isTrivialSingleIterationLoop(
path: source/slang/slang-ir-simplify-cfg.h
user: login: csyonghe
body: isTrivialSingleIterationLoop is only using `CFGSimplificationContext::domTree`, instead of exposing the entire `CFGSimplicationContext` here, just modify isTrivialSingleIterationLoop to take only `domTree`, which is available at our new callsite.
created_at: 2025-04-07T21:31:53Z
author_association: COLLABORATOR
diff_hunk: @@ -20,6 +20,7 @@ struct IRSimplificationOptions
 
     bool minimalOptimization = false;
     bool removeRedundancy = false;
+    bool tryToHoist = false;
path: source/slang/slang-ir-ssa-simplification.h
user: login: csyonghe
body: tryToHoist ==> `hoistLoopInvariantInsts`.
created_at: 2025-04-07T21:32:09Z
author_association: COLLABORATOR
diff_hunk: @@ -20,6 +21,11 @@ struct RedundancyRemovalContext
             auto terminatorInst = parentBlock->getTerminator();
             if (auto loop = as<IRLoop>(terminatorInst))
             {
+                // Don't bother hoisting if a loop has only a single trivial iteration.
+                CFGSimplificationContext ctx;
+                if (isTrivialSingleIterationLoop(ctx, func, loop))
path: source/slang/slang-ir-redundancy-removal.cpp
user: login: csyonghe
body: pass `dom` to `isTrivialSingleIterationLoop` to avoid redundant computation of the dominator tree.
created_at: 2025-04-07T21:32:59Z
author_association: COLLABORATOR
diff_hunk: @@ -14,6 +25,11 @@ struct CFGSimplificationOptions
     static CFGSimplificationOptions getFast() { return CFGSimplificationOptions{false, false}; }
 };
 
+bool isTrivialSingleIterationLoop(
path: source/slang/slang-ir-simplify-cfg.h
user: login: jarcherNV
body: I missed this simplification. Good idea. Fixed.
created_at: 2025-04-08T02:39:25Z
author_association: COLLABORATOR
diff_hunk: @@ -20,6 +20,7 @@ struct IRSimplificationOptions
 
     bool minimalOptimization = false;
     bool removeRedundancy = false;
+    bool tryToHoist = false;
path: source/slang/slang-ir-ssa-simplification.h
user: login: jarcherNV
body: Done.
created_at: 2025-04-08T02:39:32Z
author_association: COLLABORATOR
diff_hunk: @@ -20,6 +21,11 @@ struct RedundancyRemovalContext
             auto terminatorInst = parentBlock->getTerminator();
             if (auto loop = as<IRLoop>(terminatorInst))
             {
+                // Don't bother hoisting if a loop has only a single trivial iteration.
+                CFGSimplificationContext ctx;
+                if (isTrivialSingleIterationLoop(ctx, func, loop))
path: source/slang/slang-ir-redundancy-removal.cpp
user: login: jarcherNV
body: Done.
created_at: 2025-04-08T02:41:06Z
author_association: COLLABORATOR

PR comments:
user: login: jkwak-work
created_at: 2025-04-04T08:56:40Z
body: You need a test case to prevent a regression.
author_association: COLLABORATOR
user: login: jarcherNV
created_at: 2025-04-04T22:23:40Z
body: I have added a test for this to confirm that the dot is after the branch, to verify that it is not hoisted.
author_association: COLLABORATOR
user: login: jarcherNV
created_at: 2025-04-04T23:12:51Z
body: The test seems to be failing with some builds, I'm guessing due to different SPIR-V asm. I'll need to investigate.
author_association: COLLABORATOR
user: login: jarcherNV
created_at: 2025-04-07T03:23:27Z
body: /format
author_association: COLLABORATOR
user: login: slangbot
created_at: 2025-04-07T03:24:14Z
body: 🌈 Formatted, please merge the changes from [this PR](https://github.com/jarcherNV/slang/pull/2)
author_association: CONTRIBUTOR