Fix return type address space checking. (#4465)

solves: `func-resource-result-simple.slang`, related to #4291

when getting the address space of a function return we now check the return address of the inst if the return type is a non pointer type.

2 files changed, 26 insertions, 27 deletions

diff --git a/source/slang/slang-ir-specialize-address-space.cpp b/source/slang/slang-ir-specialize-address-space.cpp
index 5a1874e08..b16449f2e 100644
--- a/source/slang/slang-ir-specialize-address-space.cpp
+++ b/source/slang/slang-ir-specialize-address-space.cpp
@@ -18,24 +18,8 @@ namespace Slang
         {
         }
 
-        AddressSpace getLeafInstAddressSpace(IRInst* inst)
+        AddressSpace getAddressSpaceFromVarType(IRInst* type)
         {
-            if (as<IRGroupSharedRate>(inst->getRate()))
-                return AddressSpace::GroupShared;
-            switch (inst->getOp())
-            {
-            case kIROp_RWStructuredBufferGetElementPtr:
-                return AddressSpace::Global;
-            case kIROp_Var:
-                if (as<IRBlock>(inst->getParent()))
-                    return AddressSpace::ThreadLocal;
-                break;
-            default:
-                break;
-            }
-            auto type = unwrapAttributedType(inst->getDataType());
-            if (!type)
-                return AddressSpace::Generic;
             if (as<IRUniformParameterGroupType>(type))
             {
                 return AddressSpace::Uniform;
@@ -61,6 +45,27 @@ namespace Slang
             return AddressSpace::Generic;
         }
 
+        AddressSpace getLeafInstAddressSpace(IRInst* inst)
+        {
+            if (as<IRGroupSharedRate>(inst->getRate()))
+                return AddressSpace::GroupShared;
+            switch (inst->getOp())
+            {
+            case kIROp_RWStructuredBufferGetElementPtr:
+                return AddressSpace::Global;
+            case kIROp_Var:
+                if (as<IRBlock>(inst->getParent()))
+                    return AddressSpace::ThreadLocal;
+                break;
+            default:
+                break;
+            }
+            auto type = unwrapAttributedType(inst->getDataType());
+            if (!type)
+                return AddressSpace::Generic;
+            return getAddressSpaceFromVarType(type);
+        }
+
         AddressSpace getAddrSpace(IRInst* inst)
         {
             auto addrSpace = mapInstToAddrSpace.tryGetValue(inst);
@@ -154,13 +159,7 @@ namespace Slang
         AddressSpace getFuncResultAddrSpace(IRFunc* callee)
         {
             auto funcType = as<IRFuncType>(callee->getDataType());
-            auto ptrResultType = as<IRPtrTypeBase>(funcType->getResultType());
-            if (!ptrResultType)
-                return AddressSpace::Generic;
-            AddressSpace resultAddrSpace = AddressSpace::Generic;
-            if (ptrResultType->hasAddressSpace())
-                resultAddrSpace = (AddressSpace)ptrResultType->getAddressSpace();
-            return resultAddrSpace;
+            return getAddressSpaceFromVarType(funcType->getResultType());
         }
 
         // Return true if the address space of the function return type is changed.
@@ -316,11 +315,11 @@ namespace Slang
                                 if (addrSpace != AddressSpace::Generic)
                                 {
                                     auto funcType = as<IRFuncType>(func->getDataType());
-                                    auto ptrResultType = as<IRPtrTypeBase>(funcType->getResultType());
-                                    SLANG_ASSERT(ptrResultType);
                                     AddressSpace resultAddrSpace = getFuncResultAddrSpace(func);
                                     if (resultAddrSpace != addrSpace)
                                     {
+                                        auto ptrResultType = as<IRPtrTypeBase>(funcType->getResultType());
+                                        SLANG_ASSERT(ptrResultType);
                                         IRBuilder builder(func);
                                         auto newResultType = builder.getPtrType(ptrResultType->getOp(), ptrResultType->getValueType(), addrSpace);
                                         fixUpFuncType(func, newResultType);
diff --git a/tests/optimization/func-resource-result/func-resource-result-simple.slang b/tests/optimization/func-resource-result/func-resource-result-simple.slang
index 3ca345816..80342e569 100644
--- a/tests/optimization/func-resource-result/func-resource-result-simple.slang
+++ b/tests/optimization/func-resource-result/func-resource-result-simple.slang
@@ -3,7 +3,7 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj
 //TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute -shaderobj
-//DISABLE_TEST(compute):COMPARE_COMPUTE:-slang -shaderobj -mtl
+//TEST(compute):COMPARE_COMPUTE:-slang -shaderobj -mtl
 
 // Test that a function that returns a resource type can be
 // compiled for targets that don't natively support resource