在 HLSL 中解锁 GPU 内部架构

这篇文章最初发表在 NVIDIA 技术博客上。

NVIDIA GPU 指令集中存在一些标准图形 API 中不包含的有用内部函数。

更新自 2016 年原始博文，添加了有关 DirectX 和 Vulkan 中新的内部结构和跨供应商 API 的信息。

例如，着色器可以使用线程束 shuffle 指令在线程束中的线程之间交换数据，而无需通过共享内存，这在没有共享内存的像素着色器中尤其重要。或者，着色器可以在全局内存中对半精度浮点数执行原子添加。

我们的文章线程之间的读取：着色器内部函数向您展示了内部指令的工作原理。现在，我将带您深入了解如何让它们在 DirectX 中运行。

在标准 DirectX 或 OpenGL 中，所有这些内部结构都不可能实现。[2023 年：这不再是事实。更多信息将在本文稍后分享。]但它们在 CUDA 中得到了多年的支持和详细记录。在 DirectX 中支持它们的机制已经推出一段时间，但没有得到广泛的记录。我的系统恰好从 2014 年 10 月开始就有旧的 NVAPI 版本 343，该版本(可能是更早的版本)在 DirectX 中支持内部函数。本文介绍了在 DirectX 中使用它们的机制。

遗憾的是，与 OpenGL 或 Vulkan 不同，DirectX 没有针对特定供应商的扩展程序的原生机制。但是，仍然可以通过自定义内部函数在 DirectX 11 或 12 中使用所有这些功能。这种机制在图形驱动程序中实现，并可通过 NVAPI 库来访问。

扩展 HLSL 着色器

要使用内部函数，必须将其编码为常规 HLSL 指令的特殊序列，以便驱动识别并转换为预期操作。这些特殊序列在 NVAPI SDK 随附的其中一个头文件中提供：nvHLSLExtns.h.

这些指令序列的一个重要方面是，它们必须在不进行优化的情况下通过 HLSL 编译器，因为编译器不理解它们的真正含义，因此可以修改它们，改变它们的顺序，甚至完全删除它们。

为了防止编译器这样做，序列在 UAV 缓冲区上使用原子操作。HLSL 编译器无法优化这些指令，因为它不知道可能的依赖项，即使没有依赖项。UAV 缓冲区基本上是假的，在通过 NVIDIA GPU 驱动程序后，实际着色器不会使用它。但应用程序仍然必须为其分配 UAV 插槽，并告诉驱动程序哪个插槽。

例如，NvShfl实现 Warp shuffle 的函数类似于以下代码示例，nvHLSLExtns.h:

ini 复制代码

int NvShfl(int val, uint srcLane, int width = NV_WARP_SIZE)
{
     uint index = g_NvidiaExt.IncrementCounter();
     g_NvidiaExt[index].src0u.x  =  val;          // variable to be shuffled
     g_NvidiaExt[index].src0u.y  =  srcLane;      // source lane
     g_NvidiaExt[index].src0u.z  =  __NvGetShflMaskFromWidth(width);
     g_NvidiaExt[index].opcode   =  NV_EXTN_OP_SHFL;
	    
// result is returned as the return value of IncrementCounter on fake UAV slot
     return g_NvidiaExt.IncrementCounter();
}

使用此函数的着色器类似于以下代码示例：

less 复制代码

// Declare that the driver should use UAV 0 to encode the instruction sequences.
// It's a pixel shader with one output, so u0 is taken by the render target - use u1.
#define NV_SHADER_EXTN_SLOT u1

// On DirectX12 and Shader Model 5.1, you can also define the register space for that UAV.
#define NV_SHADER_EXTN_REGISTER_SPACE space0

// Include the header - note that the UAV slot has to be declared before including it.
#include "nvHLSLExtns.h"

Texture2D tex : register(t0);
SamplerState samp : register(s0);

float4 main(in float2 texCoord : UV) : SV_Target
{
     float4 color = tex.Sample(samp, texCoord);

     // Use NvShfl to distribute the color from lane 0 to all other lanes in the warp.
     // The NvShfl function accepts and returns uint data, so use asuint/asfloat to pass float values.
	color.r = asfloat(NvShfl(asuint(color.r), 0));
	color.g = asfloat(NvShfl(asuint(color.g), 0));
	color.b = asfloat(NvShfl(asuint(color.b), 0));
	color.a = asfloat(NvShfl(asuint(color.a), 0));

	return color;
}

