WIP Mesh shaders for SPIR-V (#3226)

* SPIR-V impl for SetMeshOutputCounts and DispatchMesh

* Unsightly fix for legalization ordering differences between GLSL and SPIR-V

* spelling

* Start a new block after terminating one in the OpEmitMeshTasksExt SPIR-V asm block

* Emit mesh shader decorations in SPIR-V

* Mesh and task shader stages for spir-v

* Output explicit gl builtins for spir-v

* Be more hygenic when SOAizing mesh outputs

* Do not create builtin paramter block for spirv mesh outputs

* Pass mesh payloads around by ref

* comment

* less expected failure

* remove unused

* Add spirv op

* Correct type query for default flat modifier

---------

Co-authored-by: Yong He <yonghe@outlook.com>

10 files changed, 259 insertions, 122 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index da729f67a..f26418855 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -7604,10 +7604,25 @@ float dot2add(float2 left, float2 right, float acc);
 //
 
 // Set the number of output vertices and primitives for a mesh shader invocation.
-__target_intrinsic(glsl, "SetMeshOutputsEXT")
 __glsl_extension(GL_EXT_mesh_shader)
 __glsl_version(450)
-void SetMeshOutputCounts(uint vertexCount, uint primitiveCount);
+void SetMeshOutputCounts(uint vertexCount, uint primitiveCount)
+{
+    __target_switch
+    {
+    case hlsl:
+        __intrinsic_asm "SetMeshOutputCounts";
+    case glsl:
+        __intrinsic_asm "SetMeshOutputsEXT";
+    case spirv:
+        return spirv_asm
+        {
+            OpCapability MeshShadingEXT;
+            OpExtension "SPV_EXT_mesh_shader";
+            OpSetMeshOutputsEXT $vertexCount $primitiveCount;
+        };
+    }
+}
 
 // Specify the number of downstream mesh shader thread groups to invoke from an amplification shader,
 // and provide the values for per-mesh payload parameters.
@@ -7625,6 +7640,16 @@ void DispatchMesh<P>(uint threadGroupCountX, uint threadGroupCountY, uint thread
         // This intrinsic doesn't take into account writing meshPayload. That
         // is dealt with separately by 'legalizeDispatchMeshPayloadForGLSL'.
         __intrinsic_asm "EmitMeshTasksEXT($0, $1, $2)";
+    case spirv:
+        return spirv_asm
+        {
+            OpCapability MeshShadingEXT;
+            OpExtension "SPV_EXT_mesh_shader";
+            OpEmitMeshTasksEXT $threadGroupCountX $threadGroupCountY $threadGroupCountZ &meshPayload;
+            // OpEmitMeshTasksExt is a terminator, so we need to start a new
+            // block to hold whatever comes after this intrinsic
+            %_ = OpLabel
+        };
     }
 }
 
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index f82a76d4a..ebcad1873 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -827,7 +827,7 @@ void GLSLSourceEmitter::_maybeEmitGLSLBuiltin(IRGlobalParam* var, UnownedStringS
         // that inline.
 
         auto paramGroupType = as<IRGLSLOutputParameterGroupType>(var->getFullType());
-        SLANG_ASSERT(paramGroupType && "Mesh shader builtin output was not a paramter group");
+        SLANG_ASSERT(paramGroupType && "Mesh shader builtin output was not a parameter group");
         auto arrayType = as<IRArrayTypeBase>(paramGroupType->getOperand(0));
         SLANG_ASSERT(paramGroupType && "Mesh shader builtin output was not an array");
         auto elementType = as<IRStructType>(arrayType->getElementType());
diff --git a/source/slang/slang-emit-spirv-ops.h b/source/slang/slang-emit-spirv-ops.h
index bc53b70ae..9bd645e5d 100644
--- a/source/slang/slang-emit-spirv-ops.h
+++ b/source/slang/slang-emit-spirv-ops.h
@@ -140,6 +140,21 @@ SpvInst* emitOpExecutionModeOutputVertices(
 
 // https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpExecutionMode
 template<typename T>
+SpvInst* emitOpExecutionModeOutputPrimitivesEXT(
+    SpvInstParent* parent,
+    IRInst* inst,
+    const T& entryPoint,
+    const SpvLiteralInteger& primitiveCount
+)
+{
+    static_assert(isSingular<T>);
+    return emitInst(
+        parent, inst, SpvOpExecutionMode, entryPoint, SpvExecutionModeOutputPrimitivesEXT, primitiveCount
+    );
+}
+
+// https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpExecutionMode
+template<typename T>
 SpvInst* emitOpExecutionModeInvocations(
     SpvInstParent* parent,
     IRInst* inst,
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 14fe2e17b..6a223ae4f 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -1769,7 +1769,6 @@ struct SPIRVEmitContext
     {
         bool needDefaultSetBindingDecoration = false;
         bool hasExplicitSetBinding = false;
-        bool isInput = false;
         for (auto rr : layout->getOffsetAttrs())
         {
             UInt index = rr->getOffset();
@@ -1786,7 +1785,6 @@ struct SPIRVEmitContext
                     varInst,
                     SpvLiteralInteger::from32(int32_t(index))
                 );
-                isInput = true;
                 break;
             case LayoutResourceKind::VaryingOutput:
                 emitOpDecorateLocation(
@@ -1872,7 +1870,8 @@ struct SPIRVEmitContext
                 && layout->getStage() == Stage::Fragment
                 && layout->usesResourceKind(LayoutResourceKind::VaryingInput))
             {
-                if (isIntegralScalarOrCompositeType(var->getDataType()))
+                const auto ptrType = as<IRPtrTypeBase>(var->getDataType());
+                if (ptrType && isIntegralScalarOrCompositeType(ptrType->getValueType()))
                     emitOpDecorate(getSection(SpvLogicalSectionID::Annotations), nullptr, varInst, SpvDecorationFlat);
             }
         }
@@ -2683,6 +2682,44 @@ struct SPIRVEmitContext
                 );
             }
             break;
+
+        case kIROp_OutputTopologyDecoration:
+            {
+                const auto o = cast<IROutputTopologyDecoration>(decoration);
+                const auto t = o->getTopology()->getStringSlice();
+                const auto m =
+                      t == "triangle" ? SpvExecutionModeOutputTrianglesEXT
+                    : t == "line" ? SpvExecutionModeOutputLinesEXT
+                    : t == "point" ? SpvExecutionModeOutputPoints
+                    : SpvExecutionModeMax;
+                SLANG_ASSERT(m != SpvExecutionModeMax);
+                emitOpExecutionMode(getSection(SpvLogicalSectionID::ExecutionModes), decoration, dstID, m);
+            }
+            break;
+
+        case kIROp_VerticesDecoration:
+            {
+                const auto c = cast<IRVerticesDecoration>(decoration);
+                emitOpExecutionModeOutputVertices(
+                    getSection(SpvLogicalSectionID::ExecutionModes),
+                    decoration,
+                    dstID,
+                    SpvLiteralInteger::from32(int32_t(c->getMaxSize()->getValue()))
+                );
+            }
+            break;
+
+        case kIROp_PrimitivesDecoration:
+            {
+                const auto c = cast<IRPrimitivesDecoration>(decoration);
+                emitOpExecutionModeOutputPrimitivesEXT(
+                    getSection(SpvLogicalSectionID::ExecutionModes),
+                    decoration,
+                    dstID,
+                    SpvLiteralInteger::from32(int32_t(c->getMaxSize()->getValue()))
+                );
+            }
+            break;
         // ...
         }
     }
