Как специалист по компьютерной графике и языкам программирования, я рад, что за последние 2 года работал над несколькими компиляторами для графических процессоров.

Это началось в 2021 году, когда я начал вносить свой вклад в taichi, библиотеку Python, которая компилирует функции Python в ядра GPU в CUDA, Metal или Vulkan. Позже я присоединился к Meta и начал работать над SparkSL, языком шейдеров, который поддерживает кроссплатформенное программирование GPU для эффектов дополненной реальности в Instagram и Facebook. Помимо личного удовольствия, я всегда считал или, по крайней мере, надеялся, что эти фреймворки на самом деле весьма полезны — они делают программирование на GPU более доступным для неспециалистов, позволяя людям создавать увлекательное графическое содержимое без необходимости осваивать сложные концепции GPU.

В моем последнем выпуске компиляторов я обратил внимание на WebGPU — графический API следующего поколения для Интернета. WebGPU обещает обеспечить высокопроизводительную графику за счет низкой нагрузки на ЦП и явного управления графическим процессором, что соответствует тенденции, начатой ​​Vulkan и D3D12 около 7 лет назад. Как и в случае с Vulkan, преимущества WebGPU в производительности достигаются за счет крутой кривой обучения. Хотя я уверен, что это не остановит талантливых программистов по всему миру от создания потрясающего контента с помощью WebGPU, я хотел дать людям возможность играть с WebGPU, не сталкиваясь с его сложностью. Так появился taichi.js.

В модели программирования taichi.js программистам не нужно рассуждать о концепциях WebGPU, таких как устройства, очереди команд, группы привязки и т. д. Вместо этого они пишут простые функции Javascript, а компилятор переводит эти функции в конвейеры вычислений или рендеринга WebGPU. Это означает, что любой может написать код WebGPU через taichi.js, если он знаком с базовым синтаксисом Javascript.

Оставшаяся часть этой статьи продемонстрирует модель программирования taichi.js с помощью программы Игра жизни. Как вы увидите, используя менее 100 строк кода, мы создадим полностью параллельную программу WebGPU, содержащую 3 вычислительных конвейера GPU и конвейер рендеринга. Полный исходный код демо можно найти здесь, а если вы хотите поиграть с кодом без настройки каких-либо локальных окружений, перейдите на эту страницу.

Игра

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

  • если у живой клетки меньше двух или больше трех живых соседей, она умирает
  • если у мертвой клетки ровно три живых соседа, она становится живой.

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

Моделирование

Давайте погрузимся в реализацию Game of Life с использованием 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()

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

Следующая строка кода, ti.addToKernelScope({...}), обеспечивает видимость переменных N, liveness и numNeighbors в ядрах taichi.js, которые представляют собой конвейеры вычислений и/или рендеринга GPU, определенные в виде функций 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 просматривает строковое представление Javascript этой лямбды и компилирует его логику в код 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(). Это просто вопрос считывания состояния жизнеспособности каждой ячейки и ее текущего количества живых соседей, а также записи нового значения живучести в соответствии с правилами игры. Как обычно, этот процесс применяется ко всем ячейкам параллельно.

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

Вершинный шейдер выполняет две функции:

  • Для каждой вершины треугольника вычислите ее конечное положение на экране (или, точнее, ее координаты «пространства отсечения»). В конвейере 3D-рендеринга это обычно включает в себя множество матричных умножений, которые преобразуют координаты модели вершины в мировое пространство, а затем в пространство камеры и, наконец, в «пространство клипа». Однако для нашего простого 2D-квадрата входные координаты вершин уже имеют правильные значения в пространстве отсечения, так что всего этого можно избежать. Все, что нам нужно сделать, это добавить фиксированное значение z, равное 0,0, и фиксированное значение w, равное 1.0 (не волнуйтесь, если не знаете, что это такое — здесь это не важно!).
          ti.outputPosition([v.x, v.y, 0.0, 1.0]);
  • Для каждой вершины создайте данные, которые будут интерполированы, а затем переданы во фрагментный шейдер. В конвейере рендеринга после выполнения вершинного шейдера для всех треугольников выполняется встроенный процесс, известный как «Растеризация». Это процесс с аппаратным ускорением, который вычисляет для каждого треугольника, какие пиксели покрываются этим треугольником. Эти пиксели также известны как «фрагменты». Для каждого треугольника программисту разрешено генерировать дополнительные данные в каждой из 3-х вершин, которые будут интерполированы на этапе растеризации. Для каждого фрагмента в пикселе соответствующий поток фрагментного шейдера получит интерполированные значения в соответствии с его расположением в треугольнике. В нашем случае фрагментному шейдеру нужно знать только местоположение фрагмента внутри 2D-квадрата, чтобы он мог получить соответствующие значения живости игры. Для этого достаточно передать в растеризатор 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. Если вы запустите программу, вы должны увидеть следующую анимацию, в которой клетки размером 128x128 эволюционируют примерно в течение 1400 поколений, прежде чем слиться в несколько видов стабилизированных организмов.

Упражнения

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

  1. [Легко] Добавьте в демо счетчик FPS! Какое значение FPS вы можете получить с текущими настройками, где N = 128? Попробуйте увеличить значение N и посмотрите, как изменится частота кадров. Сможете ли вы написать ванильную программу на Javascript, которая получает эту частоту кадров без taichi.js или WebGPU?
  2. [Средний] Что произойдет, если мы объединим countNeighbors() и updateLiveness() в одно ядро ​​и сохраним счетчик neighbors в качестве локальной переменной? Будет ли программа всегда работать корректно?
  3. [Жесткий] В taichi.js ti.kernel(..) всегда создает функцию async, независимо от того, содержит ли он вычислительные конвейеры или конвейеры рендеринга. Если вам нужно угадать, в чем смысл этой async-ности? И какой смысл звонить await на эти звонки async? Наконец, почему в определенной выше функции frame мы поставили await только для функции render(), но не для двух других?

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

Ресурсы

Конечно, этот пример Game of Life лишь поверхностно показывает, что вы можете делать с taichi.js. От симуляции жидкости в реальном времени до физических рендереров — существует множество других taichi.js программ, с которыми вы можете играть, и еще больше программ, которые вы можете написать самостоятельно. Дополнительные примеры и учебные ресурсы см.

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