这个示例看起来可能是在做一些毫无意义的事情，而且确实如此。图形应用程序中内部函数的真实用例通常很复杂。例如，Warp shuffle 可用于优化算法(如光线消除)中的内存访问。VXGI 中使用浮点原子来在体素化期间累加发射。但是，这些应用程序需要大量着色器和主机代码才能正常工作。另一方面，这个示例几乎可以插入任何像素着色器，效果很明显。

编译此着色器时，每次调用NvShfl扩展到此序列中，指定或获取寄存器名称：

scss 复制代码

imm_atomic_alloc r1.x, u1
mov r3.yz, l(0,0,31,0)
mov r3.x, r2.z
store_structured u1.xyz, r1.x, l(76), r3.xyzx
store_structured u1.x, r1.x, l(0), l(1)
imm_atomic_alloc r0.y, u1

当此着色器通过驱动程序的 JIT 编译器时，NvShfl函数映射到一个 GPU 指令：

ini 复制代码

SHFL.IDX        PT, R3, R3, RZ, 0x1f;

在 DirectX 11 中创建扩展着色器

要实际使用此着色器，必须以特殊方式创建其运行时对象。定期调用ID3D11Device::CreatePixelShader这还不够，因为驱动程序必须知道着色器打算使用内部函数。它还必须知道使用哪个 UAV 插槽。

如果您使用的是 DirectX 11，请使用NvAPI_D3D11_SetNvShaderExtnSlot函数调用之前和之后CreatePixelShader:

scss 复制代码

// Do this one time during app initialization.
NvAPI_Initialize();

ID3D11PixelShader* pShader = nullptr;
HRESULT D3DResult = E_FAIL;

// First, enable compilation of intrinsics. 
// The second parameter is the UAV slot index that is used in the shader: u1.
NvAPI_Status NvapiStatus = NvAPI_D3D11_SetNvShaderExtnSlot(pDevice, 1);
if(NvapiStatus == NVAPI_OK)
{
     // Then create the shader as usual...
     D3DResult = pDevice->CreatePixelShader(pBytecode, BytecodeLength, nullptr, &pShader);

     // And disable again by telling the driver to use an invalid UAV slot.
     NvAPI_D3D11_SetNvShaderExtnSlot(pDevice, ~0u);
}

if(FAILED(D3DResult))
{
     // ...Handle the error...
}

此方法适用于任何可以引用 UAV 的着色器。因此，在 DirectX 11.0 中，它适用于像素和计算着色器。在 DirectX 11.1 及更高版本中，它应该适用于各种着色器。

在 DirectX 12 中创建扩展的工作流状态对象

如果您使用的是 DirectX 12，则不存在单独的着色器对象，而是创建完整的工作流状态 (PSO).

还有其他各种特定于 NVIDIA 的工作流状态扩展程序可通过 NVAPI 访问，因此为了避免使用各种扩展程序创建 PSO 的功能组合爆炸， NVIDIA 仅制作了两个功能，一个用于图形，另一个用于计算，可接受使用的扩展程序列表：

NvAPI_D3D12_CreateGraphicsPipelineState
NvAPI_D3D12_CreateComputePipelineState

HLSL 扩展由NVAPI_D3D12_PSO_SET_SHADER_EXTENSION_SLOT_DESC结构。不过，整个工作流状态只有一个，因此，如果工作流中的两个或多个着色器使用内部函数，它们必须为其使用相同的 UAV 插槽。

scss 复制代码

// Do this one time during app initialization.
NvAPI_Initialize();

// Fill the PSO description structure
D3D12_GRAPHICS_PIPELINE_STATE_DESC PsoDesc;
PsoDesc.VS = { pVSBytecode, VSBytecodeLength };
// ...And so on, as usual...

// Also fill the extension structure. 
// Use the same UAV slot index and register space that are declared in the shader.
NVAPI_D3D12_PSO_SET_SHADER_EXTENSION_SLOT_DESC ExtensionDesc;       
ExtensionDesc.baseVersion = NV_PSO_EXTENSION_DESC_VER;
ExtensionDesc.psoExtension = NV_PSO_SET_SHADER_EXTNENSION_SLOT_AND_SPACE;
ExtensionDesc.version = NV_SET_SHADER_EXTENSION_SLOT_DESC_VER;
ExtensionDesc.uavSlot = 1;
ExtensionDesc.registerSpace = 0;