@@ -2822,6 +2859,8 @@ struct SPIRVEmitContext
         CASE(Geometry,  Geometry);
         CASE(Fragment,  Fragment);
         CASE(Compute,   GLCompute);
+        CASE(Mesh,      MeshEXT);
+        CASE(Amplification, TaskEXT);
 
         // TODO: Extended execution models for ray tracing, etc.
 
@@ -3038,6 +3077,23 @@ struct SPIRVEmitContext
                 SLANG_UNREACHABLE("Unimplemented system value in spirv emit.");
             }
         }
+
+        //
+        // These are system-value variables which require redeclaration in
+        // GLSL, SPIR-V makes no such distinction so we can use similar logic
+        // to above.
+        //
+        if(const auto linkageDecoration = inst->findDecoration<IRLinkageDecoration>())
+        {
+            const auto name = linkageDecoration->getMangledName();
+            if(name == "gl_PrimitiveTriangleIndicesEXT")
+                return getBuiltinGlobalVar(inst->getFullType(), SpvBuiltInPrimitiveTriangleIndicesEXT);
+            if(name == "gl_PrimitiveLineIndicesEXT")
+                return getBuiltinGlobalVar(inst->getFullType(), SpvBuiltInPrimitiveLineIndicesEXT);
+            if(name == "gl_PrimitivePointIndicesEXT")
+                return getBuiltinGlobalVar(inst->getFullType(), SpvBuiltInPrimitivePointIndicesEXT);
+        }
+
         return nullptr;
     }
 
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 15ce25c3c..6a49e9842 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -896,10 +896,6 @@ Result linkAndOptimizeIR(
     //
     // If any have survived this far, change them back to regular (decorated)
     // arrays that the emitters can deal with.
-    //
-    // TODO: This is too early for the SPIR-V backend, which requires these
-    // types for when it calls legalizeEntryPointsForGLSL (later than GLSL does
-    // above)
     legalizeMeshOutputTypes(irModule);
 
     if (options.shouldLegalizeExistentialAndResourceTypes)
diff --git a/source/slang/slang-ir-glsl-legalize.cpp b/source/slang/slang-ir-glsl-legalize.cpp
index eb808bc97..2d24a0577 100644
--- a/source/slang/slang-ir-glsl-legalize.cpp
+++ b/source/slang/slang-ir-glsl-legalize.cpp
@@ -1941,7 +1941,8 @@ static void legalizeMeshPayloadInputParam(
     IRBuilderInsertLocScope locScope{builder};
     builder->setInsertInto(builder->getModule());
 
-    const auto g = builder->emitVar(pp->getDataType(), SpvStorageClassTaskPayloadWorkgroupEXT);
+    const auto ptrType = cast<IRPtrTypeBase>(pp->getDataType());
+    const auto g = builder->createGlobalVar(ptrType->getValueType(), SpvStorageClassTaskPayloadWorkgroupEXT);
     g->setFullType(builder->getRateQualifiedType(builder->getGroupSharedRate(), g->getFullType()));
     // moveValueBefore(g, builder->getFunc());
     builder->addNameHintDecoration(g, pp->findDecoration<IRNameHintDecoration>()->getName());
@@ -2023,6 +2024,7 @@ static void legalizeMeshOutputParam(
     moveValueBefore(g, builder->getFunc());
     builder->addNameHintDecoration(g, pp->findDecoration<IRNameHintDecoration>()->getName());
     pp->replaceUsesWith(g);
+    // pp is only removed later on, so sadly we have to keep it around for now
     struct MeshOutputSpecializationCondition : FunctionCallSpecializeCondition
     {
         bool doesParamWantSpecialization(IRParam*, IRInst* arg)
@@ -2059,6 +2061,7 @@ static void legalizeMeshOutputParam(
                 IRBuilderInsertLocScope locScope{builder};
                 builder->setInsertBefore(a);
                 a->replaceUsesWith(d.irValue);
+                a->removeAndDeallocate();
                 return;
             }
             // Otherwise, go through the uses one by one and see what we can do
@@ -2134,7 +2137,11 @@ static void legalizeMeshOutputParam(
                     SLANG_UNEXPECTED("Unhandled use of mesh output parameter during GLSL legalization");
                 }
             });
+
+            SLANG_ASSERT(!a->hasUses());
+            a->removeAndDeallocate();
         };
+
         assignUses(globalOutputVal, l);
     });
 
