VK_AMDX_shader_enqueue.proposal

This extension adds the ability for developers to enqueue mesh pipelines and compute shader workgroups from other compute shaders.

Problem Statement

Applications are increasingly using more complex renderers, often incorporating multiple compute passes that classify, sort, or otherwise preprocess input data. These passes may be used to determine how future work is performed on the GPU; but triggering that future GPU work requires either a round trip to the host, or going through buffer memory and using indirect commands. Host round trips necessarily include more system bandwidth and latency as command buffers need to be built and transmitted back to the GPU. Indirect commands work well in many cases, but they have little flexibility when it comes to determining what is actually dispatched; they must be enqueued ahead of time, synchronized with heavy API barriers, and execute with a single pre-recorded pipeline.

Whilst latency can be hidden and indirect commands can work in many cases where additional latency and bandwidth is not acceptable, recent engine developments such as Unreal 5’s Nanite technology explicitly require the flexibility of shader selection and low latency. A desirable solution should be able to have the flexibility required for these systems, while keeping the execution loop firmly on the GPU.

Solution Space

Three main possibilities exist:

1. Extend indirect commands
2. VK_NV_device_generated_commands
3. Shader enqueue

More flexible indirect commands could feasibly allow things like shader selection, introduce more complex flow control, or include indirect state setting commands. The main issue with these is that these always require parameters to be written through regular buffer memory, and that buffer memory has to be sized for each indirect command to handle the maximum number of possibilities. As well as the large allocation size causing memory pressure, pushing all that data through buffer memory will reduce the bandwidth available for other operations. All of this could cause bottlenecks elsewhere in the pipeline. Hypothetically a new interface for better scheduling/memory management could be introduced, but that starts looking a lot like option 3.

Option 2 - implementing a cross-vendor equivalent of VK_NV_device_generated_commands would be a workable solution that adds both flexibility and avoids a CPU round trip. The reason it has not enjoyed wider support is due to concerns about how the commands are generated - it uses a tokenised API which has to be processed by the GPU before it can be executed. For existing GPUs this can mean doing things like running a single compute shader invocation to process each token stream into a runnable command buffer, adding both latency and bandwidth on the GPU.

Option 3 - OpenCL and CUDA have had some form of shader enqueue API for a while, where the focus has typically been primarily on enabling developers and on compute workloads. From a user interface perspective these have had a decent amount of battle testing and is quite a popular and flexible interface.

This proposal is built around something like Option 3, but extended to be explicit and performant.

Proposal

API Changes

Graph Pipelines

In order to facilitate dispatch of multiple shaders from the GPU, the implementation needs some information about how pipelines will be launched and synchronized. This proposal introduces a new execution graph pipeline that defines execution paths between multiple shaders, and allows dynamic execution of different shaders.

VkResult vkCreateExecutionGraphPipelinesAMDX(
    VkDevice                                        device,
    VkPipelineCache                                 pipelineCache,
    uint32_t                                        createInfoCount,
    const VkExecutionGraphPipelineCreateInfoAMDX*   pCreateInfos,
    const VkAllocationCallbacks*                    pAllocator,
    VkPipeline*                                     pPipelines);

typedef struct VkExecutionGraphPipelineCreateInfoAMDX {
    VkStructureType                             sType;
    const void*                                 pNext;
    VkPipelineCreateFlags                       flags;
    uint32_t                                    stageCount;
    const VkPipelineShaderStageCreateInfo*      pStages;
    const VkPipelineLibraryCreateInfoKHR*       pLibraryInfo;
    VkPipelineLayout                            layout;
    VkPipeline                                  basePipelineHandle;
    int32_t                                     basePipelineIndex;
} VkExecutionGraphPipelineCreateInfoAMDX;

Shaders defined by pStages and any pipelines in pLibraryInfo→pLibraries define the possible nodes of the graph. The linkage between nodes however is defined wholly in shader code, though may be overridden by specialization constants in many cases.

