Пишем Surfgrad, высокопроизводительную библиотеку для автоматического дифференцирования выражений при помощи WebGPU.

Я работаю в компании Nomic, и многие из моих коллег заняты созданием больших TSNE-подобных визуализаций, работающих в браузере. При визуализации таких двумерных карт возникает две проблемы: проецировать эти конструкции (напр. TSNE и UMAP) в 2D-координатную систему протекает медленно и требует больших затрат оперативной памяти, особенно по мере того, как вы увеличиваете датасет и пытаетесь визуализировать в браузере миллионы точек данных, не расплавив при этом ноутбук невзначай.

Отобразить в браузере миллионы точек данных, не расплавив компьютер — та ещё задача. Мне доводилось слышать, что многие проблемы с масштабированием удаётся решать при помощи инструмента Deepscatter, разработанного Беном Шмидтом.

Но многие из таких разговоров, которые мне известны, вертятся вокруг Typescript и великолепия WebGPU как такового. Готовя эту статью, я не смог найти ни одной библиотеки для автоматического дифференцирования выражений, которая была бы написана с применением WebGPU. Но было бы упущением не назвать здесь два репозитория с функционально схожим наполнением: webGPT (библиотека на основе трансформеров, приспособлена только для логического вывода) и webgpu-blas (ядра для быстрого перемножения матриц под webGPU). Поэтому в качестве самообразования и, желая получше изучить WebGPU и Typescript, я решил написать Surfgrad, высокопроизводительную библиотеку для автоматического дифференцирования выражений под управлением WebGPU. Она обеспечивает тензорные операции в браузере. Как понятно по названию и по принципу работы, она во многом сделана по примеру tinygrad и micrograd.

В этом посте я расскажу, как мне удалось оптимизировать ядро WebGPU для простого перемножения матриц до достижения арифметической интенсивности более чем в 1ТФЛОПС. В рамках этого проекта я не пытался создать максимально быструю библиотеку для автоматического дифференцирования выражений, а хотел продемонстрировать нюансы WebGPU, и в чём этот инструмент может отличаться от CUDA.

Пожалуй, в будущем при помощи Surfgrad даже можно будет запускать свежие модели Llama.

Что такое WebGPU

WebGPU — это API, специально разработанный для всех желающих писать код для выполнения на GPU, который действовал бы на любом компьютере или смартфоне, где установлен браузер. Ранее требовалось колдовать над WebGL, чтобы выполнить нагрузки для машинного обучения, необходимые, например, для  отображения невидимого холста и считывания чисел как цветовых значений. Теперь, когда GPU на ноутбуках стали мощнее, на них можно задействовать вычислительные ядра (то есть, подавать данные на вход и получать данные на выходе без всяких дополнительных фокусов). Так, сообщается, что мощность M3 Pro от Apple  ~7ТФЛОПС. На этом устройстве даже можно запускать Llama3.2 (с ONNX) в браузере и обрабатывать 85 токенов в секунду.

WebGPU создавался с целью перевести «вычислительный» шейдер в ранг сущностей первого класса и подготовить почву для внутрибраузерной приватной разработки с опорой на машинное обучение.

Вычислительные (а также вершинные и фрагментные) шейдеры пишутся на языке WGSL. Язык WGSL позволяет разработчикам написать один шейдер, компилируемый в код на низкоуровневых языках, например SPIR-V для Vulkan и MSL для Metal.

Сравнение WebGPU и CUDA

NVIDIA — наиболее популярный вариант для работы с железом, а CUDA, его API, отчасти объясняет такую популярность, но этот API работает только на аппаратном обеспечении от NVIDIA.

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

Также WebGPU на полголовы превосходит CUDA потому, что может компилировать код и в другие языки для GPU, например, Vulkan и Metal. Это своеобразный React Native для вычислительных шейдеров GPU.

Основы вычислительных шейдеров GPU

Мельчайшая интересующая нас единица — это поток, в котором выполняется вычислительный шейдер.

Рабочие группы (workGroups) — это группы потоков. Мы объединяем потоки в такие группы и параллельно выполняем потоки каждой группы (часто такая структура в CUDA называется threadBlock). Потоки одной рабочей группы могут совместно обращаться к одной и той же области памяти.

