Привет, Хабр! Сегодня хочу предложить вашему вниманию перевод на русский язык статьи моего коллеги и хорошего приятеля Dunfan Lu. Он создал taichi.js
- open-source фреймворк для программирования графики на WebGPU, и написал подробный туториал о том, как его использовать на примере знаменитой "Игры жизни". Уверен, эта сложная и красивая работа на стыке технологий рендеринга и компиляции не оставит вас равнодушными. - пр. переводчика.
Я рад, что как специалисту по компьютерной графике и компиляторам, за последние 2 года мне удалось поработать над несколькими компиляторами для графических процессоров (GPU). В 2021 году я начал работать над taichi, библиотекой на языке Python, которая транслирует функции Python в ядра CUDA, Metal или Vulkan. Позже я присоединился к Meta, где начал работать над SparkSL, языком шейдеров, который обеспечивает кросс-платформенное программирование GPU для AR в Instagram и Facebook. Помимо личного удовольствия, которое я получил от этой работы, я всегда считал или, по меньшей мере, надеялся, что эти фреймворки-компиляторы будут полезны. Ведь они были спроектированы, чтобы сделать программирование на GPU более доступным для неспециалистов, позволяя людям создавать привлекательный графический контент без необходимости глубоко разбираться в сложных концепциях GPU.
Работая над моим последним компилятором, я обратил внимание на WebGPU — графический API следующего поколения для веба. WebGPU обещал обеспечить высокопроизводительную графику за счет снижения нагрузки на CPU и прямого доступа к GPU, что соответствовало тренду, начатому Vulkan и D3D12 около 7 лет назад. Как и в случае с Vulkan, для достижения улучшенной производительности в WebGPU, необходимо пройти длинную кривую обучения. Хотя я уверен, что это не остановит талантливых программистов по всему миру от создания потрясающего контента при помощи WebGPU, я хотел помочь людям преодолеть начальную сложность вхождения в WebGPU, предоставив песочницу для работы с ним. Так появился taichi.js
.
В модели программирования taichi.js
разработчикам не нужно разбираться в таких концепциях WebGPU, как устройства (devices), очереди команд (command queues) или группы привязки (bind groups). Вместо этого они пишут простые функции на JavaScript, а компилятор переводит эти функции в вычислительные конвейеры (compute pipelines) или конвейеры для рендеринга (render pipelines). Это означает, что любой, кто знаком с базовым синтаксисом JavaScript, может написать код WebGPU при помощи taichi.js
.
Далее в этой статье будет продемонстрирована модель программирования taichi.js
на примере программы «Игра жизни». Вы увидите, как используя менее 100 строк кода, можно создать полностью параллельную программу на WebGPU, содержащую 3 вычислительных конвейера GPU и 1 конвейер рендеринга. Полный исходный код примера можно найти здесь, а если вы хотите поэкспериментировать с кодом без необходимости настраивать окружение, можете сделать это здесь.
Игра
"Игра жизни" — классический пример клеточного автомата, системы клеток, которые со временем эволюционируют по простым правилам. Игра была изобретена математиком Джоном Конвеем в 1970 году и с тех пор не теряет популярности среди программистов и математиков. Игра происходит на двухмерной сетке, где каждая ячейка может быть либо живой, либо мертвой.
Правила игры просты:
если у живой клетки меньше двух или больше трех живых соседей, она умирает;
если у мертвой клетки ровно три живых соседа, она становится живой.
Несмотря на свою простоту, “Игра жизни” может демонстрировать удивительное поведение. Начиная с любого случайного начального состояния, игра часто сходится к состоянию, в котором доминируют несколько паттернов, которые напоминает виды (species), пережившие эволюцию.
Симуляция
Рассмотрим в реализацию "Игры жизни" с помощью taichi.js
. Для начала импортируем библиотеку taichi.js
с кратким именем ti
и определим асинхронную функцию main()
, которая будет содержать всю логику программы. В main()
мы начнем с вызова ti.init()
, который инициализирует библиотеку и контексты WebGPU внутри нее.
import * as ti from "path/to/taichi.js"
let main = async () => {
await ti.init();
...
};
main()
Затем определим структуры данных, необходимые для симуляции “Игры жизни”:
let N = 128;
let liveness = ti.field(ti.i32, [N, N])
let numNeighbors = ti.field(ti.i32, [N, N])
ti.addToKernelScope({ N, liveness, numNeighbors });
Здесь мы определили две переменные, liveness
и numNeighbors
, которые определяются как ti.field
. В taichi.js
«field» — это, по сути, n-мерный массив, размерность которого указывается во втором аргументе ti.field()
. Тип элемента массива определяется в первом аргументе. В данном случае это ti.i32
, обозначающий 32-битные целые числа. Однако элементы ti.field
могут быть и другими, более сложными типами, включая векторы, матрицы и даже структуры.
Следующая строка кода, ti.addToKernelScope({...})
, обеспечивает видимость переменных N
, liveness
и numNeighbors
в ядрах (kernels), которые являются вычислительными конвейерами и/или конвейерами рендеринга, определенными в форме функций JavaScript. В качестве примера, рассмотрим следующее ядро init
, которое используется для заполнения клеток в сетке начальными значениями живучести, где каждая клетка изначально имеет 20%-й шанс быть живой:
let init = ti.kernel(() => {
for (let I of ti.ndrange(N, N)) {
liveness[I] = 0
let f = ti.random()
if (f < 0.2) {
liveness[I] = 1
}
}
})
init()
Ядро init()
создается при помощи вызова ti.kernel()
с лямбда-функцией JavaScript в качестве аргумента. Под капотом taichi.js
просматривает строковое представление этой лямбда-функции и компилирует ее логику в код WebGPU. Здесь лямбда-функция содержит цикл for
, индекс I
которого проходит через ti.ndrange(N, N)
. Это означает, что I
примет N
xN
разных значений в диапазоне от [0, 0]
до [N-1, N-1]
.
А дальше начинается магия — в taichi.js
все циклы for
верхнего уровня в ядре будут распараллелены. В частности, для каждого возможного значения индекса цикла taichi.js
выделит одну нить вычислительного шейдера WebGPU. В нашем случае мы выделяем по одной нити для каждой ячейки в симуляции “Игры жизни”, инициализируя ее случайным образом. Случайность обеспечивается функцией ti.random()
, одной из многих функций, предоставляемых в библиотеке taichi.js для использования в ядре. Полный список этих функций доступен в документации taichi.js.
Определив начальное состояние игры, перейдем к эволюции. Следующие два ядра taichi.js
, определяют ее так:
let countNeighbors = ti.kernel(() => {
for (let I of ti.ndrange(N, N)) {
let neighbors = 0
for (let delta of ti.ndrange(3, 3)) {
let J = (I + delta - 1) % N
if ((J.x != I.x || J.y != I.y) && liveness[J] == 1) {
neighbors = neighbors + 1;
}
}
numNeighbors[I] = neighbors
}
});
let updateLiveness = ti.kernel(() => {
for (let I of ti.ndrange(N, N)) {
let neighbors = numNeighbors[I]
if (liveness[I] == 1) {
if (neighbors < 2 || neighbors > 3) {
liveness[I] = 0;
}
}
else {
if (neighbors == 3) {
liveness[I] = 1;
}
}
}
})
Как и ядро init()
, которое мы рассматривали ранее, эти два ядра также имеют циклы for
верхнего уровня, перебирающие ячейки сетки, и которые распараллеливаются компилятором. В countNeighbors()
для каждой клетки мы рассматриваем 8 соседних клеток и подсчитываем, сколько из этих соседей «живы». Количество живых соседей хранится в поле numNeighbors
. Обратите внимание, что при переборе соседей цикл for (let delta of ti.ndrange(3, 3)) {...}
не распараллеливается, поскольку это не цикл верхнего уровня. Индекса цикла delta
находится в диапазоне от [0, 0]
до [2, 2]
и используется для смещения исходного индекса ячейки I
. Мы не выходим за границы массивов при помощи деления по модулю N
(для читателей, которые любят топологические модели: это означает, что игра имеет тороидальные граничные условия).
Подсчитав количество соседей для каждой ячейки, мы переходим к обновлению их состояния живучести в ядре updateLiveness()
. Мы просто считываем значения liveness
и текущего количества живых соседей для каждой ячейки и записываем новое значение живучести в соответствии с правилами игры. Как и ранее, этот процесс применяется ко всем ячейкам параллельно.
На этом, по сути, завершается реализация логики симуляции в игре. Теперь рассмотрим, как определить конвейер рендеринга WebGPU для отображения игры на веб-странице.
Рендеринг
Написание кода рендеринга в taichi.js
чуть сложнее, чем написание вычислительных ядер общего назначения, и требует некоторого понимания вершинных шейдеров, фрагментных шейдеров и конвейера растеризации в целом. Однако, простая модель программирования taichi.js
делает эти концепции чрезвычайно простыми для работы и анализа.
Прежде чем что-либо рисовать, нам нужен доступ к канвасу. Предполагая, что канвас с именем result_canvas
существует в HTML, следующие строки кода создадут объект ti.CanvasTexture
, представляющий собой текстуру, в которую мы будем рисовать при помощи taichi.js
.
let htmlCanvas = document.getElementById('result_canvas');
htmlCanvas.width = 512;
htmlCanvas.height = 512;
let renderTarget = ti.canvasTexture(htmlCanvas);
На нашем канвасе мы нарисуем квадрат и 2D-сетку игры в этом квадрате. В GPU геометрия для рендеринга представлена в виде треугольников. Так квадрат, который мы пытаемся визуализировать, будет представлен в виде двух треугольников. Эти два треугольника определены в поле ti.field
, в котором хранятся координаты 6 вершин:
let vertices = ti.field(ti.types.vector(ti.f32, 2), [6]);
await vertices.fromArray([
[-1, -1],
[1, -1],
[-1, 1],
[1, -1],
[1, 1],
[-1, 1],
]);
Как и в случае с полями liveness
и numNeighbors
, нам необходимо явно объявить, что переменные renderTarget
и vertices
должны быть видны в ядрах:
ti.addToKernelScope({ vertices, renderTarget });
Теперь у нас есть все данные, необходимые для реализации конвейера рендеринга. Вот реализация самого конвейера:
let render = ti.kernel(() => {
ti.clearColor(renderTarget, [0.0, 0.0, 0.0, 1.0]);
for (let v of ti.inputVertices(vertices)) {
ti.outputPosition([v.x, v.y, 0.0, 1.0]);
ti.outputVertex(v);
}
for (let f of ti.inputFragments()) {
let coord = (f + 1) / 2.0;
let texelIndex = ti.i32(coord * (liveness.dimensions - 1));
let live = ti.f32(liveness[texelIndex]);
ti.outputColor(renderTarget, [live, live, live, 1.0]);
}
});
Внутри ядра render()
мы начинаем с очистки renderTarget
полностью черным цветом, представленным в RGBA как [0.0, 0.0, 0.0, 1.0]
.
Далее мы определяем два цикла for верхнего уровня, которые, как вы уже знаете, будут распараллеленны в WebGPU. Однако, в отличие от предыдущих циклов, в которых мы перебирали объекты при помощи ti.ndrange
, эти циклы перебирают ti.inputVertices(vertices)
и ti.inputFragments()
соответственно. Это означает, что эти циклы будут скомпилированы в вершинные и фрагментные шейдеры WebGPU, которые работая вместе образуют конвейер рендеринга.
Вершинный шейдер выполняет две функции:
Для каждой вершины треугольника он вычисляет ее конечное положение на экране (или, точнее, ее координаты в пространстве отсечения, clip space). В конвейере 3D-рендеринга это обычно включает в себя несколько матричных умножений, которые преобразуют координаты вершины в мировое пространство, затем в пространство камеры и, наконец, в пространстве отсечения. Однако, для нашего простого 2D-квадрата входные координаты вершин уже имеют правильные значения в пространстве отсечения, так что умножения матриц можно избежать. Все, что нам необходимо сделать, это определить значение z равным 0.0, а w равным 1.0 (не волнуйтесь, если не знаете, что это такое — здесь это не важно!).
ti.outputPosition([v.x, v.y, 0.0, 1.0]);
Для каждой вершины он формирует данные, которые будут интерполированы, а затем переданы во фрагментный шейдер. В конвейере рендеринга после выполнения вершинного шейдера для всех треугольников выполняется встроенный процесс, известный как растеризация. Растеризация - это аппаратно-ускоренный процесс, который вычисляет для каждого треугольника, какие пиксели покрываются этим треугольником. Эти пиксели также известны как «фрагменты». Для каждого треугольника можно генерировать дополнительные данные в каждой из 3-х вершин. Эти данные будут интерполированы на этапе растеризации. Для каждого фрагмента соответствующая нить фрагментного шейдера получит интерполированные значения в соответствии с расположением фрагмента в треугольнике (подробнее про растеризацию и интерполяцию можно почитать, например, здесь - пр. переводчика).
В нашем случае фрагментному шейдеру нужно знать только местоположение фрагмента внутри 2D-квадрата, чтобы он мог получить соответствующие значения
liveness
. Для этого достаточно передать в растеризатор 2D-координату вершины, и фрагментный шейдер получит интерполированное 2D-местоположение самого пикселя:ti.outputVertex(v);
Перейдем к фрагментному шейдеру:
for (let f of ti.inputFragments()) {
let coord = (f + 1) / 2.0;
let cellIndex = ti.i32(coord * (liveness.dimensions - 1));
let live = ti.f32(liveness[cellIndex]);
ti.outputColor(renderTarget, [live, live, live, 1.0]);
}
Значение f
— это интерполированное местоположение пикселя, полученное из вершинного шейдера. Используя это значение, фрагментный шейдер будет получать значение liveness
для этого пикселя. Это делается путем преобразования координат пикселя f
в диапазон [0, 0] ~ [1, 1]
и сохранения этой координаты в переменной coord
. Затем происходит умножение на размер поля liveness
, что дает индекс ячейки cellIndex
. Наконец, мы получаем признак live
для этой клетки, который равен 0
, если клетка мертва, и 1
, если она жива. В конце мы выводим значение RGBA этого пикселя в renderTarget
, где все компоненты R, G, B равны live
, а компонент A равен 1
для полной непрозрачности.
Когда конвейер рендеринга определен, все, что осталось — это собрать все воедино, вызывая ядра симуляции и конвейер рендеринга в каждом кадре:
async function frame() {
countNeighbors()
updateLiveness()
await render();
requestAnimationFrame(frame);
}
await frame();
Вот и все! Мы завершили WebGPU реализацию “Игры жизни“ с использованием taichi.js
. Если вы запустите программу, вы должны увидеть анимацию, в которой 128x128 клеток эволюционируют примерно в течение 1400 поколений, прежде чем слиться в несколько видов стабилизированных организмов.
Упражнения
Надеюсь, демо вам понравилось! Если да, то у меня есть для вас несколько дополнительных упражнений и вопросов, над которыми я предлагаю вам подумать (кстати, быстро поэкспериментировать с кодом можно здесь).
[Легко] Добавьте в демо счетчик FPS! Какое значение FPS вы можете получить с текущими настройками, где
N = 128
? Попробуйте увеличить значениеN
и посмотрите, как изменится частота кадров. Смогли бы вы написать программу на чистом JavaScript, которая выдают такую же частоту кадров безtaichi.js
или без WebGPU?[Средне] Что произойдет, если мы объединим
countNeighbors()
иupdateLiveness()
в одно ядро и сохраним счетчик соседейneighbors
как локальную переменную? Будет ли программа всегда работать корректно?[Сложно] В
taichi.js
ti.kernel(..)
всегда создает асинхронную функцию, независимо от того, содержит ли она вычислительные конвейеры или конвейеры рендеринга. Попробуйте догадаться, в чем смысл использованияasync
здесь? И в чем смысл вызоваawait
для этих асинхронных вызовов? Наконец, почему в функцииframe
, определенной выше, мы поместилиawait
только для функцииrender()
, но не для двух других?
Последние 2 вопроса особенно интересны, так как они касаются не только внутреннего устройства компилятора и среды выполнения фреймворка taichi.js, но и принципов программирования для GPU. Дайте мне знать ваш ответ!
Ресурсы
Конечно, этот пример с “Игрой жизни” лишь поверхностно показывает, что можно сделать при помощи taichi.js
. Существуют другие программы на taichi.js
, от моделирования жидкости в реальном времени до рендеринга на основе законов физики (physically based renderers), которые вы можете использовать, и более того, которые вы можете написать самостоятельно.
Дополнительные примеры и учебные ресурсы:
Удачного кодирования!