@@ -2173,119 +2180,127 @@ static void legalizeMeshOutputParam(
     //
 
     // First, collect the subset of outputs being used
-    auto isMeshOutputBuiltin = [](IRInst* g)
-    {
-        if(const auto s = composeGetters<IRStringLit>(
-            g,
-            &IRInst::findDecoration<IRImportDecoration>,
-            &IRImportDecoration::getMangledNameOperand
-            ))
-        {
-            const auto n = s->getStringSlice();
-            if (n == "gl_Position" ||
-                n == "gl_PointSize" ||
-                n == "gl_ClipDistance" ||
-                n == "gl_CullDistance" ||
-                n == "gl_PrimitiveID" ||
-                n == "gl_Layer" ||
-                n == "gl_ViewportIndex" ||
-                n == "gl_CullPrimitiveEXT" ||
-                n == "gl_PrimitiveShadingRateEXT")
+    const bool isSPIRV = codeGenContext->getTargetFormat() == CodeGenTarget::SPIRV
+        || codeGenContext->getTargetFormat() == CodeGenTarget::SPIRVAssembly;
+    if(!isSPIRV)
+    {
+        auto isMeshOutputBuiltin = [](IRInst* g)
+        {
+            if(const auto s = composeGetters<IRStringLit>(
+                g,
+                &IRInst::findDecoration<IRImportDecoration>,
+                &IRImportDecoration::getMangledNameOperand
+                ))
             {
-                return s;
+                const auto n = s->getStringSlice();
+                if (n == "gl_Position" ||
+                    n == "gl_PointSize" ||
+                    n == "gl_ClipDistance" ||
+                    n == "gl_CullDistance" ||
+                    n == "gl_PrimitiveID" ||
+                    n == "gl_Layer" ||
+                    n == "gl_ViewportIndex" ||
+                    n == "gl_CullPrimitiveEXT" ||
+                    n == "gl_PrimitiveShadingRateEXT")
+                {
+                    return s;
+                }
+            }
+            return (IRStringLit*)nullptr;
+        };
+        auto leaves = globalOutputVal.leafAddresses();
+        struct BuiltinOutputInfo
+        {
+            IRInst* param;
+            IRStringLit* nameDecoration;
+            IRType* type;
+            IRStructKey* key;
+        };
+        List<BuiltinOutputInfo> builtins;
+        for(auto leaf : leaves)
+        {
+            if(auto decoration = isMeshOutputBuiltin(leaf))
+            {
+                builtins.add({leaf, decoration, nullptr, nullptr});
             }
         }
-        return (IRStringLit*)nullptr;
-    };
-    auto leaves = globalOutputVal.leafAddresses();
-    struct BuiltinOutputInfo
-    {
-        IRInst* param;
-        IRStringLit* nameDecoration;
-        IRType* type;
-        IRStructKey* key;
-    };
-    List<BuiltinOutputInfo> builtins;
-    for(auto leaf : leaves)
-    {
-        if(auto decoration = isMeshOutputBuiltin(leaf))
+        if(builtins.getCount() == 0)
+        {
+            return;
+        }
+        const auto _locScope = IRBuilderInsertLocScope{builder};
+        builder->setInsertBefore(func);
+        auto meshOutputBlockType = builder->createStructType();
         {
-            builtins.add({leaf, decoration, nullptr, nullptr});
+            const auto _locScope2 = IRBuilderInsertLocScope{builder};
+            builder->setInsertInto(meshOutputBlockType);
+            for(auto& builtin : builtins)
+            {
+                auto t = composeGetters<IRType>(
+                    builtin.param,
+                    &IRInst::getFullType,
+                    &IROutTypeBase::getValueType,
+                    &IRArrayTypeBase::getElementType);
+                auto key = builder->createStructKey();
+                auto n = builtin.nameDecoration->getStringSlice();
+                builder->addImportDecoration(key, n);
+                builder->createStructField(meshOutputBlockType, key, t);
+                builtin.type = t;
+                builtin.key = key;
+            }
+        }
+
+        // No emitter actually handles GLSLOutputParameterGroupTypes, this isn't a
+        // problem as it's used as an intrinsic.
+        // GLSL does permit redeclaring these particular ones, so it might be nice to
+        // add a linkage decoration instead of it being an intrinsic in the event
+        // that we start outputting the linkage decoration instead of it being an
+        // intrinsic in the event that we start outputting these.
+        auto blockParamType = builder->getGLSLOutputParameterGroupType(
+            builder->getArrayType(meshOutputBlockType, meshOutputType->getMaxElementCount()));
+        auto blockParam = builder->createGlobalParam(blockParamType);
+        bool isPerPrimitive = as<IRPrimitivesType>(meshOutputType);
+        auto typeName = isPerPrimitive ? "gl_MeshPerPrimitiveEXT" : "gl_MeshPerVertexEXT";
+        auto arrayName = isPerPrimitive ? "gl_MeshPrimitivesEXT" : "gl_MeshVerticesEXT";
+        builder->addTargetIntrinsicDecoration(
+            meshOutputBlockType,
+            CapabilitySet(CapabilityAtom::GLSL),
+            UnownedStringSlice(typeName));
+        builder->addImportDecoration(blockParam, UnownedStringSlice(arrayName));
+        if(isPerPrimitive)
+        {
+            builder->addDecoration(blockParam, kIROp_GLSLPrimitivesRateDecoration);
+        }
+        // // While this is probably a correct thing to do, LRK::VaryingOutput
+        // // isn't really used for redeclaraion of builtin outputs, and assumes
+        // // that it's got a layout location, the correct fix might be to add
+        // // LRK::BuiltinVaryingOutput, but that would be polluting LRK for no
+        // // real gain.
+        // //
+        // IRVarLayout::Builder varLayoutBuilder{builder, IRTypeLayout::Builder{builder}.build()};
+        // varLayoutBuilder.findOrAddResourceInfo(LayoutResourceKind::VaryingOutput);
+        // varLayoutBuilder.setStage(Stage::Mesh);
+        // builder->addLayoutDecoration(blockParam, varLayoutBuilder.build());
+
+        for(auto builtin : builtins)
+        {
+            traverseUsers(builtin.param, [&](IRInst* u)
+            {
+                auto p = as<IRGetElementPtr>(u);
+                SLANG_ASSERT(p && "Mesh Output sentinel parameter wasn't used as an array");
+
+                IRBuilderInsertLocScope locScope{builder};
+                builder->setInsertBefore(p);
+                auto e = builder->emitElementAddress(builder->getPtrType(meshOutputBlockType), blockParam, p->getIndex());
+                auto a = builder->emitFieldAddress(builder->getPtrType(builtin.type), e, builtin.key);
+
+                p->replaceUsesWith(a);
+            });
         }
     }
-    if(builtins.getCount() == 0)
-    {
-        return;
-    }
-    const auto _locScope = IRBuilderInsertLocScope{builder};
-    builder->setInsertBefore(func);
-    auto meshOutputBlockType = builder->createStructType();
-    {
-        const auto _locScope2 = IRBuilderInsertLocScope{builder};
-        builder->setInsertInto(meshOutputBlockType);
-        for(auto& builtin : builtins)
-        {
-            auto t = composeGetters<IRType>(
-                builtin.param,
-                &IRInst::getFullType,
-                &IROutTypeBase::getValueType,
-                &IRArrayTypeBase::getElementType);
-            auto key = builder->createStructKey();
-            auto n = builtin.nameDecoration->getStringSlice();
-            builder->addImportDecoration(key, n);
-            builder->createStructField(meshOutputBlockType, key, t);
-            builtin.type = t;
-            builtin.key = key;
-        }
-    }
-
-    // No emitter actually handles GLSLOutputParameterGroupTypes, this isn't a
-    // problem as it's used as an intrinsic.
-    // GLSL does permit redeclaring these particular ones, so it might be nice to
-    // add a linkage decoration instead of it being an intrinsic in the event
-    // that we start outputting the linkage decoration instead of it being an
-    // intrinsic in the event that we start outputting these.
-    auto blockParamType = builder->getGLSLOutputParameterGroupType(
-        builder->getArrayType(meshOutputBlockType, meshOutputType->getMaxElementCount()));
-    auto blockParam = builder->createGlobalParam(blockParamType);
-    bool isPerPrimitive = as<IRPrimitivesType>(meshOutputType);
-    auto typeName = isPerPrimitive ? "gl_MeshPerPrimitiveEXT" : "gl_MeshPerVertexEXT";
-    auto arrayName = isPerPrimitive ? "gl_MeshPrimitivesEXT" : "gl_MeshVerticesEXT";
-    builder->addTargetIntrinsicDecoration(
-        meshOutputBlockType,
-        CapabilitySet(CapabilityAtom::GLSL),
-        UnownedStringSlice(typeName));
-    builder->addImportDecoration(blockParam, UnownedStringSlice(arrayName));
-    if(isPerPrimitive)
-    {
-        builder->addDecoration(blockParam, kIROp_GLSLPrimitivesRateDecoration);
-    }
-    // // While this is probably a correct thing to do, LRK::VaryingOutput
-    // // isn't really used for redeclaraion of builtin outputs, and assumes
-    // // that it's got a layout location, the correct fix might be to add
-    // // LRK::BuiltinVaryingOutput, but that would be polluting LRK for no
-    // // real gain.
-    // //
-    // IRVarLayout::Builder varLayoutBuilder{builder, IRTypeLayout::Builder{builder}.build()};
-    // varLayoutBuilder.findOrAddResourceInfo(LayoutResourceKind::VaryingOutput);
-    // varLayoutBuilder.setStage(Stage::Mesh);
-    // builder->addLayoutDecoration(blockParam, varLayoutBuilder.build());
-
-    for(auto builtin : builtins)
-    {
-        traverseUsers(builtin.param, [&](IRInst* u)
-        {
-            auto p = as<IRGetElementPtr>(u);
-            SLANG_ASSERT(p && "Mesh Output sentinel parameter wasn't used as an array");
-
-            IRBuilderInsertLocScope locScope{builder};
-            builder->setInsertBefore(p);
-            auto e = builder->emitElementAddress(builder->getPtrType(meshOutputBlockType), blockParam, p->getIndex());
-            auto a = builder->emitFieldAddress(builder->getPtrType(builtin.type), e, builtin.key);
-
-            p->replaceUsesWith(a);
-        });
-    }
+
+    SLANG_ASSERT(!g->hasUses());
+    g->removeAndDeallocate();
 }
 
 void legalizeEntryPointParameterForGLSL(
@@ -2992,10 +3007,16 @@ void legalizeDispatchMeshPayloadForGLSL(IRModule* module)
                 // parameter and store into the value being passed to this
                 // call.
                 builder.setInsertInto(module->getModuleInst());
-                const auto v = builder.emitVar(payloadType, SpvStorageClassTaskPayloadWorkgroupEXT);
+                const auto v = builder.createGlobalVar(payloadType, SpvStorageClassTaskPayloadWorkgroupEXT);
                 v->setFullType(builder.getRateQualifiedType(builder.getGroupSharedRate(), v->getFullType()));
                 builder.setInsertBefore(call);
-                builder.emitStore(v, payload);
+                builder.emitStore(v, builder.emitLoad(payload));
+
+                // Then, make sure that it's this new global which is being
+                // passed into the call to DispatchMesh, this is unimportant
+                // for GLSL which ignores such a parameter, but the SPIR-V
+                // backend depends on it being the global
+                call->getArgs()[3].set(v);
             }
         }
     });
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 37d4e7e18..815966218 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -3668,6 +3668,9 @@ public:
     IRFunc* createFunc();
     IRGlobalVar* createGlobalVar(
         IRType* valueType);