Shaders in pStages must be in the GLCompute execution model, and may have the CoalescingAMDX execution mode.

Pipelines in pLibraries can be compute pipelines, graphics pipelines, or other execution graph pipelines. Compute and graphics pipelines must be created with the VK_PIPELINE_CREATE_2_LIBRARY_BIT_KHR and VK_PIPELINE_CREATE_2_EXECUTION_GRAPH_BIT_AMDX flag bits. Execution graph pipelines used as libraries must be created with the VK_PIPELINE_CREATE_2_LIBRARY_BIT_KHR flag bit.

VK_PIPELINE_CREATE_2_EXECUTION_GRAPH_BIT_AMDX = 0x100000000ULL

Each shader in an execution graph is associated with a name and an index, which are used to identify the target shader when dispatching a payload. The VkPipelineShaderStageNodeCreateInfoAMDX provides options for specifying how the shader is specified with regards to its entry point name and index, and can be chained to the VkPipelineShaderStageCreateInfo structure.

const uint32_t VK_SHADER_INDEX_UNUSED_AMDX = 0xFFFFFFFF;

typedef struct VkPipelineShaderStageNodeCreateInfoAMDX {
    VkStructureType                             sType;
    const void*                                 pNext;
    const char*                                 pName;
    uint32_t                                    index;
} VkPipelineShaderStageNodeCreateInfoAMDX;

• index sets the index value for a shader.
• pName allows applications to override the name specified in SPIR-V by OpEntryPoint.

If pName is NULL then the original name is used, as specified by VkPipelineShaderStageCreateInfo::pName. If index is VK_SHADER_INDEX_UNUSED_AMDX then the original index is used, either as specified by the ShaderIndexAMDX Execution Mode, or 0 if that too is not specified. If this structure is not provided, pName defaults to NULL, and index defaults to VK_SHADER_INDEX_UNUSED_AMDX.

When dispatching from another shader, the index is dynamic and can be specified in uniform control flow - however the name must be statically declared as a decoration on the payload. Allowing the index to be set dynamically lets applications stream shaders in and out dynamically, by simply changing constant data and relinking the graph pipeline from new libraries. Shaders with the same name and different indexes must consume identical payloads and have the same execution model. Shaders with the same name in an execution graph pipeline must have unique indexes.

When dispatching from another shader, any declared input payload for the dispatched node must be less than or equal to the size of the output payload in the dispatching node. Additionally, if an input payload is declared in the dispatched shader, the input and output payloads must specify members with the same decorations at the same offsets.

Graphics Pipeline State

When adding a graphics pipeline to an execution graph pipeline, applications must specify a graphics pipeline with a complete set of state, and the VK_PIPELINE_CREATE_2_LIBRARY_BIT_KHR and VK_PIPELINE_CREATE_2_EXECUTION_GRAPH_BIT_AMDX flags set. Graphics pipelines must only include mesh shaders; vertex shader pipelines or mesh pipelines with task shaders are not supported. When creating such a graphics pipeline from libraries as an interaction with VK_EXT_graphics_pipeline_library, those libraries must also have been created with those flags.

For graphics pipelines defined in this way, only the following dynamic state is allowed:

• VK_DYNAMIC_STATE_VIEWPORT
• VK_DYNAMIC_STATE_SCISSOR
• VK_DYNAMIC_STATE_LINE_WIDTH
• VK_DYNAMIC_STATE_DEPTH_BIAS
• VK_DYNAMIC_STATE_BLEND_CONSTANTS
• VK_DYNAMIC_STATE_DEPTH_BOUNDS
• VK_DYNAMIC_STATE_VIEWPORT_WITH_COUNT
• VK_DYNAMIC_STATE_SCISSOR_WITH_COUNT
• VK_DYNAMIC_STATE_SAMPLE_LOCATIONS_EXT
• VK_DYNAMIC_STATE_FRAGMENT_SHADING_RATE_KHR

