如何使用 python 和 numba 在 RTX GPU 中对 NVIDIA 的张量核心进行编程？

Posted 2023-03-23

【中文标题】如何使用 python 和 numba 在 RTX GPU 中对 NVIDIA 的张量核心进行编程？

【英文标题】：How to program NVIDIA's tensor cores in RTX GPU with python and numba?

【发布时间】：2020-10-10 15:47:34

【问题描述】：

我有兴趣在 python 中使用来自 NVIDIA RTX GPU 的张量核心，以从它在某些科学计算中的加速中受益。 Numba 是一个很棒的库，它允许为 cuda 编程内核，但我还没有找到如何使用张量内核。可以用 Numba 完成吗？如果没有，我应该使用什么？

【参考方案1】：

....我还没有找到如何使用张量核心。可以用 Numba 完成吗？

没有。 Numba 目前在设备代码中没有半精度支持或张量核心内在函数。

如果没有，我应该使用什么？

我认为你会被困在用本机 C++ 方言编写内核代码，然后使用 PyCUDA 之类的东西来运行从该 C++ 方言编译的设备代码。

【讨论】：

我希望 Numba 将来支持这一点。那么我将不得不使用 pyCUDA。

@DiegoPalacios：在他们的 github 跟踪器上打开功能请求。除非你告诉他们，否则他们不会知道有兴趣