+    IRGlobalVar* createGlobalVar(
+        IRType*         valueType,
+        IRIntegerValue  addressSpace);
     IRGlobalParam* createGlobalParam(
         IRType* valueType);
     
diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index c45378fa3..f0517ed98 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -4,6 +4,7 @@
 #include "slang-ir-glsl-legalize.h"
 
 #include "slang-ir-clone.h"
+#include "slang-ir-legalize-mesh-outputs.h"
 #include "slang-ir.h"
 #include "slang-ir-insts.h"
 #include "slang-emit-base.h"
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 5fa558e15..675667dbc 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -4288,6 +4288,20 @@ namespace Slang
         return globalVar;
     }
 
+    IRGlobalVar* IRBuilder::createGlobalVar(
+        IRType*         valueType,
+        IRIntegerValue  addressSpace)
+    {
+        auto ptrType = getPtrType(kIROp_PtrType, valueType, addressSpace);
+        IRGlobalVar* globalVar = createInst<IRGlobalVar>(
+            this,
+            kIROp_GlobalVar,
+            ptrType);
+        _maybeSetSourceLoc(globalVar);
+        addGlobalValue(this, globalVar);
+        return globalVar;
+    }
+
     IRGlobalParam* IRBuilder::createGlobalParam(
         IRType* valueType)
     {
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 452b24c30..ce577de6e 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -2515,9 +2515,15 @@ void addCallArgsForParam(
     /// Compute the direction for a parameter based on its declaration
 ParameterDirection getParameterDirection(VarDeclBase* paramDecl)
 {
-    if( paramDecl->hasModifier<RefModifier>() )
+    if( paramDecl->hasModifier<RefModifier>()
+        || paramDecl->hasModifier<HLSLPayloadModifier>() )
     {
-        // The AST specified `ref`:
+        // The AST specified `ref` or `payload`:
+
+        // The payload types are a groupshared variable, and we really don't
+        // want to copy that into registers in every invocation on platforms
+        // where this matters, so treat them as by-reference here.
+
         return kParameterDirection_Ref;
     }
     if( paramDecl->hasModifier<InOutModifier>() )