When these dynamic states are specified, this state is captured from the command buffer state at the point the execution graph is dispatched, and applies to all nodes that have that state set dynamically executed as part of that dispatch. All graphics pipelines in an execution graph must use the same set of dynamic states. Applications can dynamically choose any other state at runtime by selecting between pipelines with different state when dispatching, but the underlying pipelines must be created statically.

When included as a library in an execution graph pipeline, the node is defined by the first shader in the graphics pipeline.

Scratch Memory

Implementations may need scratch memory to manage dispatch queues or similar when executing a pipeline graph, and this is explicitly managed by the application.

typedef struct VkExecutionGraphPipelineScratchSizeAMDX {
    VkStructureType                     sType;
    void*                               pNext;
    VkDeviceSize                        minSize;
    VkDeviceSize                        maxSize;
    VkDeviceSize                        sizeGranularity;
} VkExecutionGraphPipelineScratchSizeAMDX;

VkResult vkGetExecutionGraphPipelineScratchSizeAMDX(
    VkDevice                                    device,
    VkPipeline                                  executionGraph,
    VkExecutionGraphPipelineScratchSizeAMDX*    pSizeInfo);

Applications can query the required amount of scratch memory for a given pipeline, and the address of a buffer of that size must be provided when calling vkCmdDispatchGraphAMDX. The amount of scratch memory needed by a given pipeline is related to the number and size of payloads across the whole graph; while the exact relationship is implementation dependent, reducing the number of unique nodes (different name string) and size of payloads can reduce scratch memory consumption.

A range of sizes are returned by the implementation; any size between minSize and maxSize can be used, though the actual memory consumed will be snapped to minSize + an integer multiple of sizeGranularity. Choosing any value less than the maximum size will reduce memory pressure but will likely result in degraded performance.

Buffers created for this purpose must use the new buffer usage flags:

VK_BUFFER_USAGE_EXECUTION_GRAPH_SCRATCH_BIT_AMDX
VK_BUFFER_USAGE_2_EXECUTION_GRAPH_SCRATCH_BIT_AMDX

Scratch memory needs to be initialized against a graph pipeline before it can be used with that graph for the first time, using the following command:

void vkCmdInitializeGraphScratchMemoryAMDX(
    VkCommandBuffer                             commandBuffer,
    VkPipeline                                  executionGraph,
    VkDeviceAddress                             scratch,
    VkDeviceSize                                scratchSize);

This command initializes it for the execution graph pipeline executionGraph with the specified scratchSize. Scratch memory will need to be re-initialized if it is going to be re-used with a different execution graph pipeline, but can be used with the same pipeline repeatedly without re-initialization. Scratch memory initialization can be synchronized using the compute pipeline stage VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT and shader write access flag VK_ACCESS_SHADER_WRITE_BIT.

Dispatch a graph

Once an execution graph has been created and scratch memory has been initialized for it, the following commands can be used to execute the graph:

typedef struct VkDispatchGraphInfoAMDX {
    uint32_t                                    nodeIndex;
    uint32_t                                    payloadCount;
    VkDeviceOrHostAddressConstAMDX              payloads;
    uint64_t                                    payloadStride;
} VkDispatchGraphInfoAMDX;

typedef struct VkDispatchGraphCountInfoAMDX {
    uint32_t                                    count;
    VkDeviceOrHostAddressConstAMDX              infos;
    uint64_t                                    stride;
} VkDispatchGraphCountInfoAMDX;

void vkCmdDispatchGraphAMDX(
    VkCommandBuffer                             commandBuffer,
    VkDeviceAddress                             scratch,
    VkDeviceSize                                scratchSize,
    const VkDispatchGraphCountInfoAMDX*         pCountInfo);

void vkCmdDispatchGraphIndirectAMDX(
    VkCommandBuffer                             commandBuffer,
    VkDeviceAddress                             scratch,
    VkDeviceSize                                scratchSize,
    const VkDispatchGraphCountInfoAMDX*         pCountInfo);

