[Перевод] taichi.js: Программируем на WebGPU без боли24.05.2023 15:31

Привет, Хабр! Сегодня хочу предложить вашему вниманию перевод на русский язык статьи моего коллеги и хорошего приятеля 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 примет NxN разных значений в диапазоне от [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. Если вы запустите программу, вы должны увидеть анимацию, в которой 128×128 клеток эволюционируют примерно в течение 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), которые вы можете использовать, и более того, которые вы можете написать самостоятельно.

Дополнительные примеры и учебные ресурсы:

Удачного кодирования!