Fix spirv generation for using output stream in a function. (#3806)

* Fix spirv generation for using output stream in a function.

* polish.

source/slang/slang-ir-glsl-legalize.cpp

@@ -2495,76 +2495,53 @@ void legalizeEntryPointParameterForGLSL(
             stage,
             pp);
 
-        // TODO: a GS output stream might be passed into other
+        // A GS output stream might be passed into other
         // functions, so that we should really be modifying
         // any function that has one of these in its parameter
-        // list (and in the limit we should be leagalizing any
-        // type that nests these...).
+        // list.
         //
-        // For now we will just try to deal with `Append` calls
-        // directly in this function.
-
-        for( auto bb = func->getFirstBlock(); bb; bb = bb->getNextBlock() )
-        {
-            for( auto ii = bb->getFirstInst(); ii; ii = ii->getNextInst() )
+        HashSet<IRInst*> workListSet;
+        List<IRFunc*> workList;
+        workList.add(func);
+        workListSet.add(func);
+        for (Index i = 0; i < workList.getCount(); i++)
+        {
+            auto f = workList[i];
+            for (auto bb = f->getFirstBlock(); bb; bb = bb->getNextBlock())
             {
-                // Is it a call?
-                if(ii->getOp() != kIROp_Call)
-                    continue;
-
-                // Is it calling the append operation?
-                auto callee = ii->getOperand(0);
-                for(;;)
+                for (auto ii = bb->getFirstInst(); ii; ii = ii->getNextInst())
                 {
-                    // If the instruction is `specialize(X,...)` then
-                    // we want to look at `X`, and if it is `generic { ... return R; }`
-                    // then we want to look at `R`. We handle this
-                    // iteratively here.
-                    //
-                    // TODO: This idiom seems to come up enough that we
-                    // should probably have a dedicated convenience routine
-                    // for this.
-                    //
-                    // Alternatively, we could switch the IR encoding so
-                    // that decorations are added to the generic instead of the
-                    // value it returns.
-                    //
-                    switch(callee->getOp())
-                    {
-                    case kIROp_Specialize:
-                        {
-                            callee = cast<IRSpecialize>(callee)->getOperand(0);
-                            continue;
-                        }
+                    // Is it a call?
+                    if (ii->getOp() != kIROp_Call)
+                        continue;
+
+                    // Is it calling the append operation?
+                    auto callee = getResolvedInstForDecorations(ii->getOperand(0));
+                    if (callee->getOp() != kIROp_Func)
+                        continue;
 
-                    case kIROp_Generic:
+                    if (getBuiltinFuncName(callee) != UnownedStringSlice::fromLiteral("GeometryStreamAppend"))
+                    {
+                        // If we are calling a function that takes a output stream as a parameter,
+                        // we need to add it to the work list to be processed.
+                        for (UInt a = 1; a < ii->getOperandCount(); a++)
                         {
-                            auto genericResult = findGenericReturnVal(cast<IRGeneric>(callee));
-                            if(genericResult)
+                            if (as<IRHLSLStreamOutputType>(ii->getOperand(a)->getDataType()))
                             {
-                                callee = genericResult;
-                                continue;
+                                if (workListSet.add(callee))
+                                    workList.add(as<IRFunc>(callee));
+                                break;
                             }
                         }
-
-                    default:
-                        break;
+                        continue;
                     }
-                    break;
-                }
-                if(callee->getOp() != kIROp_Func)
-                    continue;
 
-                if (getBuiltinFuncName(callee) != UnownedStringSlice::fromLiteral("GeometryStreamAppend"))
-                {
-                    continue;
-                }
+                    // Okay, we have a declaration, and we want to modify it!
 
-                // Okay, we have a declaration, and we want to modify it!
+                    builder->setInsertBefore(ii);
 
-                builder->setInsertBefore(ii);
-
-                assign(builder, globalOutputVal, ScalarizedVal::value(ii->getOperand(2)));
+                    assign(builder, globalOutputVal, ScalarizedVal::value(ii->getOperand(2)));
+                }
             }
         }
 

source/slang/slang-ir-specialize-resources.cpp

@@ -1236,6 +1236,8 @@ bool isIllegalGLSLParameterType(IRType* type)
     }
     if (as<IRMeshOutputType>(type))
         return true;
+    if (as<IRHLSLStreamOutputType>(type))
+        return true;
     return false;
 }
 

tests/spirv/geometry-shader-sub-func.slang

@@ -0,0 +1,44 @@
+// geometry-shader-sub-func.slang
+
+//TEST:SIMPLE(filecheck=SPIRV): -target spirv -emit-spirv-directly -O0
+
+struct CoarseVertex
+{
+    float4 position     : POSITION;
+    float3 color        : COLOR;
+    uint   id           : ID;
+}
+
+struct RasterVertex
+{
+    float4 position : POSITION;
+    float3 color    : COLOR;
+    uint   id       : SV_RenderTargetArrayIndex;
+}
+
+// SPIRV: OpEntryPoint
+// SPIRV: OpStore %outputStream_position
+// SPIRV: OpEmitVertex
+void appendVertex(TriangleStream<RasterVertex> outputStream, RasterVertex vert)
+{
+    outputStream.Append(vert);
+}
+
+[shader("geometry")]
+[maxvertexcount(3)]
+void main(
+    triangle CoarseVertex coarseVertices[3],
+    inout TriangleStream<RasterVertex> outputStream,
+    uint primitiveID : SV_PrimitiveID)
+{
+    [ForceUnroll]
+    for(int ii = 0; ii < 3; ++ii)
+    {
+        CoarseVertex coarseVertex = coarseVertices[ii];
+        RasterVertex rasterVertex;
+        rasterVertex.position   = coarseVertex.position;
+        rasterVertex.color      = coarseVertex.color;
+        rasterVertex.id         = coarseVertex.id + primitiveID;
+        appendVertex(outputStream, rasterVertex);
+    }
+}