void vkCmdDispatchGraphIndirectCountAMDX(
    VkCommandBuffer                             commandBuffer,
    VkDeviceAddress                             scratch,
    VkDeviceSize                                scratchSize,
    VkDeviceAddress                             countInfo);

Each of the above commands enqueues payloads for an array of nodes in the bound execution graph pipeline, according to the contents of the VkDispatchGraphCountInfoAMDX and VkDispatchGraphInfoAMDX structures.

vkCmdDispatchGraphAMDX takes all of its arguments from the host pointers. VkDispatchGraphCountInfoAMDX::infos.hostAddress is a pointer to an array of VkDispatchGraphInfoAMDX structures, with stride equal to VkDispatchGraphCountInfoAMDX::stride and VkDispatchGraphCountInfoAMDX::count elements.

vkCmdDispatchGraphIndirectAMDX consumes most parameters on the host, but uses the device address for VkDispatchGraphCountInfoAMDX::infos, and also treating payloads parameters as device addresses.

vkCmdDispatchGraphIndirectCountAMDX consumes countInfo on the device and all child parameters also use device addresses.

Data consumed via a device address must be from buffers created with the VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT and VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT flags. payloads is a pointer to a linear array of payloads in memory, with a stride equal to payloadStride. payloadCount may be 0. The range of memory from scratch up to scratchSize may be used by the implementation to hold temporary data during graph execution, and can be synchronized using the compute pipeline stage and shader write access flags.

These dispatch commands must not be called in protected command buffers or secondary command buffers.

The size of the payload provided for each dispatched node must be at least as large as the NodePayloadAMDX declaration in the node, and the layout of the payload data in memory will be interpreted as it is laid out in the selected node’s shader, including any member decorations. In particular, this means for nodes that consume indirect parameters from the payload, those parameters must be provided in the correct location as specified in the shader. For example, for a compute shader that does not include a StaticNumWorkgroupsAMDX or CoalescingAMDX declaration, each dispatch will consume a payload structure containing a member decorated with PayloadDispatchIndirectAMDX that indicates the number of workgroups to dispatch in each dimension.

Node payload members must be explicitly laid out with offset and array stride decorations, both in the input and output.

• If the dispatched shader uses GLCompute or MeshEXT Execution Model, then it is allowed to not specify the input payload. In this case, the payload is defined implicitly as follows:
  • If the StaticNumWorkgroupsAMDX or CoalescingAMDX execution modes are specified, the payload is empty.
  • Otherwise, the payload is a structure with a single member that is a vector of three 32-bit unsigned integers.

Payloads are always read (including built-in values) according to the input payload definition - the output payload definition must have the same size as the expected input, but does not otherwise need to match. Applications must take care to ensure that values are where they expect them.

The nodeIndex is a unique integer identifier identifying a specific shader name and shader index (defined by VkPipelineShaderStageNodeCreateInfoAMDX) added to the executable graph pipeline. vkGetExecutionGraphPipelineNodeIndexAMDX can be used to query the identifier for a given node:

VkResult vkGetExecutionGraphPipelineNodeIndexAMDX(
    VkDevice                                        device,
    VkPipeline                                      executionGraph,
    const VkPipelineShaderStageNodeCreateInfoAMDX*  pNodeInfo,
    uint32_t*                                       pNodeIndex);

pNodeInfo specifies the shader name and index as set up when creating the pipeline, with the associated node index returned in pNodeIndex. When used with this function, pNodeInfo→pName must not be NULL.

Execution graph pipelines and their resources are bound using a new pipeline bind point:

VK_PIPELINE_BIND_POINT_EXECUTION_GRAPH_AMDX

Properties

The following new properties are added to Vulkan:

typedef VkPhysicalDeviceShaderEnqueuePropertiesAMDX {
    VkStructureType                     sType;
    void*                               pNext;
    uint32_t                            maxExecutionGraphDepth;
    uint32_t                            maxExecutionGraphShaderOutputNodes;
    uint32_t                            maxExecutionGraphShaderPayloadSize;
    uint32_t                            maxExecutionGraphShaderPayloadCount;
    uint32_t                            executionGraphDispatchAddressAlignment;
    uint32_t                            maxExecutionGraphWorkgroupCount[3];
    uint32_t                            maxExecutionGraphWorkgroups;
} VkPhysicalDeviceShaderEnqueuePropertiesAMDX;

Each limit is defined as follows:

• maxExecutionGraphDepth defines the maximum node chain length in the graph, and must be at least 32. A node that is dispatched with an API command is at depth 1 and the node that receives a payload from it is at depth 2, and so on. If a node uses tail recursion, each recursive call increases the depth by 1 as well.
• maxExecutionGraphShaderOutputNodes specifies the maximum number of unique nodes that can be dispatched from a single shader, and must be at least 256.
• maxExecutionGraphShaderPayloadSize specifies the maximum total size of payload declarations in a shader, and must be at least 32KB.
• maxExecutionGraphShaderPayloadCount specifies the maximum number of output payloads that can be initialized in a single workgroup, and must be at least 256.
• executionGraphDispatchAddressAlignment specifies the alignment of non-scratch VkDeviceAddress arguments consumed by graph dispatch commands, and must be no more than 4 bytes.
• maxExecutionGraphWorkgroupCount[3] describes the maximum number of local workgroups that a shader can be dispatched with, and must be at least (65535, 65535, 65535) for the X, Y, and Z dimensions, respectively.
• maxExecutionGraphWorkgroups describes the total number of local workgroups that a shader can be dispatched with and must be at least 16777215.

Features

The following new features are added to Vulkan:

typedef VkPhysicalDeviceShaderEnqueueFeaturesAMDX {
    VkStructureType                     sType;
    void*                               pNext;
    VkBool32                            shaderEnqueue;
    VkBool32                            shaderMeshEnqueue;
} VkPhysicalDeviceShaderEnqueueFeaturesAMDX;

The shaderEnqueue feature enables the ability to enqueue compute shader workgroups from other compute shaders. The shaderMeshEnqueue feature enables the ability to enqueue mesh nodes in an execution graph.

SPIR-V Changes

A new capability is added:

A new storage class is added:

An entry point must only declare one variable in the NodePayloadAMDX storage class in its interface.

New execution modes are added:

A shader module declaring ShaderEnqueueAMDX capability must only be used in execution graph pipelines created by vkCreateExecutionGraphPipelinesAMDX command.

MaxNodeRecursionAMDX must be specified if a shader re-enqueues itself, which takes place if that shader allocates and enqueues a payload for the same node name and index. Other forms of recursion are not allowed.

An application must not dispatch the shader with a number of workgroups in any dimension greater than the values specified by MaxNumWorkgroupsAMDX.

StaticNumWorkgroupsAMDX allows the declaration of the number of workgroups to dispatch to be coded into the shader itself, which can be useful for optimizing some algorithms. When a compute shader is dispatched using existing vkCmdDispatchGraph* commands, the workgroup counts specified there are overridden. When enqueuing such shaders with a payload, these arguments will not be consumed from the payload before application-specified data begins.

The values of MaxNumWorkgroupsAMDX and StaticNumWorkgroupsAMDX must be less than or equal to VkPhysicalDeviceShaderEnqueuePropertiesAMDX::maxExecutionGraphWorkgroupCount.

The product of the X, Y, and Z values of MaxNumWorkgroupsAMDX and StaticNumWorkgroupsAMDX must be less than or equal to VkPhysicalDeviceShaderEnqueuePropertiesAMDX::maxExecutionGraphWorkgroups.

The arguments to each of these execution modes must be a constant 32-bit integer value, and may be supplied via specialization constants.

When a GLCompute or MeshEXT shader is being used in an execution graph, NumWorkgroups must not be used.