WebGPU может одновременно диспетчеризовать множество таких рабочих групп, а в терминологии CUDA такая структура называется «грид» (состоящий из threadBlocks).

Как и в CUDA, сами рабочие группы и их диспетчеризация определяются в 3D. Размер workGroup определяется в @workgroup_size(x, y, z),  где число потоков в рабочей группе равно  x * y * z.

Как написать быстрое перемножение матриц

Именно на перемножение матриц тратится большинство операций с плавающей точкой в секунду (ФЛОПс) при работе с большими языковыми моделями, например GPT-4 и Llama. Это базовый примитив для большинства рабочих нагрузок, связанных с машинным обучением и логическим выводом.

Предусмотренная в WebGPU нативная поддержка перемножения матриц распространяется только на малые матрицы, бесполезные при таких рабочих нагрузок, которые характерны для современного глубокого обучения — здесь матрицы могут быть велики. Для справки: в Llama 3.1 70B используются матрицы размером (8192x28672).

Коротко остановимся на математической нотации.

Перемножение матриц

Начнём с того, что определить перемножение матриц можно через три матрицы: A, B, C.

Матрица на матрицу 

Если A — это матрица размером m × n, а B — матрица размером n × p,

то произведение матриц C = AB (при определении не ставится ни точка, ни знак умножения) определяется как матрица размером m × p,

такая, что

для i = 1, ..., m и j = 1, ..., p.

Следовательно, чтобы получить запись произведения вида c_ij ⁠ci, нужно почленно умножить записи из i-й строки A на записи из j-го столбца B и просуммировать эти n произведений. Иными словами, c_ij – это ⁠скалярное произведение i-й строки A и j-го столбца B.

Общее количество ФЛОПс, необходимых для перемножения матриц, равно 2 * M * K * N, поскольку для каждой операции требуется как умножение, так и сложение.

Определяем нижнюю границу нашего ядра

Развивая отличный пример из статьи Саймона Бёма, имеем две матрицы размером 4092x4092, а затем прибавляем к ним ещё одну матрицу размером 4092x4092. Аналогично, имеем:

1. Всего ФЛОПС: 137 миллиардов

2. Общий объём данных для считывания: 201 МБ

3. Общий объём данных для сохранения: 67 МБ

Оговорюсь, что я занимаюсь разработкой на Mac M2 Pro с арифметической интенсивностью ~6 ТФЛОП/с и с  пропускной способностью памяти в 200 ГБ/с.

Итак, минимальное время на работу вычислительного ядра составляет:

(137GFLOP) / (6TFLOPS/s) = 22 мс

А на обращение к памяти:

(267MB) / (200GB/s) = 1,34 мс

Поэтому у нас есть ограничение по вычислительным возможностям (коэффициент ~16x!).

Пишем ядро

Ядро 1: упрощённая реализация

Вот как проще всего вычислить скалярное произведение между матрицами A и B и записать результат в C: для каждой строки в матрице A (формы M), перебираем столбцы A (формы K) и умножаем на соответствующее значение B. В Python код выглядит так:

def matmul(a, b, c):
    """
    Выполнить упрощённое перемножение матриц: C = A * B
    
    :параметр a: входная матрица A формы (m, k)
    :параметр b: входная матрица B формы (k, n)
    :param c: результирующая матрица C формы (m, n), в которой сохраняется результат
    """
    m = len(a)
    k = len(a[0])
    n = len(b[0])
    
    # Перемножаем матрицы
    for i in range(m):
        for j in range(n):
            c[i][j] = 0
            for l in range(k):
                c[i][j] += a[i][l] * b[l][j]

Примерно как и в вышеприведённом коде Python определяем входные значения. Оговоримся, что для выполнения WebGPU из Typescript требуется немало стереотипного кода, и этот вопрос я оставлю самым любознательным для самостоятельного изучения: https://webgpufundamentals.org/webgpu/lessons/webgpu-fundamentals.html. Кроме того, в WGSL поддерживается ряд типов.

struct Dimensions {
  M: u32,
  K: u32,
  N: u32,
}

@group(0) @binding(0) var<uniform> dimensions: Dimensions;
@group(0) @binding(1) var<storage, read> a: array<f32>;
@group(0) @binding(2) var<storage, read> b: array<f32>;
@group(0) @binding(3) var<storage, read_write> result: array<f32>;

А вот наше вычислительное ядро:

@compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let index = global_id.x;
  let row = index / dimensions.N;
  let col = index % dimensions.N;

  if (index < dimensions.M * dimensions.N) {
    var sum = 0.0;
    for (var i: u32 = 0u; i < dimensions.K; i = i + 1u) {
      sum = sum + a[row * dimensions.K + i] * b[i * dimensions.N + col];
    }
    result[row * dimensions.N + col] = sum;
  }
}

Функционально этот код эквивалентен вышеприведённому коду Python! В workgroup_size(1) мы определяем, насколько велика будет рабочая группа (напоминаю, что все данные представляются в 3D).

Итак, поскольку каждая рабочая группа представлена всего одним потоком, она обрабатывает один result[i, j].

Чтобы вычислить полную матрицу, необходимо взять столько записей, сколько содержится в матрице, и вызвать dispatchWorkgroups. Ради упрощения статьи и сокращения листингов, я убрал значительную часть стереотипного кода (нужного для настройки буферов GPU). Здесь я сосредоточусь только на тех деталях, которые важны для оптимизации ядер, написанных на WGSL.

pass.dispatchWorkgroups(a.shape[0] * b.shape[1]) 

где a.shape == M, b.shape[1] == N (почти) для любой матрицы MxN.

Теперь, как показано ниже, открывается большой простор для улучшений!

Самая большая квадратная матрица, которую можно вычислить — это 128x128, всё дело в ограничениях, заложенных в WebGPU (подробнее об этом ниже). Удаётся выйти всего на 1,64 ГФЛОПС, это лишь бледная тень от теоретического максимума в 6 ТФЛОПС.

Почему это ядро работает так медленно? Фактически каждая рабочая группа вычисляет всего одну запись из общего количества в 16 384 элементов (128^2). Даже при условии параллельного выполнения каждая рабочая группа загружает собственную копию матриц. Издержки при запуске дополнительных рабочих групп, пожалуй, даже выше, чем если увеличить количество рабочих потоков в нашей рабочей группе и, следовательно, вычислять в ней больше результатов. Причём, ни одна из рабочих групп не сможет воспользоваться преимуществами кэширования входных данных.

Ядро 2: Дайте ещёёёёёё потоков!

Первое ядро позволяло нам вычислять лишь небольшие квадратные матрицы. Дело в том, что ограничено количество рабочих групп (maxComputeWorkgroupsPerDimension), которые можно одновременно диспетчеризовать.

Поскольку мы запускаем отдельную рабочую группу для обработки одной записи, матрица 256x256 не вписывается в допустимые пределы!

Помните эту часть кода?