// Put the pointer to the extension into an array. There can be multiple extensions enabled at one time.
// Other supported extensions are: 
       //     - Extended rasterizer state
       //  - Pass-through geometry shader, implicit or explicit
       //  - Depth bound test
       const NVAPI_D3D12_PSO_EXTENSION_DESC* pExtensions[] = { &ExtensionDesc };

// Now create the PSO.
ID3D12PipelineState* pPSO = nullptr;
NvAPI_Status NvapiStatus = NvAPI_D3D12_CreateGraphicsPipelineState(pDevice, &PsoDesc, ARRAYSIZE(pExtensions), pExtensions, &pPSO);

if(NvapiStatus != NVAPI_OK)
     {
        // ...Handle the error...
     }
}

查询 GPU 功能支持

最后，在尝试使用内部函数之前，您可能想知道应用所用的设备是否实际上支持这些内部函数。有两个 NVAPI 函数可以告诉您：

NvAPI_D3D11_IsNvShaderExtnOpCodeSupported
NvAPI_D3D12_IsNvShaderExtnOpCodeSupported

我们opCodeparameter 标识您感兴趣的特定操作。操作代码在nvShaderExtnEnums.hNVAPI SDK 随附的文件。例如，要测试 DirectX 11 设备是否支持 Warp shuffle，请使用以下代码示例：

ini 复制代码

#include "nvShaderExtnEnums.h"

bool bSupported = false;
NvAPI_Status NvapiStatus = NvAPI_D3D11_IsNvShaderExtnOpCodeSupported(pDevice, NV_EXTN_OP_SHFL, &bSupported);

if(NvapiStatus == NVAPI_OK && bSupported)
{
     // Yay, the device is no older than 2012!
}

2023 年更新：新的内部函数和跨供应商 API

NVIDIA GPU 支持的内部函数并不仅限于线程束 shuffle。事实上，线程束 shuffle 和相关函数现在可以通过 DirectX 12 和 Vulkan 中的跨供应商内部函数获得，因此无需使用 NVAPI。有关 DirectX 12 波内部函数的更多信息，请参阅Wave 内部函数。有关 Vulkan 子组操作的更多信息，请参阅Vulkan 子组教程。

NVIDIA GPU 支持的内部函数的完整列表可在名为 nvHLSLExtns.h 的文件中找到，现已在 GitHub 上提供。此文件中声明的函数可细分为几个通用类别：

较旧的线程束运算：shuffle、vote、ballot、通道索引 (NvShfl*, NvAny, NvAll, NvBallot, NvGetLaneId)
更新的线程束运算：波形匹配 (NvWaveMatch). NvWaveMatch返回线程束中活动通道的遮罩，这些通道的参数值与当前通道相同。
特殊寄存器访问权限(NvGetSpecial)
FP16、FP32 和 Uint64 变量上的扩展原子运算 (NvInterlocked*)
可变速率着色(NvGetShadingRate, NvEvaluateAttribute*)
纹理足迹评估(NvFootprint*)
WaveMultiPrefix 函数(NvWaveMultiPrefix*这些函数只是基于其他内部函数构建的算法。
光线追踪微图扩展程序(NvRtMicroTriangle*, NvRtMicroVertex*)
光线追踪着色器执行重排序(NvHitObject, NvReorderThread)

更新：编译具有正确选项的着色器

目前， NVIDIA GPU 驱动存在一个影响 HLSL 内部函数的已知问题。具体来说，如果着色器使用D3DCOMPILE_SKIP_OPTIMIZATION标志或/Od传递给 FXC 的命令行选项。如果您看到内部函数不起作用，请确保未指定此标志。

结束语

有关 NVAPI 函数和结构的更多信息，请参阅 NVAPI 头文件中的注释。有关更多用例和内部函数示例，请参阅以下资源：

线程间读取：着色器内置函数
NVIDIA NVAPI SDK：现在可在 GitHub /NVIDIA/nvapi 库中找到！
在 Kepler 架构上实现更快的并行归约
CUDA 专业技巧：使用 Kepler Shuffle
CUDA Pro 提示：利用 Warp 聚合原子操作进行优化过滤
Shuffle：技巧与窍门（GTC 会议）

阅读原文