Allow duplication of OpTypeNodePayloadArrayAMDX

This commit is mainly to allow OpTypeNodePayloadArrayAMDX to be
duplicated in SPIRV when using WorkGraphs.

AMD spec requires that there should be dedicated type declaration with
OpTypeNodePayloadArrayAMDX for each "node". And each of them should be
decorated with PayloadNodeNameAMDX.
https://docs.vulkan.org/features/latest/features/proposals/VK_AMDX_shader_enqueue.html

This behavior is very different from how "types" are declared. Normally
types are declared only once and they are shared across the whole
program. But for the work-graphs, each "RecordType" belongs to each node
and they need to be declared for each node.

Consider the following HLSL example where a producer triggers two
consumer nodes. Each node will be differentiated by "NodeID" attribute.

```
[Shader("node")]
[NodeLaunch("broadcasting")]
[NumThreads(4,1,1)]
void myFancyNode(...,
    [NodeID("myNode1")] NodeOutput<MY_RECORD> myRecords1,
    [NodeID("myNode2")] NodeOutput<MY_RECORD> myRecords2,
    ...)
{
    ThreadNodeOutputRecords<MY_RECORD> myRecord1 = myRecords1.GetThreadNodeOutputRecords(1);
    myRecord1.myData = 1;
    myRecord1.OutputCompelete();

    ThreadNodeOutputRecords<MY_RECORD> myRecord2 = myRecords2.GetThreadNodeOutputRecords(1);
    myRecord2.myData = 1;
    myRecord2.OutputCompelete();
}
```

Even though both are using the same type, "MY_RECORD", each of them
needs to be declared separately. The generated SPIR-V code for the
HLSL above is something like the following,
```
%str_node1 = OpConstantStringAMDX "myNode1";
%str_node2 = OpConstantStringAMDX "myNode1";

%type_node1 = OpTypeNodePayloadArrayAMDX %type_myRecord;
%type_node2 = OpTypeNodePayloadArrayAMDX %type_myRecord;

OpDecorateId %type_node1 PayloadNodeNameAMDX %str_node1;
OpDecorateId %type_node2 PayloadNodeNameAMDX %str_node2;
```
Note that `%type_node1` and `%type_node2` are equivalent, but they must
be declared separately because each of them will be decorated with
different nodeID string.

Author: jkwak-work

source/slang/core.meta.slang

@@ -4211,6 +4211,25 @@ attribute_syntax [QuadDerivatives] : QuadDerivativesAttribute;
 __attributeTarget(FuncDecl)
 attribute_syntax [RequireFullQuads] : RequireFullQuadsAttribute;
 
-__generic<T>
-typealias NodePayloadPtr = Ptr<T, $( (uint64_t)AddressSpace::NodePayloadAMDX)>;
+// Work-graphs
+
+/// @internal
+/// __SPIRVNodePayloadArray is to emit OpTypeNodePayloadArrayAMDX and
+/// SpvDecorationPayloadNodeNameAMDX for it.
+///
+__generic<RecordType, let nodeID : int>
+__intrinsic_type($(kIROp_SPIRVNodePayloadArrayType))
+struct __SPIRVNodePayloadArray;
+
+/// @internal
+/// NodePayloadArrayPtr is a pointer type for "NodePayloadAMDX" storage-class and it points to an array of work-graph nodes
+///
+__generic<RecordType, let nodeID : int>
+typealias NodePayloadArrayPtr = Ptr<__SPIRVNodePayloadArray<RecordType, nodeID>, $( (uint64_t)AddressSpace::NodePayloadAMDX)>;
+
+/// @internal
+/// NodePayloadPtr is a pointer type for "NodePayloadAMDX" storage-class and it points to a `RecordType`.
+///
+__generic<RecordType>
+typealias NodePayloadPtr = Ptr<RecordType, $( (uint64_t)AddressSpace::NodePayloadAMDX)>;
 

source/slang/hlsl.meta.slang

@@ -24371,28 +24371,85 @@ int8_t4_packed pack_clamp_s8(int16_t4 unpackedValue)
 
 // Work-graphs
 