@compute @workgroup_size(1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { 

Можно сократить количество назначаемых рабочих групп, увеличив количество потоков в каждой рабочей группе!

Если обновить код,

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { 

то можно сократить количество назначаемых рабочих групп на каждое измерение:

const WORKGROUP_SIZE = 256;
pass.dispatchWorkgroups((a.shape[0] * b.shape[1]) / WORKGROUP_SIZE);

Почему 256? Что ж, а это другой предел :)

Увеличивая размер рабочей группы, можно повысить производительность ядра в 200 раз!

Ядро 3: Вычисление при помощи рабочих групп в 2D

Пока все эти вычисления остаются «одномерными», мы можем вычислить лишь матрицу ограниченного размера. Всё дело в другом ограничении: maxComputeWorkgroupsPerDimension.

Даже не внося особенных изменений в код, можно распределить работу на два измерения и так обойти этот предел — соответственно, запустить больше рабочих групп, при этом более крупных. В таком случае нам поддаётся перемножение матриц размером 4096x4096.

Обновляем @workgroup_size(8, 8), проверяем границы

@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let row = global_id.x;
  let col = global_id.y;

  if (row < dimensions.M && col < dimensions.N) {
    var sum : f32 = 0.0;
    for (var i: u32 = 0u; i < dimensions.K; i = i + 1u) {
      sum = sum + a[row * dimensions.K + i] * b[i * dimensions.N + col];
    }
    result[row * dimensions.N + col] = sum;
  }
}

и диспетчеризуем рабочие группы в двух измерениях:

const WORKGROUP_SIZE = 16; 
pass.dispatchWorkgroups(    
          Math.ceil(a.shape[0]  / WORKGROUP_SIZE), 
          Math.ceil(b.shape[1] / WORKGROUP_SIZE),
);    

Но это ядро работает ещё медленнее исходного! Что же происходит?

Если немного изменить код

@compute @workgroup_size(8, 8)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let row = global_id.y;
  let col = global_id.x;

то производительность ядра значительно увеличится.

Почему так? Просто нам удалось более эффективно задействовать кэшированные входные данные. Измерение x в global_invocation_id  увеличивается раньше, чем измерение y, поэтому в каждой рабочей группе становится больше потоков, работающих с одной и той же строкой в матрице A. В противном случае переменная строки затирается при каждом вызове внутри рабочей группы, и каждому потоку приходится потратить несколько циклов, чтобы прочитать информацию из глобальной памяти, а не из кэша.

Ядро 4: разбиение ядра (kernel tiling)

Также следует учитывать, какое именно количество работы выполняет каждый поток.

До сих пор в этой статье рассматривалось, как каждый поток обрабатывает ровно одну запись. Но сравните издержки: одно дело каждый раз запускать рабочую группу, а другое — вычислять более одного элемента на поток!

Если второй вариант окажется быстрее первого, то мы должны наблюдать большое ускорение. Примерно это и рекомендуется делать при сборке вычислительных ядер в материале от Apple.

В качестве проверки попробуем вычислять четыре результата в каждом потоке (речь о разбиении 1x4).

const BLOCKSIZE: u32 = 16;
const TILESIZE: u32 = 4;
@compute @workgroup_size(BLOCKSIZE, BLOCKSIZE)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let row = global_id.y;
    let col = global_id.x * TILESIZE;

    if (row >= dimensions.M || col >= dimensions.N) {
        return;
    }

    var sum00: f32 = 0.0;
    var sum01: f32 = 0.0;
    var sum02: f32 = 0.0;
    var sum03: f32 = 0.0;

    for (var i: u32 = 0u; i < dimensions.K; i = i + 1u) {
        let a_elem = a[row * dimensions.K + i];
        sum00 = sum00 + a_elem * b[i * dimensions.N + col];
        sum01 = sum01 + a_elem * b[i * dimensions.N + col + 1u];
        sum02 = sum02 + a_elem * b[i * dimensions.N + col + 2u];
        sum03 = sum03 + a_elem * b[i * dimensions.N + col + 3u];
    }

    result[row * dimensions.N + col] = sum00;
    result[row * dimensions.N + col + 1u] = sum01;
    result[row * dimensions.N + col + 2u] = sum02;
    result[row * dimensions.N + col + 3u] = sum03;
}

Кажется, что ядро почти не изменилось, с той оговоркой, что мы размотали вычисления и стали рассчитывать результаты TILESIZE по потокам.

Можно развить этот приём и вычислять в каждом потоке 2D-результаты! Теперь мы будем вычислять не по 4 элемента на каждую строку, а по 4 элемента для 4 строк (речь о 2D -фрагменте).

const BLOCKSIZE: u32 = 16;
const TILE_M: u32 = 4;  // Размер фрагмента в измерении M 
const TILE_N: u32 = 4;  // Размер фрагмента в измерении N 

@compute @workgroup_size(BLOCKSIZE, BLOCKSIZE)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let row = global_id.y * TILE_M;
    let col = global_id.x * TILE_N;

    // initialize the array with all 0s
    var sums: array<array<f32, TILE_N>, TILE_M>;
    for (var i = 0u; i < TILE_M; i++) {
        for (var j = 0u; j < TILE_N; j++) {
            sums[i][j] = 0.0;
        }
    }

    // Вычисляем 2D-фрагмент
    for (var k = 0u; k < dimensions.K; k++) {
        // для каждой строки
        for (var i = 0u; i < TILE_M; i++) {
            let a_element = a[(row + i) * dimensions.K + k];
            // вычисляем скалярное произведение
            for (var j = 0u; j < TILE_N; j++) {
                let b_element = b[k * dimensions.N + (col + j)];
                sums[i][j] += a_element * b_element;
            }
        }
    }

    // Записываем результаты
    for (var i = 0u; i < TILE_M; i++) {
        for (var j = 0u; j < TILE_N; j++) {
            let output_row = row + i;
            let output_col = col + j;
            if (output_row < dimensions.M && output_col < dimensions.N) {
                result[output_row * dimensions.N + output_col] = sums[i][j];
            }
        }
    }
}

