Как запрограммировать тензорные ядра NVIDIA в графическом процессоре RTX с python и numba?

Question

Мне интересно использовать тензорные ядра графических процессоров NVIDIA RTX в python, чтобы получить выгоду от его ускорения в некоторых научных c вычислениях. Numba - отличная библиотека, которая позволяет программировать ядра для cuda, но я не нашел, как использовать тензорные ядра. Можно ли это сделать с помощью Numba? Если нет, что мне использовать?

Answer 1

.... Не нашел как использовать тензорные ядра. Можно ли это сделать с помощью Numba?

Нет. В настоящее время Numba не имеет поддержки половинной точности или встроенных функций тензорного ядра c, доступных в коде устройства.

Если нет, что мне следует использовать?

Я думаю вам придется писать код ядра на собственном диалекте C ++, а затем использовать что-то вроде PyCUDA для запуска кода устройства, скомпилированного на этом диалекте C ++.