Fix scope fixing for address insts. (#2724)

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

5 files changed, 25 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);
 
 }