Теперь каждый поток вычисляет участок результирующей матрицы, клетку размером 4x4, и мы действительно видим, что это ядро работает немного лучше предыдущего.

Удивительно, но 2D-разбиение получается достаточно медленным. Почему мы не учли время, которое при этом тратится, и не сравнили его со временем на запуск рабочих групп, которые могли бы успеть сделать что-то полезное? Кроме того, почему обработка идёт медленнее, чем при распределении по одной единице на поток?

Ядро 5: размотка

Чтобы ответить на последний вопрос, нам потребуется углубиться в скомпилированные ядра WebGPU.

Некоторые компиляторы автоматически разматывают цикл, если границы цикла известны во время компиляции. Но мы же пишем ядро общего назначения для работы с входными данными переменной формы!

Кроме того, когда пишешь код на WGSL, ты никак не можешь контролировать директивы компилятора.

Если посмотреть ассемблерный бит-код из Metal, то видно, что в наборе инструкций по-прежнему содержится цикл for!

%51 = phi i32 [ 0, %41 ], [ %61, %50 ]
%52 = add i32 %37, %51
%53 = zext i32 %52 to i64
%54 = getelementptr inbounds [1 x float], ptr addrspace(1) %3, i64 0, i64 %53
%55 = load float, ptr addrspace(1) %54, align 4, !tbaa !27, !alias.scope !43, !noalias !44
%56 = zext i32 %51 to i64
%57 = getelementptr inbounds %struct.type_5, ptr %7, i64 0, i32 0, i64 %49, i32 0, i64 %56
%58 = load float, ptr %57, align 4, !tbaa !27
%59 = fmul fast float %55, %48
%60 = fadd fast float %58, %59
store float %60, ptr %57, align 4, !tbaa !27
%61 = add nuw nsw i32 %51, 1
%62 = icmp eq i32 %61, 4
br i1 %62, label %38, label %50 // ветвящийся цикл for

Тогда как после размотки код WGSL компилируется в следующий:

...
%141 = fmul fast float %112, %103
%142 = fadd fast float %141, %82
%143 = fmul fast float %116, %103
%144 = fadd fast float %143, %81
%145 = fmul fast float %120, %103
%146 = fadd fast float %145, %80
%147 = fmul fast float %124, %103
%148 = fadd fast float %147, %79
%149 = fmul fast float %112, %107
%150 = fadd fast float %149, %78
%151 = fmul fast float %116, %107
%152 = fadd fast float %151, %77
%153 = fmul fast float %120, %107
%154 = fadd fast float %153, %76
%155 = fmul fast float %124, %107
%156 = fadd fast float %155, %75
%157 = add nuw i32 %91, 1
%158 = icmp eq i32 %157, %27
br i1 %158, label %159, label %74 

Поскольку мы выполняем размотку вручную, GPU удаётся сократить издержки: он обходится без инициализации и инкремента внутреннего цикла, задействует параллелизм на уровне инструкций, а также компенсировать издержки, возникающие при использовании сравнительно немногочисленных рабочих групп, так как в каждом потоке он успевает выполнить больше работы. Когда мы использовали наш цикл в ядре #4, оно не могло опираться на эти оптимизации, поэтому работа шла медленнее, чем при запуске дополнительных рабочих групп (ядро #3).

А если сделать грид размером 8x8, то получим трёхкратное ускорение по сравнению с циклом 4x4 и даже превзойдём 1ТФЛОП!

Заключение

Мы хорошо поработали и смогли собрать высокопроизводительное ядро для перемножения матриц, которое работает в 1000 раз быстрее неоптимизированного ядра и приближается к пиковым значениям, теоретически достижимым на Apple M2 Pro.

Учитывая, как часто обновляется WebGPU, здесь возможны и другие оптимизации! Например, мы не воспользовались подгруппами, новой фичей, появившейся в Chrome 125. Она должна обеспечивать ускоренный доступ к памяти и совместное использование ресурсов на уровне подгрупп, чтобы лишний раз не повторять вычисления.

Большое спасибо Абхишайке Махаджану (который ведёт потрясающий блог) и Эльману Мансимову за обратную связь и за то, что вдохновили меня написать эту статью!