Fix scope fixing for address insts. (#2724)

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

7 files changed, 106 insertions, 12 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index a31c16505..356f1c7ce 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -2835,7 +2835,7 @@ void CLikeSourceEmitter::emitFunctionBody(IRGlobalValueWithCode* code)
     // storage for derived structures like the region tree (and logic
     // for invalidating them when a transformation would break them).
     //
-    fixValueScoping(regionTree);
+    fixValueScoping(regionTree, [this](IRInst* inst) {return shouldFoldInstIntoUseSites(inst); });
 
     // Now emit high-level code from that structured region tree.
     //
diff --git a/source/slang/slang-ir-deduplicate.cpp b/source/slang/slang-ir-deduplicate.cpp
index 168451479..f0d0f57dd 100644
--- a/source/slang/slang-ir-deduplicate.cpp
+++ b/source/slang/slang-ir-deduplicate.cpp
@@ -35,10 +35,6 @@ namespace Slang
         }
     }
 
-    void addHoistableInst(
-        IRBuilder* builder,
-        IRInst* inst);
-
     void IRDeduplicationContext::tryHoistInst(IRInst* inst)
     {
         List<IRInst*> workList;
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 88a5f6075..9bb66823b 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -3975,6 +3975,10 @@ inline IRTargetIntrinsicDecoration* findBestTargetIntrinsicDecoration(
 }
 
 
+void addHoistableInst(
+    IRBuilder* builder,
+    IRInst* inst);
+
 }
 
 #endif
diff --git a/source/slang/slang-ir-restructure-scoping.cpp b/source/slang/slang-ir-restructure-scoping.cpp
index 17ed44162..d7195a718 100644
--- a/source/slang/slang-ir-restructure-scoping.cpp
+++ b/source/slang/slang-ir-restructure-scoping.cpp
@@ -211,7 +211,8 @@ void defaultInitializeVar(
 static void fixValueScopingForInst(
     IRInst*         def,
     SimpleRegion*   defRegion,
-    RegionTree*     regionTree)
+    RegionTree*     regionTree,
+    bool            isInstAlwaysFolded)
 {
     // This algorithm should not consider "phi nodes" for now,
     // because the emit logic will already create variables for them.
@@ -305,8 +306,18 @@ static void fixValueScopingForInst(
         // If we've gotten this far, we know that `u` is a "bad"
         // use of `def`, and needs fixing.
         //
-        // We will use a temporary variable to resolve the bad scoping,
-        // creating it on-demand when we ecounter a first "bad" use, and
+        // For insts that are always fold into use sites, we try to hoist them
+        // to as early as possible, and then leave it there.
+        // 
+        if (isInstAlwaysFolded)
+        {
+            def->removeFromParent();
+            addHoistableInst(&builder, def);
+            continue;
+        }
+
+        // For non-hoistable insts, we will use a temporary variable to resolve
+        // the bad scoping, creating it on-demand when we ecounter a first "bad" use, and
         // then re-using that temporary for any subsequent bad uses.
         //
         if( !tmp )
@@ -433,7 +444,7 @@ static void fixValueScopingForInst(
     }
 }
 
-void fixValueScoping(RegionTree* regionTree)
+void fixValueScoping(RegionTree* regionTree, const Func<bool, IRInst*>& shouldAlwaysFoldInst)
 {
     // We are going to have to walk through every instruction
     // in the code of the function to detect an bad cases.
@@ -475,8 +486,8 @@ void fixValueScoping(RegionTree* regionTree)
         for (auto inst = block->getFirstOrdinaryInst(); inst; inst = nextInst)
         {
             nextInst = inst->getNextInst();
-
-            fixValueScopingForInst(inst, parentRegion, regionTree);
+            bool isInstAlwaysFolded = shouldAlwaysFoldInst(inst);
+            fixValueScopingForInst(inst, parentRegion, regionTree, isInstAlwaysFolded);
         }
     }
 }
diff --git a/source/slang/slang-ir-restructure-scoping.h b/source/slang/slang-ir-restructure-scoping.h
index 6c9266754..f629d0013 100644
--- a/source/slang/slang-ir-restructure-scoping.h
+++ b/source/slang/slang-ir-restructure-scoping.h
@@ -1,10 +1,12 @@
 // slang-ir-restructure-scoping.h
 #pragma once
+#include "../core/slang-func-ptr.h"
 
 namespace Slang
 {
 
 class RegionTree;
+struct IRInst;
 
 /// Fix cases where a value might be used in a non-nested region.
 ///
@@ -19,6 +21,6 @@ class RegionTree;
 /// to survive across blocks are communicated through variables
 /// declared at a sufficiently broad scope.
 ///
-void fixValueScoping(RegionTree* regionTree);
+void fixValueScoping(RegionTree* regionTree, const Func<bool, IRInst*>& shouldAlwaysFoldInst);
 
 }
diff --git a/tests/bugs/addr-scope-fix.slang b/tests/bugs/addr-scope-fix.slang
new file mode 100644
index 000000000..8a58b7daf
--- /dev/null
+++ b/tests/bugs/addr-scope-fix.slang
@@ -0,0 +1,76 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+struct HitInfo
+{
+    float val;
+    float getVal() { return val; }
+}
+struct PathInfo
+{
+    HitInfo hit;
+    [mutating]
+    void setVal(float x)
+    {
+        hit.val = x;
+    }
+}
+
+void someOp(HitInfo hitInfo)
+{
+    outputBuffer[0] = hitInfo.val;
+}
+
+void func(int x)
+{
+    PathInfo path = {};
+    int i = 0;
+    if (x < 10)
+    {
+        // The `elementAddr(path, hit)` is first defined in this block.
+        // Note that this block dominates all remaining blocks in this
+        // function, so all future references to path.hit will use
+        // `elementAddr` inst defined here.
+        //
+        // The bug here is that when we emit code, the emit logic will
+        // find that the elementAddr inst is not in a valid region for
+        // future use and will try to create a local var for it. This will
+        // result a pointer typed var being created, and leads to invalid
+        // hlsl/glsl code.
+        // 
+        // The fix is to hoist all always-fold insts to earliest point
+        // in the program instead of creating a var for it.
+        //
+        someOp(path.hit);
+    }
+    else
+    {
+        // This return makes the true block dominate the rest of the blocks
+        // from a region that does not "dominate" the rest code regions.
+        return;
+    }
+
+    while (i < 3)
+    {
+        if (i < 2)
+        {
+            path.setVal(1);
+        }
+        else
+        {
+            path.setVal(2);
+        }
+
+        // This is the point where we are using path.hit again.
+        outputBuffer[i] = path.hit.getVal();
+        i++;
+    }
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    func(4);
+}
\ No newline at end of file
diff --git a/tests/bugs/addr-scope-fix.slang.expected.txt b/tests/bugs/addr-scope-fix.slang.expected.txt
new file mode 100644
index 000000000..73616c399
--- /dev/null
+++ b/tests/bugs/addr-scope-fix.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+1.0
+1.0
+2.0
+