Fix invalid code generation for when using nested resource specialization (#4751)

4 files changed, 174 insertions, 50 deletions

diff --git a/source/slang/slang-ir-specialize-resources.cpp b/source/slang/slang-ir-specialize-resources.cpp
index c4be278df..5e15098d4 100644
--- a/source/slang/slang-ir-specialize-resources.cpp
+++ b/source/slang/slang-ir-specialize-resources.cpp
@@ -102,6 +102,20 @@ bool specializeResourceParameters(
     return result;
 }
 
+void inlineAllCallsOfFunction(IRFunc* func)
+{
+    traverseUses(func, [&](IRUse* use)
+        {
+            auto user = use->getUser();
+            auto call = as<IRCall>(user);
+            if (!call)
+                return;
+            if (call->getCallee() != func)
+                return;
+            inlineCall(call);
+        });
+}
+
     /// A pass to specialize resource-typed function outputs
 struct ResourceOutputSpecializationPass
 {
@@ -116,11 +130,27 @@ struct ResourceOutputSpecializationPass
     TargetRequest*  targetRequest;
     IRModule*       module;
 
-    // Functions that requires specialization but are currently unspecializable.
-    List<IRFunc*>* unspecializableFuncs;
+    /// Functions that requires specialization but are currently unspecializable.
+    HashSet<IRFunc*>* unspecializableFuncs;
+
+    /// Functions that required specialization and were specialized.
+    HashSet<IRFunc*> specializedFuncs;
+
+    enum class SpecializeFuncResult
+    {
+        OtherFuncFailed = -2,
+        ThisFuncFailed = -1,
+        Ok = 1,
+    };
+
+    bool failedResult(SpecializeFuncResult val)
+    {
+        return val < SpecializeFuncResult::Ok;
+    }
 
     bool processModule()
     {
+        specializedFuncs.clear();
         bool changed = false;
 
         // The main logic consists of iterating over all functions
@@ -140,6 +170,12 @@ struct ResourceOutputSpecializationPass
 
     bool processFunc(IRFunc* oldFunc)
     {
+        // Avoid re-computing by checking our 'processFunc' cache.
+        if (specializedFuncs.contains(oldFunc))
+            return true;
+        if (unspecializableFuncs->contains(oldFunc))
+            return false;
+
         // We don't want to waste any effort on functions that don't merit
         // specialization, so the first step is to identify if the function
         // has any outputs that use resource types.
@@ -185,7 +221,8 @@ struct ResourceOutputSpecializationPass
         // us to rewrite call sites.
         //
         FuncInfo funcInfo;
-        if( SLANG_FAILED(specializeFunc(newFunc, funcInfo)) )
+        SpecializeFuncResult result = specializeFunc(newFunc, funcInfo);
+        if( failedResult(result) )
         {
             // Even though we deterined that we *should* specialize
             // this function, we were not able to because of some
@@ -209,7 +246,10 @@ struct ResourceOutputSpecializationPass
             // messages can be front-end rather than back-end errors.
             //
             newFunc->removeAndDeallocate();
-            unspecializableFuncs->add(oldFunc);
+            // Check if `oldFunc` is the reason for failing,
+            // Otherwise don't add to 'unspecializableFuncs'
+            if(result == SpecializeFuncResult::ThisFuncFailed)
+                unspecializableFuncs->add(oldFunc);
             return false;
         }
 
@@ -274,6 +314,7 @@ struct ResourceOutputSpecializationPass
         {
             specializeCallSite(oldCall, newFunc, funcInfo);
         }
+        specializedFuncs.add(oldFunc);
         return true;
     }
 
@@ -485,7 +526,7 @@ struct ResourceOutputSpecializationPass
 
     // We now turn to the code that fills in the `FuncInfo` structure.
 
-    Result specializeFunc(IRFunc* func, FuncInfo& outFuncInfo)
+    SpecializeFuncResult specializeFunc(IRFunc* func, FuncInfo& outFuncInfo)
     {
         // To specialize a function, we attempt to specialize
         // all the applicable parameters and the function result.
@@ -516,26 +557,28 @@ struct ResourceOutputSpecializationPass
             nextParam = param->getNextParam();
 
             ParamInfo paramInfo;
-            SLANG_RETURN_ON_FAIL(maybeSpecializeParam(param, paramInfo, outFuncInfo));
+            auto result = maybeSpecializeParam(param, paramInfo, outFuncInfo);
+            if (failedResult(result)) return result;
             outFuncInfo.oldParams.add(paramInfo);
         }
 
-        SLANG_RETURN_ON_FAIL(maybeSpecializeResult(func, outFuncInfo.result, outFuncInfo));
+        auto result = maybeSpecializeResult(func, outFuncInfo.result, outFuncInfo);
+        if (failedResult(result)) return result;
 
-        return SLANG_OK;
+        return SpecializeFuncResult::Ok;
     }
 
     // The logic for specializing a function result (the return value) is
     // simpler than that for parameters, so we will look at it first.
 
-    Result maybeSpecializeResult(IRFunc* func, ReturnValueInfo& outResultInfo, FuncInfo& ioFuncInfo)
+    SpecializeFuncResult maybeSpecializeResult(IRFunc* func, ReturnValueInfo& outResultInfo, FuncInfo& ioFuncInfo)
     {
         // If the result type of the function isn't a resource type,
         // then we don't need to specialize the result, and we
         // can succeed without doing anything.
         //
         if( !isResourceType(func->getResultType()) )
-            return SLANG_OK;
+            return SpecializeFuncResult::Ok;
 
         // Otherwise, we know that we will need to produce specialization
         // information in `outResultInfo` or fail in the attempt.
@@ -584,7 +627,8 @@ struct ResourceOutputSpecializationPass
             // or to match a new `return` value against previous
             // ones, then the specialization process will fail.
             //
-            SLANG_RETURN_ON_FAIL(specializeOutputValue(value, outResultInfo, ioFuncInfo));
+            auto result = specializeOutputValue(value, outResultInfo, ioFuncInfo);
+            if (failedResult(result)) return result;
 
             // We will replace the `return <value>;` operation with
             // a simple `return;`, because the new specialized function
@@ -599,7 +643,7 @@ struct ResourceOutputSpecializationPass
         // `outResultInfo` and return successfully.
         //
         completeOutputValue(outResultInfo, ioFuncInfo);
-        return SLANG_OK;
+        return SpecializeFuncResult::Ok;
     }
 
     void prepareOutputValue(OutputInfo& ioValueInfo, FuncInfo& ioFuncInfo)
@@ -632,7 +676,7 @@ struct ResourceOutputSpecializationPass
         ioValueInfo.newOutputParamCount = ioFuncInfo.newOutputParams.getCount() - ioValueInfo.firstNewOutputParamIndex;
     }
 
-    Result specializeOutputValue(IRInst* value, OutputInfo& ioOutputInfo, FuncInfo& ioFuncInfo)
+    SpecializeFuncResult specializeOutputValue(IRInst* value, OutputInfo& ioOutputInfo, FuncInfo& ioFuncInfo)
     {
         // This function is called or each `value` that might be written
         // to the output identified by `ioOutputInfo`.
@@ -660,7 +704,7 @@ struct ResourceOutputSpecializationPass
         // the same opcode.
         //
         if(value->getOp() != representative->getOp())
-            return SLANG_FAIL;
+            return SpecializeFuncResult::ThisFuncFailed;
 
         // Furthermore, only certain instructions are amenable to
         // specialization, because in general we cannot reproduce
@@ -679,7 +723,7 @@ struct ResourceOutputSpecializationPass
             // Any opcode we do not specifically enable should cause
             // specialization to fail.
             //
-            return SLANG_FAIL;
+            return SpecializeFuncResult::ThisFuncFailed;
 
         case kIROp_GlobalParam:
             // A direct reference to a global shader parameter is
@@ -688,8 +732,8 @@ struct ResourceOutputSpecializationPass
             // We do need to require that all values used for the
             // same output refer to the *same* global parameter.
             //
-            if(value != representative) return SLANG_FAIL;
-            return SLANG_OK;
+            if(value != representative) return SpecializeFuncResult::ThisFuncFailed;
+            return SpecializeFuncResult::Ok;
 
         // TODO: There are a number of additional cases that we should
         // enable here.
@@ -734,19 +778,19 @@ struct ResourceOutputSpecializationPass
     // is more involved than that for the function `return` value, so we
     // put it off until we'd discussed the shared subroutines.
 
-    Result maybeSpecializeParam(IRParam* param, ParamInfo& outParamInfo, FuncInfo& ioFuncInfo)
+    SpecializeFuncResult maybeSpecializeParam(IRParam* param, ParamInfo& outParamInfo, FuncInfo& ioFuncInfo)
     {
-        // We only want to specialize in the cse where the parameter
+        // We only want to specialize in the case where the parameter
         // is an `out` or `inout` (both inherit from `IROutTypeBase`),
         // and the pointed-to type is a resource.
         //
         auto paramType = param->getDataType();
         auto outType = as<IROutTypeBase>(paramType);
         if(!outType)
-            return SLANG_OK;
+            return SpecializeFuncResult::Ok;
         auto valueType = outType->getValueType();
         if(!isResourceType(valueType))
-            return SLANG_OK;
+            return SpecializeFuncResult::Ok;
 
         prepareOutputValue(outParamInfo, ioFuncInfo);
 
@@ -802,7 +846,9 @@ struct ResourceOutputSpecializationPass
             outParamInfo.oldArgMode = ParamInfo::OldArgMode::Ignore;
         }
 
-        // Next, we want to identify all the places in the function
+        // Before we change something (and likely break this 
+        // function if something fails after a change) we want
+        // to identify all the places in the function
         // that `store` to the given output parameter.
         //
         // Note: this logic is subtly depending on the structure
@@ -836,17 +882,44 @@ struct ResourceOutputSpecializationPass
         // TODO: We should decide on an encoding for the behavior of
         // `out`/`inout` parameters that doesn't have as many "gotcha" cases.
         //
+        // We will also now recursively specialize all `IRCall` inside a 'parent function' 
+        // when trying to specialize a 'parent function'. This is to ensure we do not remove
+        // a parameter SSA needs for SSA'ing a localVar into a globalVar (and DCE requires 
+        // to not DCE an important 'IRCall').
+        // 
+        SpecializeFuncResult recursiveSpecializationResult = SpecializeFuncResult::Ok;
         List<IRStore*> stores;
         traverseUses(param, [&](IRUse* use)
             {
                 auto user = use->getUser();
-                auto store = as<IRStore>(user);
-                if (!store)
+                switch (user->getOp())
+                {
+                case kIROp_Store:
+                {
+                    auto store = as<IRStore>(user);
+                    if (store->ptr.get() != param)
+                        return;
+                    stores.add(store);
+                    return;
+                }
+                case kIROp_Call:
+                {
+                    // This call may require an inline if it fails to specialize
+                    IRFunc* func = as<IRFunc>(as<IRCall>(user)->getCallee());
+                    if (!func)
+                        return;
+                    
+                    if(!processFunc(func))
+                    {
+                        recursiveSpecializationResult = SpecializeFuncResult::OtherFuncFailed;
+                    }
                     return;
-                if (store->ptr.get() != param)
+                }
+                default:
                     return;
-                stores.add(store);
+                };
             });
+        if (failedResult(recursiveSpecializationResult)) return recursiveSpecializationResult;
 
         // Having identified the places where a value is stored to
         // the output parameter, we iterate over those values to
@@ -856,7 +929,8 @@ struct ResourceOutputSpecializationPass
         for(auto store : stores)
         {
             auto value = store->val.get();
-            SLANG_RETURN_ON_FAIL(specializeOutputValue(value, outParamInfo, ioFuncInfo));
+            auto result = specializeOutputValue(value, outParamInfo, ioFuncInfo);
+            if (failedResult(result)) return result;
 
             // Given our assumptions about how `store`s to output
             // parameters are used, we can eliminate all these `store`s
@@ -875,7 +949,7 @@ struct ResourceOutputSpecializationPass
         param->removeAndDeallocate();
 
         completeOutputValue(outParamInfo, ioFuncInfo);
-        return SLANG_OK;
+        return SpecializeFuncResult::Ok;
     }
 
     void specializeCallSite(
@@ -1125,7 +1199,7 @@ struct ResourceOutputSpecializationPass
 bool specializeResourceOutputs(
     CodeGenContext*         codeGenContext,
     IRModule*               module,
-    List<IRFunc*>&          unspecializableFuncs)
+    HashSet<IRFunc*>&          unspecializableFuncs)
 {
     auto targetRequest = codeGenContext->getTargetReq();
     if(isD3DTarget(targetRequest) || isKhronosTarget(targetRequest))
@@ -1170,7 +1244,7 @@ bool specializeResourceUsage(
     for (;;)
     {
         bool changed = true;
-        List<IRFunc*> unspecializableFuncs;
+        HashSet<IRFunc*> unspecializableFuncs;
         while (changed)
         {
             changed = false;
@@ -1201,18 +1275,8 @@ bool specializeResourceUsage(
 
         // Inline unspecializable resource output functions and then continue trying.
         for (auto func : unspecializableFuncs)
-        {
-            traverseUses(func, [&](IRUse* use)
-            {
-                auto user = use->getUser();
-                auto call = as<IRCall>(user);
-                if (!call)
-                    return;
-                if (call->getCallee() != func)
-                    return;
-                inlineCall(call);
-            });
-        }
+            inlineAllCallsOfFunction(func);
+
         simplifyIR(codeGenContext->getTargetProgram(), irModule,
             IRSimplificationOptions::getFast(codeGenContext->getTargetProgram()));
     }
diff --git a/tests/language-feature/resource-specialization-nested-specialization.slang b/tests/language-feature/resource-specialization-nested-specialization.slang
new file mode 100644
index 000000000..2b88e7611
--- /dev/null
+++ b/tests/language-feature/resource-specialization-nested-specialization.slang
@@ -0,0 +1,58 @@
+//TEST:SIMPLE(filecheck=CHECK_DXIL): -target dxil -profile sm_6_0 -entry computeMain -stage compute -DMEMBER_FUNCTION_CALL
+//TEST:SIMPLE(filecheck=CHECK_DXIL): -target dxil -profile sm_6_0 -entry computeMain -stage compute
+
+//CHECK_DXIL: computeMain
+
+struct Grid
+{
+    uint bufSize;
+    StructuredBuffer<uint> buf;
+};
+
+struct GridGeo
+{
+    Grid grids[2];
+
+    void getGrid(uint index, out Grid grid)
+    {
+        grid = grids[index];
+    }
+};
+
+struct Scene
+{
+    GridGeo gridGeo;
+
+    void getGrid_BAD(uint index, out Grid grid)
+    {
+        gridGeo.getGrid(index, grid);
+    }
+
+    void getGrid_GOOD(uint index, out Grid grid)
+    {
+        grid = gridGeo.grids[index];
+    }
+};
+
+ParameterBlock<Scene> gScene;
+RWStructuredBuffer<uint> gridBuffers[2];
+RWStructuredBuffer<uint> outputBuffer;
+
+void direct_getGrid_BAD(uint index, out Grid grid)
+{
+    gScene.gridGeo.getGrid(index, grid);
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 threadId: SV_DispatchThreadID)
+{
+   
+    Grid grid;
+
+#ifdef MEMBER_FUNCTION_CALL
+    direct_getGrid_BAD(1, grid);
+#else 
+    gScene.getGrid_BAD(1, grid);
+#endif
+    gridBuffers[0][1] = grid.buf[1];
+}
diff --git a/tests/current-bugs/resource-struct-out.slang b/tests/language-feature/resource-specialization-struct-out.slang
index d47b2ec7c..b525dafb9 100644
--- a/tests/current-bugs/resource-struct-out.slang
+++ b/tests/language-feature/resource-specialization-struct-out.slang
@@ -1,15 +1,16 @@
-//DISABLE_TEST:SIMPLE:-target hlsl -entry computeMain -profile cs_6_2
+//TEST:SIMPLE(filecheck=CHECK_DXIL):-target dxil -entry computeMain -profile cs_6_2
+//CHECK_DXIL: computeMain
 
-// This test demonstrates out parameter with a struct & resource type crashes 
+// This test demonstrates out parameter with a struct & resource type.
     
-RWTexture1D<float> g_t;
+RWTexture1D<int> g_t;
 
 RWStructuredBuffer<int> outputBuffer;
 
 struct Thing
 {
     int a;
-    RWTexture1D<float> t;
+    RWTexture1D<int> t;
 };
 
 void setThing(out Thing t)
diff --git a/tests/current-bugs/resource-struct-return.slang b/tests/language-feature/resource-specialization-struct-return.slang
index 8d0508097..ca01f55d2 100644
--- a/tests/current-bugs/resource-struct-return.slang
+++ b/tests/language-feature/resource-specialization-struct-return.slang
@@ -1,14 +1,15 @@
-//DISABLE_TEST:SIMPLE:-target hlsl -entry computeMain -profile cs_6_2
+//TEST:SIMPLE(filecheck=CHECK_DXIL):-target dxil -entry computeMain -profile cs_6_2
+//CHECK_DXIL: computeMain
 
-// This test demonstrates returning struct with resource causes internal compiler error
+// This test demonstrates returning struct with resource.
 
-RWTexture1D<float> g_t;
+RWTexture1D<int> g_t;
 RWStructuredBuffer<int> outputBuffer;
 
 struct Thing
 {
     int a;
-    RWTexture1D<float> t;
+    RWTexture1D<int> t;
 };
 
 Thing makeThing()