如何通过Vulkan使用Nvidia的Tensor核心

Question

如何使用Vulkan利用Nvidia的张量核心（在计算着色器中？！）？

有Nvidia Programming Tensor Cores in CUDA 9撰写的这篇文章，但这显然是针对CUDA的。我对CUDA不太熟悉，但看起来必须采取一些措施才能在Tensor内核上启用计算，例如必须将算法设置为某种特殊类型，并且将某些数学类型设置为值{{1} }。我想知道，是否还可以在其他API中使用Tensor核心加速功能，而我对Vulkan尤其感兴趣。

更具体地说，我想深入研究一下去噪更多的过滤器。据我了解，过滤器主要需要的是Tensor核心能够加速的数学运算，即矩阵乘法和累加运算。

Answer 1

Nvidia最近添加了一些新扩展，其中一个扩展名为VK_NV_COOPERATIVE_MATRIX，这将允许在Vulkan内部使用张量核心。

我相信是在昨天{strong>昨天添加了glslang处理此新功能的功能，这就是为什么到目前为止(see here)还没有看到此功能的原因：

以下是一些使用示例：

https://github.com/KhronosGroup/glslang/blob/4605e2ed2b2b1acbe157d365c3c528367b8b168f/Test/spv.coopmat.comp

https://github.com/KhronosGroup/glslang/blob/4605e2ed2b2b1acbe157d365c3c528367b8b168f/Test/spv.1.3.coopmat.comp

#version 450 core
#extension GL_KHR_memory_scope_semantics : enable
#extension GL_NV_cooperative_matrix : enable
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : enable

#pragma use_variable_pointers

layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

layout(set = 0, binding = 0) coherent buffer Block {
    float y[1024*1024];
    float x[];
} block;


void main()
{
    fcoopmatNV<32, gl_ScopeSubgroup, 16, 8> m = fcoopmatNV<32, gl_ScopeSubgroup, 16, 8>(0.0);

    m = m + m;
    m = m - m;
    m = -m;
    m = 2.0*m;
    m = m*2.0;

    coopMatLoadNV(m, block.x, 16, 128, false);
    coopMatStoreNV(m, block.x, 16, 128, false);
}

这似乎与在CUDA中的完成方式非常相似，要求将显式的内存传输到张量核可以运行的内存中。

因此要使用它们，您需要在vulkan中使用VK_NV_COOPERATIVE_MATRIX和在glsl中使用GL_NV_COOPERATIVE_MATRIX。

编辑：

j00hi提到现在有一​​个nvidia blog post关于如何使用这些张量核心的信息。

Answer 2

Tensor核心是一个利基功能，可能无法使其成为Vulkan扩展。 您仍然可以使用CUDA进行张量核心加速计算，并在CUDA和Vulkan上下文之间共享数据。

检查此示例：cuda vulkan interop

请注意，由于在启动CUDA内核和处理来自Vulkan端的结果之间需要同步，因此性能可能会受到影响。您将必须评估应用程序中的成本。