When CoalescingAMDX is used, it has the following effects on a compute shader’s inputs and outputs:

• The WorkgroupId built-in is always (0,0,0)
• NB: This affects related built-ins like GlobalInvocationId
• So similar to StaticNumWorkgroupsAMDX, no dispatch size is consumed from the payload-specified
• The input in the NodePayloadAMDX storage class must have a type of OpTypeNodePayloadArrayAMDX.
• This input must be decorated with NodeMaxPayloadsAMDX, indicating the number of payloads that can be received.
• The number of payloads received can be queried through OpNodePayloadArrayLengthAMDX

When SharesInputWithAMDX is declared, the node will be dispatched whenever the node identified by it is dispatched, with the same input payload. The following limitations apply for sharing nodes in this way:

• Nodes must only share with a node that does not declare SharesInputWithAMDX
• No more than 256 nodes in a graph can share the same input (including the base node)
• Applications must not directly dispatch any node with the SharesInputWithAMDX execution mode.
• Input payloads must be decorated with NonWritable if SharesInputWithAMDX is declared.
• Emitting a payload to a shared node multiplies all of the payload resources by the number of shared nodes, as they count against values in VkPhysicalDeviceShaderEnqueuePropertiesAMDX.

If IsApiEntryAMDX is set to false, vkCmdDispatchGraph* commands must not reference this node.

New decorations are added:

The following new built-ins are provided:

If the Execution Model is GLCompute or MeshEXT, and neither the StaticNumWorkgroupsAMDX or CoalescingAMDX execution modes are specified, if an input payload is specified it must include a member with the PayloadDispatchIndirectAMDX decoration, indicating the number of workgroups to dispatch in each dimension.

New constant instructions are added to allow specialization of string variables, which are used for linkage between shaders.

A new payload type is defined that can be allocated dynamically and then enqueued for a node:

Decorations on this type indicate which node this type will be dispatched to and how it consumes resources. Once a payload array type has been declared and all relevant decorations specified, they can be allocated using:

Once a payload array is allocated, it can be enqueued to the identified node by calling OpEnqueueNodePayloadsAMDX. Enqueues are performed in the same manner as the vkCmdDispatchGraph* API commands. If the node receiving the payloads has the CoalescingAMDX execution mode, there is no guarantee what set of payloads are visible to the same workgroup.

The shader must not enqueue payloads to a shader with the same name as this shader unless the index identifies this node and MaxNodeRecursionAMDX is declared with a sufficient depth. Shaders with the same name and different indexes can each recurse independently.

Once this has been called, accessing any element of Payload Array is undefined behavior.

The length of Payload Array can be queried at any point by calling:

Before allocating payloads, applications can determine whether allocating payloads is possible for a particular node index:

• If a payload type is decorated with PayloadNodeSparseArrayAMDX, applications can determine whether a node exists at a particular index.
• If a payload type is decorated with PayloadNodeNameAMDX that matches the current node, applications can determine whether a node at a particular index has reached its max recursion depth.
• In all other cases, the payload can be allocated.

Payloads enqueued in this way will be provided to the node through the NodePayloadAMDX storage class in the shader. These payloads can be read by the receiving node, but also can be written for a limited amount of communication between multiple workgroups enqueued for the same node. It is a data race if one workgroup writes to a particular element of the payload and another workgroup accesses it in any way, with one exception; once all nodes have finished writing, it is safe for the last node to read those values. Workgroups can indicate that they have finished writing to the payload by calling:

Once this has been called for a given payload, writing values into that payload by the current invocation/workgroup is undefined behavior.

Issues

How does this extension interact with device groups?

It works the same as any other dispatch commands - work is replicated to all devices unless applications split the work themselves. There is no automatic scheduling between devices.

What dynamic state should be allowed?

Proposed: Support a subset of dynamic state.

For now, this specification exposes basic "value" state - primarily things where there is only a value to modify rather than a mode switch or state enable.