From 1e1a49ccf595dcc99bd9792a47199ec89d5b4370 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Fri, 11 Jul 2025 16:54:43 -0700
Subject: Fixup address spaces after inlining. (#7731)

* Fixup address spaces after inlining.

* add -O0
---
 source/slang/slang-ir-specialize-address-space.cpp | 65 ++++++++++++++++++++++
 1 file changed, 65 insertions(+)

(limited to 'source/slang/slang-ir-specialize-address-space.cpp')

diff --git a/source/slang/slang-ir-specialize-address-space.cpp b/source/slang/slang-ir-specialize-address-space.cpp
index ae0542734..2bc1de775 100644
--- a/source/slang/slang-ir-specialize-address-space.cpp
+++ b/source/slang/slang-ir-specialize-address-space.cpp
@@ -412,4 +412,69 @@ void specializeAddressSpace(IRModule* module, InitialAddressSpaceAssigner* addrS
     AddressSpaceContext context(module, addrSpaceAssigner);
     context.processModule();
 }
+
+void propagateAddressSpaceFromInsts(List<IRInst*>&& workList)
+{
+    HashSet<IRInst*> visited;
+    auto addUserToWorkList = [&](IRInst* inst)
+    {
+        for (auto use = inst->firstUse; use; use = use->nextUse)
+        {
+            auto user = use->getUser();
+            if (visited.add(user))
+                workList.add(user);
+        }
+    };
+    for (auto item : workList)
+    {
+        visited.add(item);
+    }
+    for (Index i = 0; i < workList.getCount(); i++)
+    {
+        auto inst = workList[i];
+        IRBuilder builder(inst);
+        auto instPtrType = as<IRPtrTypeBase>(inst->getDataType());
+        if (!instPtrType)
+            continue;
+        for (auto use = inst->firstUse; use; use = use->nextUse)
+        {
+            auto user = use->getUser();
+            builder.setInsertBefore(user);
+            switch (user->getOp())
+            {
+            case kIROp_Loop:
+            case kIROp_UnconditionalBranch:
+                {
+                    auto branch = as<IRUnconditionalBranch>(user);
+                    UIndex phiIndex = (UIndex)(use - branch->getArgs());
+                    auto param = getParamAt(branch->getTargetBlock(), phiIndex);
+                    if (!param)
+                        continue;
+                    user = param;
+                    break;
+                }
+            }
+            switch (user->getOp())
+            {
+            case kIROp_FieldAddress:
+            case kIROp_GetElementPtr:
+            case kIROp_GetOffsetPtr:
+            case kIROp_Param:
+                {
+                    auto valueType = tryGetPointedToType(&builder, user->getDataType());
+                    if (!valueType)
+                        continue;
+                    auto newType = builder.getPtrTypeWithAddressSpace(valueType, instPtrType);
+                    if (newType != user->getDataType())
+                    {
+                        user->setFullType(newType);
+                        addUserToWorkList(user);
+                    }
+                    break;
+                }
+            }
+        }
+    }
+}
+
 } // namespace Slang
-- 
cgit v1.2.3