+///@public
+/// Set of zero or more records for each thread.
+///
+__generic<RecordType, let nodeID : int>
+struct ThreadNodeOutputRecords
+{
+    NodePayloadArrayPtr<RecordType, nodeID> alloc;
+
+    [ForceInline]
+    [require(spirv)] // [require(spirv, thread)]
+    __init(int recordCount, int nodeIndex = 0)
+    {
+        alloc = spirv_asm
+        {
+            result = OpAllocateNodePayloadsAMDX $$NodePayloadArrayPtr<RecordType, nodeID> Workgroup $recordCount $nodeIndex;
+        };
+    }
+
+    [ForceInline]
+    [require(hlsl_spirv)] // [require(hlsl_spirv, thread)]
+    void OutputComplete()
+    {
+        __target_switch
+        {
+        case hlsl: __intrinsic_asm ".OutputComplete";
+        case spirv: spirv_asm { OpEnqueueNodePayloadsAMDX $alloc; };
+        }
+    }
+
+    __subscript(uint index) -> NodePayloadPtr<RecordType>
+    {
+        [ForceInline]
+        get { return Get(index); }
+    }
+
+    [ForceInline]
+    [require(hlsl_spirv)] // [require(hlsl_spirv, thread)]
+    NodePayloadPtr<RecordType> Get(int index = 0)
+    {
+        __target_switch
+        {
+        case hlsl: __intrinsic_asm ".Get";
+        case spirv:
+            NodePayloadArrayPtr<RecordType, nodeID> ptr = spirv_asm
+            {
+                %ptr  = OpTypePointer Function $$__SPIRVNodePayloadArray<RecordType, nodeID>;
+                %var  = OpVariable %ptr Function;
+                        OpStore %var $alloc;
+                %Load = OpLoad $$NodePayloadArrayPtr<RecordType, nodeID> %var;
+                result : $$NodePayloadArrayPtr<RecordType, nodeID> = OpAccessChain %Load $index
+            };
+            return reinterpret<NodePayloadPtr<RecordType>>(ptr);
+        }
+    }
+};
+
 //@public:
-/// read-only input to Broadcasting launch node.
-__generic<T>
-//TODO: DispatchNodeInputRecord should be available only for broadcasting node shader.
-//[require(broadcasting_node)]
-[require(spirv)]
+/// Read-only input to Broadcasting launch node.
+///
+__generic<RecordType, let nodeID : int>
 struct DispatchNodeInputRecord
 {
     /// Provide an access to a record object that only holds a single record.
-    NodePayloadPtr<T> Get()
+    [ForceInline]
+    [require(hlsl_spirv)] // [require(hlsl_spirv, broadcasting)]
+    NodePayloadPtr<RecordType> Get()
     {
         int index = 0;
         __target_switch
         {
+        case hlsl: __intrinsic_asm ".Get";
         case spirv:
-            return spirv_asm
+            NodePayloadArrayPtr<RecordType, nodeID> ptr = spirv_asm
             {
-                %in_payload_t = OpTypeNodePayloadArrayAMDX $$T;
-                %in_payload_ptr_t = OpTypePointer NodePayloadAMDX %in_payload_t;
-                %var = OpVariable %in_payload_ptr_t NodePayloadAMDX;
-                result : $$NodePayloadPtr<T> = OpAccessChain %var $index;
+                %ptr = OpTypePointer NodePayloadAMDX $$__SPIRVNodePayloadArray<RecordType, nodeID>;
+                %var = OpVariable %ptr NodePayloadAMDX;
+                result : $$NodePayloadArrayPtr<RecordType, nodeID> = OpAccessChain %var $index;
             };
+            return reinterpret<NodePayloadPtr<RecordType>>(ptr);
         }
     }
 };

source/slang/slang-emit-spirv-ops.h

