nvidia - 如何通过 Vulkan 使用 Nvidia 的 Tensor Core

Question

如何使用 Vulkan 来利用 Nvidia 的张量核心(在计算着色器中？!)？

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

更具体地说，我想进一步深入研究降噪滤波器。据我了解，过滤器主要需要张量核心能够加速的数学运算，即矩阵乘法和累加运算。

Answer 1

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

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关于如何使用这些张量核心。