如何通过 Vulkan 使用 Nvidia Tensor Cores
Posted
技术标签:
【中文标题】如何通过 Vulkan 使用 Nvidia Tensor Cores【英文标题】:How to use Nvidia's Tensor Cores via Vulkan 【发布时间】:2019-06-16 21:44:50 【问题描述】:如何使用 Vulkan 来利用 Nvidia 的张量核心(在计算着色器中?!)?
Nvidia Programming Tensor Cores in CUDA 9 有这篇文章,但这显然是针对 CUDA。我对 CUDA 不太熟悉,但看起来必须采取一些措施来启用 Tensor 核心上的计算,例如必须将算法设置为某种特殊类型,并且必须将某些数学类型设置为值 CUDNN_TENSOR_OP_MATH
。我想知道,是否也可以从其他 API 使用 Tensor 核心加速,我对 Vulkan 特别感兴趣。
更具体地说,我想深入研究用于降噪的过滤器。据我了解,过滤器主要需要 Tensor 核心能够加速的数学运算,即矩阵乘法和累加运算。
【问题讨论】:
实际上,NVIDIA 是应该回答这个问题的。我们没有必要推测。 IMO 提出这个问题的最佳地点是github.com/KhronosGroup/Vulkan-Ecosystem/issues。 github.com/KhronosGroup/Vulkan-Docs/issues/686 有一些切题的讨论,但我建议不要进一步污染该问题。 【参考方案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 来介绍如何使用这些张量核心。
【讨论】:
英伟达发表了一篇题为“Machine Learning Acceleration in Vulkan with Cooperative Matrices”的博客文章,所以是的,这似乎是要走的路。感谢您提供信息和链接!【参考方案2】:张量核心是一个小众功能,可能不会使其成为 Vulkan 扩展。 您仍然可以使用 CUDA 进行张量核心加速计算,并在 CUDA 和 Vulkan 上下文之间共享数据。
查看此示例:cuda vulkan interop
请注意,由于在启动 CUDA 内核和使用 Vulkan 端的结果之间需要同步,因此性能可能会受到影响。您必须评估申请中的成本。
【讨论】:
Nvidia 已经通过 SPIR-V 和 Vulkan 添加了支持,所以这是完全错误的khronos.org/registry/vulkan/specs/1.1-extensions/html/…以上是关于如何通过 Vulkan 使用 Nvidia Tensor Cores的主要内容,如果未能解决你的问题,请参考以下文章
至于 vulkan,GeForce RTX 3090 的 GPU 驱动程序不起作用