@@ -2571,12 +2571,36 @@ template<typename T>
 SpvInst* emitOpTypeNodePayloadArray(IRInst* inst, const T& type)
 {
     static_assert(isSingular<T>);
-    return emitInstMemoized(
+    return emitInst(
         getSection(SpvLogicalSectionID::ConstantsAndTypes),
         inst,
         SpvOpTypeNodePayloadArrayAMDX,
         kResultID,
         type);
 }
 
+// https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/AMD/SPV_AMDX_shader_enqueue.html#OpTypeNodePayloadArrayAMDX
+SpvInst* emitOpConstantString(IRInst* inst, const UnownedStringSlice& str)
+{
+    return emitInstMemoized(
+        getSection(SpvLogicalSectionID::ConstantsAndTypes),
+        inst,
+        SpvOpConstantStringAMDX,
+        kResultID,
+        str);
+}
+
+// https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpDecorateId
+template<typename T1, typename T2>
+SpvInst* emitOpDecoratePayloadNodeName(IRInst* inst, const T1& target, const T2& id)
+{
+    return emitInst(
+        getSection(SpvLogicalSectionID::Annotations),
+        inst,
+        SpvOpDecorateId,
+        target,
+        SpvDecorationPayloadNodeNameAMDX,
+        id);
+}
+
 #endif // SLANG_IN_SPIRV_EMIT_CONTEXT

source/slang/slang-emit-spirv.cpp

@@ -1542,8 +1542,7 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                 else if (storageClass == SpvStorageClassNodePayloadAMDX)
                 {
                     auto spvValueType = ensureInst(valueType);
-                    auto spvNodePayloadType = emitOpTypeNodePayloadArray(inst, spvValueType);
-                    valueTypeId = getID(spvNodePayloadType);
+                    valueTypeId = getID(spvValueType);
                 }
                 else
                 {
@@ -1909,6 +1908,24 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
         case kIROp_IndicesType:
         case kIROp_PrimitivesType:
             return nullptr;
+        case kIROp_SPIRVNodePayloadArrayType:
+            if (auto nodePayloadArrayType = as<IRSPIRVNodePayloadArrayType>(inst))
+            {
+                auto newType = emitOpTypeNodePayloadArray(
+                    inst, nodePayloadArrayType->getRecordType());
+
+                Slang::StringBuilder str;
+                str << "NodeID_" << uint32_t(nodePayloadArrayType->getNodeID()->getValue());
+                SpvInst* spvStr = emitOpConstantString(nullptr, str.getUnownedSlice());
+                (void)spvStr;
+
+                auto r = emitOpDecoratePayloadNodeName(
+                    nullptr,
+                    newType,
+                    spvStr);
+                (void)r;
+                return newType;
+            }
         default:
             {
                 if (as<IRSPIRVAsmOperand>(inst))
@@ -7750,6 +7767,7 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                         return getSection(SpvLogicalSectionID::Annotations);
                     case SpvOpTypeNodePayloadArrayAMDX:
                         return getSection(SpvLogicalSectionID::ConstantsAndTypes);
+
                     default:
                         return defaultParent;
                     }

source/slang/slang-ir-inst-defs.h

@@ -265,6 +265,9 @@ INST(ExpandTypeOrVal, ExpandTypeOrVal, 1, HOISTABLE)
 // A type that identifies it's contained type as being emittable as `spirv_literal.
 INST(SPIRVLiteralType, spirvLiteralType, 1, HOISTABLE)
 
+// WorkGraphs
+INST(SPIRVNodePayloadArrayType, spirvNodePayloadArrayType, 2, HOISTABLE)
+
 // A TypeType-typed IRValue represents a IRType.
 // It is used to represent a type parameter/argument in a generics.
 INST(TypeType, type_t, 0, HOISTABLE)

source/slang/slang-ir.h

@@ -1752,6 +1752,14 @@ struct IRSPIRVLiteralType : IRType
     IRType* getValueType() { return static_cast<IRType*>(getOperand(0)); }
 };
 
+struct IRSPIRVNodePayloadArrayType : IRType
+{
+    IR_LEAF_ISA(SPIRVNodePayloadArrayType)
+
+    IRType* getRecordType() { return static_cast<IRType*>(getOperand(0)); }
+    IRIntLit* getNodeID() { return static_cast<IRIntLit*>(getOperand(1)); }
+};
+
 struct IRPtrTypeBase : IRType
 {
     IRType* getValueType() { return (IRType*)getOperand(0); }

tests/workgraphs/consumer.slang

@@ -1,25 +1,35 @@
-//TEST:SIMPLE(filecheck=CHK): -target spirv-asm -stage compute -entry main -skip-spirv-validation
+//TEST:SIMPLE(filecheck=CHK): -target spirv-asm -stage compute -entry computeMain -skip-spirv-validation
+
+RWStructuredBuffer<int> outputBuffer;
+
 struct RecordData
 {
     int myData;
 };
 
 [shader("compute")]
 [numthreads(1, 1, 1)]
-void main(uint3 dispatchThreadId : SV_GroupThreadID)
+void computeMain(uint3 dispatchThreadId : SV_GroupThreadID)
 {
     spirv_asm
     {
-        OpExecutionMode $main ShaderIndexAMDX $(0);
-        OpExecutionMode $main StaticNumWorkgroupsAMDX $(1) $(1) $(1);
+        OpExecutionMode $computeMain ShaderIndexAMDX $(0);
+        OpExecutionMode $computeMain StaticNumWorkgroupsAMDX $(1) $(1) $(1);
     };
 
-    DispatchNodeInputRecord<RecordData> inputData;
+    //TODO: "myNodeID" is supposed to come from [NodeID("name")] attribute of the entry point.
+    // Until it is implemented properly, we will use an int-type generic argument.
+    //
+    #define myNodeID 0
+    DispatchNodeInputRecord<RecordData, myNodeID> inputData;
 
-    let recordData = inputData.Get();
-    int myData = recordData.myData;
+    int myData = inputData.Get().myData;
+    outputBuffer[dispatchThreadId.x] = myData;
 }
 
+//CHK: OpCapability ShaderEnqueueAMDX
+//CHK: OpExtension "SPV_AMDX_shader_enqueue"
+
 //CHK: ; Types, variables and constants
 //CHK: [[MemberType:%[a-zA-Z_0-9]+]] = OpTypeInt 32 1
 //CHK: [[StructType:%[a-zA-Z_0-9]+]] = OpTypeStruct [[MemberType]]

tests/workgraphs/producer.slang

@@ -0,0 +1,53 @@
+//TEST:SIMPLE(filecheck=CHK): -target spirv-asm -stage compute -entry computeMain -skip-spirv-validation
+
+RWStructuredBuffer<int> inputBuffer;
+
+struct RecordData
+{
+    int myData;
+}
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    spirv_asm
+    {
+        OpExecutionMode $computeMain ShaderIndexAMDX $(0);
+        OpExecutionMode $computeMain StaticNumWorkgroupsAMDX $(1) $(1) $(1);
+    };
+
+    //TODO: "myNodeID" is supposed to come from [NodeID("name")] attribute on "node" variable.
+    // Until it is implemented properly, we will use an int-type generic argument.
+    //
+    #define myNodeID 0
+    ThreadNodeOutputRecords<RecordData, myNodeID> node = { 1 };
+
+    node.Get().myData = inputBuffer[0];
+    node.OutputComplete();
+}
+
+//CHK: OpCapability ShaderEnqueueAMDX
+//CHK: OpExtension "SPV_AMDX_shader_enqueue"
+
+//CHK-DAG: [[RecordType:%[a-zA-Z_0-9]+]] = OpTypeNodePayloadArrayAMDX %RecordData
+//CHK-DAG: [[PtrType:%[a-zA-Z_0-9]+]] = OpTypePointer NodePayloadAMDX [[RecordType]]
+//CHK-DAG: [[NodeID:%[a-zA-Z_0-9]+]] = OpConstantStringAMDX "NodeID_0"
+//CHK-DAG: OpDecorateId [[RecordType]] PayloadNodeNameAMDX [[NodeID]]
+
+//CHK-NOT: = OpTypeNodePayloadArrayAMDX
+//CHK-NOT: = OpConstantStringAMDX
+//CHK-NOT: OpDecorateID {{.*}} PayloadNodeNameAMDX
+
+// ThreadNodeOutputRecords::__init()
+//CHK: [[Alloc:%[a-zA-Z_0-9]+]] = OpAllocateNodePayloadsAMDX [[PtrType]] %
+//CHK-NOT: = OpAllocateNodePayloadsAMDX
+
+// ThreadNodeOutputRecords::Get()
+//CHK: [[Load:%[a-zA-Z_0-9]+]] = OpLoad [[PtrType]] %var
+//CHK: = OpAccessChain [[PtrType]] [[Load]] %
+
+// ThreadNodeOutputRecords::OutputComplete()
+//CHK: OpEnqueueNodePayloadsAMDX [[Alloc]]
+//CHK-NOT: OpEnqueueNodePayloadsAMDX
+