Web-разработка

Несложное WebGPU-программирование с использованием taichi.js

07.12.2024

Будучи фанатом компьютерной графики и языков программирования, могу похвастаться тем, что в последние два года работал над несколькими GPU-компиляторами.

Сначала в 2021 году внес личный вклад в разработку taichi — библиотеки Python, которая компилирует Python-функции в GPU-ядра на 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, как устройства, очереди команд, группы привязки и т. д. Достаточно создавать обычные JavaScript-функции, а компилятор преобразует эти функции в конвейеры вычислений или рендеринга WebGPU. Таким образом, любой программист, овладевший базовым синтаксисом JavaScript, может писать код WebGPU через taichi.js.

В этой статье будет продемонстрирована модель taichi.js-программирования на примере создания игры «Жизнь» (Game of Life). Вы увидите, как можно с помощью менее 100 строк разработать полностью распараллеленную WebGPU-программу, содержащую три конвейера GPU-вычислений и конвейер GPU-рендеринга. Полный исходный код демонстрации можно найти здесь, а, чтобы опробовать код, не устанавливая никаких локальных сред, перейдите на эту страницу.

Игра

Игра «Жизнь» является классическим примером клеточного автомата — системы клеток, которые развиваются со временем по простым правилам. Он был изобретен математиком Джоном Конвеем в 1970 году и с тех пор стал любимой игрой как компьютерщиков, так и математиков. Игра ведется на двухмерной сетке, где каждая клетка может быть живой или мертвой. Правила игры просты:

если у живой (living) клетки меньше двух или больше трех живых соседей, она умирает;

если у мертвой (dead) клетки есть ровно три живых соседа, она становится живой.

Несмотря на свою простоту, игра «Жизнь» может демонстрировать удивительное поведение. Начиная с любого случайного исходного состояния, игра часто сходится к состоянию, в котором доминируют несколько паттернов, как будто это «виды», выжившие в результате эволюции.

Симуляция

Приступим к реализации игры «Жизнь» с помощью taichi.js. Сначала импортируем библиотеку taichi.js (сокращенно — ti) и определим async-функцию main(), которая будет содержать всю логику. Внутри main() начнем с вызова ti.init(), который инициализирует библиотеку и ее WebGPU-контексты.

import * as ti from "path/to/taichi.js"

let main = async () => {
    await ti.init();
    ...
};

main()

После вызова ti.init() определим структуры данных, необходимые для симуляции игры «Жизнь»:

    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 «поле» — это n-мерный массив, размерность которого задается во втором аргументе ti.field(). Тип элемента массива задается в первом аргументе. В данном случае ti.i32 указывает на 32-битные целые числа. Однако элементами поля могут быть и более сложные типы, включая векторы, матрицы и структуры.

Строка кода ti.addToKernelScope({...}) гарантирует, что переменные N, liveness и numNeighbors будут видны в «ядрах» taichi.js, которые представляют собой конвейеры GPU-вычислений и/или GPU-рендеринга, определенные в виде JavaScript-функций. В качестве примера можно привести следующее ядро init , которое используется для заполнения клеток сетки начальными значениями liveness, где каждая клетка имеет 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 компилирует логику строкового представления этой JavaScript-лямбды в код WebGPU. Лямбда содержит цикл for, индекс которого I проводит итерацию через ti.ndrange(N, N). Это означает, что I будет принимать NxN различных значений в диапазоне от [0, 0] до [N-1, N-1].

И тут начинается волшебство: все верхнеуровневые for-циклы в ядре будут распараллелены в taichi.js. Точнее, для каждого возможного значения индекса цикла taichi.js выделяет один поток шейдера вычислений WebGPU для его выполнения. В данном случае выделим один GPU-поток для каждой клетки из симуляции игры «Жизнь», инициализируя ее в случайном состоянии liveness. Случайность обеспечивается функцией 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. (Для читателя, склонного к топологии, это означает, что игра имеет тороидальные граничные условия).

Подсчитав количество соседей для каждой клетки, обновляем их состояния живости (liveness) в ядре updateLiveness(). Это простое дело — считать состояние живости каждой клетки и текущее количество ее живых соседей и записать обратно новое значение живости в соответствии с правилами игры. Как обычно, данный процесс применяется ко всем клеткам параллельно.

На этом реализация логики симуляции игры завершена. Далее рассмотрим, как определить конвейер рендеринга WebGPU для отрисовки эволюции игры на веб-странице.

Рендеринг

Написание кода рендеринга в taichi.js несколько сложнее, чем написание кода для вычислительных ядер общего назначения, и требует некоторого понимания вершинных шейдеров, фрагментных шейдеров и конвейеров растеризации в целом. Однако благодаря простой модели taichi.js-программирования с этими концепциями довольно легко разобраться и работать.

Прежде чем рисовать что-либо, нужно получить доступ к холсту, на котором будем это делать. Предполагая, что в HTML существует холст с именем result_canvas, создадим с помощью следующих строк кода объект ti.CanvasTexture, представляющий собой кусок текстуры, на который можно наложить конвейер рендеринга taichi.js.

    let htmlCanvas = document.getElementById('result_canvas');
    htmlCanvas.width = 512;
    htmlCanvas.height = 512;
    let renderTarget = ti.canvasTexture(htmlCanvas);

