Матричное умножение
Возможность писать «слитые» (fused) ядра (ядра с обобщением этапов вычислений) для операций и сокращений по элементам важна, но недостаточна, учитывая, что в нейронных сетях задачи умножения матриц занимают важное место. Как оказалось, Triton отлично подходит и для них, достигая пиковой производительности менее чем в 25 строк Python. С другой стороны, реализация чего-то подобного в CUDA потребует гораздо больше усилий и даже, скорее всего, приведёт к снижению производительности
Умножение матриц при помощи Triton:
Одним из важных преимуществ рукописных ядер матричного умножения является то, что они могут быть настроены по желанию, чтобы разместить слитые преобразования их входных данных (например, слайсинга) и выходных (например, Leaky ReLU). Без Triton нетривиальные модификации ядер матричного умножения недоступны разработчикам без исключительного опыта программирования GPU.
V100 производительность тензорных ядер при умножении матриц с соответствующим образом настроенными значениями для BLOCK_MM, BLOCK_NN, BLOCK_KK, GROUP_MM.
Бэкенд компилятора
Мы обнаружили, что использование блочных представлений программ Triton-IR позволяет нашему компилятору автоматически выполнять широкий спектр важных оптимизаций.
Данные могут автоматически укладываться в общую память, просматривая операнды вычислительно интенсивных операций на уровне блоков (например, tl.dot), и выделяться/синхронизироваться с помощью стандартных методов анализа живучести.
Компилятор Triton выделяет общую память, анализируя живой диапазон блочных переменных, задействованных в вычислительно интенсивных операциях.
С другой стороны, программы Triton могут быть эффективно и автоматически распараллелены как (1) между SM — путём одновременного выполнения различных экземпляров ядра, так и (2) внутри SM — путём анализа пространства итераций каждой операции на уровне блоков и адекватного разбиения его на различные SIMD-блоки, как показано ниже.
Автоматическое распараллеливание в Triton. Каждая операция на уровне блоков определяет блочное пространство итераций, которое автоматически распараллеливается для использования ресурсов, доступных на потоковом мультипроцессоре (SM).