На холсте отобразим квадрат и нарисуем 2D-сетку игры. В графических процессорах геометрия, подлежащая рендерингу, представляется в виде треугольников. В данном случае квадрат, который необходимо отрисовать, будет представлен в виде двух треугольников. Эти два треугольника определены в поле ti.field, которое хранит координаты каждой из шести вершин двух треугольников:

    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, чтобы они были видны в ядрах GPU в taichi.js:

    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]);
        }
    });

Далее определяем два верхнеуровневых цикла for-loops, которые, как вы уже знаете, являются циклами, распараллеливаемыми в WebGPU. Однако, в отличие от предыдущих циклов, в которых выполнялись итерации над объектами ti.ndrange, в этих циклах выполняем итерации над ti.inputVertices(vertices) и ti.inputFragments() соответственно. Это указывает на то, что данные циклы будут скомпилированы в «вершинные шейдеры» и «фрагментные шейдеры» WebGPU, которые работают вместе как конвейер рендеринга.

У вершинного шейдера есть две обязанности.

Для каждой вершины треугольника вычислить ее конечное местоположение на экране, точнее, координаты «Clip Space» (координаты усеченного пространства). В конвейере 3D-рендеринга это обычно включает совокупность матричных умножений, которые преобразуют координаты вершин модели в «World space» (координаты мирового пространства), затем в «Camera space» (координаты пространства камеры) и, наконец, в «Clip Space». Однако для простого 2D-квадрата входные координаты вершин уже имеют правильные значения в «Clip Space», поэтому можно избежать этих вычислений. Все, что нужно сделать, — добавить фиксированное значение z, равное 0,0, и фиксированное значение w, равное 1.0 (не волнуйтесь, если не понимаете, что такое z и w, — это не важно).

          ti.outputPosition([v.x, v.y, 0.0, 1.0]);

Для каждой вершины генерировать данные для интерполяции, которые затем передаются во фрагментный шейдер. В конвейере рендеринга встроенный процесс, известный как «Растеризация» («Rasterization»), выполняется для всех треугольников после выполнения вершинного шейдера. Это процесс с аппаратным ускорением, который вычисляет для каждого треугольника, какие пиксели покрывает этот треугольник. Эти пиксели также известны как «фрагменты».

Для каждого треугольника программист может сгенерировать дополнительные данные в каждой из трех вершин, которые будут интерполированы на этапе растеризации. Для каждого фрагмента пикселя соответствующий поток фрагментного шейдера будет получать интерполированные значения в соответствии с его местоположением в треугольнике. В нашем случае фрагментному шейдеру нужно знать только местоположение фрагмента в 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 — это интерполированное местоположение пикселя, переданное из вершинного шейдера. Используя это значение, фрагментный шейдер будет искать состояние живости клетки, покрывающей этот пиксель. Для этого нужно сначала преобразовать координаты пикселя f в диапазон [0, 0] ~ [1, 1] и записать их в переменную coord. Затем она умножается на размерность поля liveness, в результате чего получается индекс покрывающей клетки.

Наконец, получаем значение 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 завершена.

Если вы запустите программу, то увидите следующую анимацию, в которой 128×128 клеток эволюционируют в течение примерно 1400 поколений, прежде чем сходятся к нескольким видам стабилизированных организмов.

Упражнения

Надеюсь, эта демонстрация показалась вам интересной! Если да, предлагаю выполнить дополнительные упражнения, чтобы поэкспериментировать, и подумать над вопросами, чтобы углубиться в тему. (Кстати, чтобы быстро поэкспериментировать с кодом, перейдите на эту страницу).

[Облегченный уровень] Добавьте в демонстрацию счетчик FPS (счетчик кадров в секунду)! Какое значение FPS можно получить при текущих настройках, где N = 128? Попробуйте увеличить значение N и посмотрите, как изменится частота кадров. Можно ли написать ванильную JavaScript-программу, которая обеспечит такую частоту кадров без taichi.js или WebGPU?

[Средний уровень] Что произойдет, если объединить countNeighbors() и updateLiveness() в одно ядро и сохранить счетчик neighbors в качестве локальной переменной? Будет ли программа работать всегда корректно?

[Уровень повышенной сложности] В taichi.js ti.kernel(..) всегда производит async-функцию независимо от того, что содержит ядро — конвейеры вычислений или конвейеры рендеринга. Догадались ли вы, в чем смысл этой асинхронности? И в чем смысл вызова await на этих async-вызовах? Наконец, почему в функции frame, определенной выше, указано await только для функции render(), но не для двух других?

Последние два вопроса особенно интересны, поскольку затрагивают внутреннюю работу компилятора и времени выполнения фреймворка taichi.js, а также принципы GPU-программирования.

Ресурсы

Конечно, пример с игрой «Жизнь» — лишь малая часть того, что можно сделать с помощью taichi.js. Он предоставляет возможности поэкспериментировать с множеством других программ. Еще больше программ — от моделирования текучих сред в реальном времени до рендеринга на основе физики — вы можете написать сами с помощью taichi.js и следующих ресурсов: