Программирование графических процессоров Apple через Go и язык Metal Shading

Программирование графических процессоров Apple с использованием языка Go и Metal Shading

Изучение Go, Cgo, металлического языка теней, металлических шейдеров производительности и бенчмаркинг различных подходов к умножению матриц

Ниже я опишу процесс использования cgo для взаимодействия между Go и нативным C, как это можно использовать для взаимодействия с привязками Objective-C для фреймворка металлических шейдеров производительности от Apple, как взаимодействовать с пользовательскими кодами GPU (шейдерами), написанными на металлическом языке теней, и, наконец, сравнить все это с ручным и Go-основанным умножением матриц на OpenBLAS. Этот код был написан для работы на моем ноутбуке M2 MacBook.

Внешний вид исходного кода, доступного здесь на GitHub, выглядит следующим образом:

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

• GPU и параллельное выполнение с плавающей запятой
• Основы работы с Metal GPU
• Металлический язык теней
• Привязки Objective-C
• Фреймворк металлических шейдеров производительности
• Go и cgo
• Реализация на Go и OpenBLAS
• Результаты

GPU и параллельное выполнение с плавающей запятой

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

• Сотрудничающие графовые нейронные сети
• Веб-интерфейс речевого API Что работает, что нет, и как улучшить его, связав его с языковой моделью GPT
• Большие модели встречаются с большим количеством данных Spark и LLM в гармонии

Графические процессоры по своей конструкции чрезвычайно эффективны в выполнении массовых параллельных операций с плавающей запятой, которые требуют высокой пропускной способности памяти. Мой ноутбук M2 MacBook имеет 8 ядер CPU и 8 ядер GPU, но для сравнения, Nvidia RTX 4090 содержит 16384 ядра, а H100 содержит 16896 ядер CUDA с сотнями дополнительных специализированных тензорных ядер. Графические процессоры обычно поддерживают обработку SIMD, что позволяет выполнить одновременно одну и ту же инструкцию на нескольких наборах данных.

Кроме графики, умножение матриц и задачи линейной алгебры в целом имеют преимущество от параллельной обработки благодаря своим высокопараллельным алгоритмам. Это, в свою очередь, поддерживает основные рабочие нагрузки машинного обучения, такие как обучение и вывод [1] [2]].

CUDA вероятно является самой известной платформой для программирования GPU, специфичной для аппаратного обеспечения Nvidia. Также существуют математические фреймворки для OpenGL. Фреймворки, такие как TensorFlow и PyTorch, могут легко и относительно прозрачно интегрироваться с аппаратным обеспечением GPU. This – интересная статья о повышении производительности поддержки фреймворков GPU, основанных на Metal, в библиотеке spaCy NLP.

Metal GPU Basics

Программирование непосредственного вычисления на GPU не так просто, как написание кода для процессоров устройства. При работе с фреймворком Metal Apple, грубая последовательность операций для выполнения кода на GPU выглядит следующим образом:

• Найти соответствующее GPU устройство
• Создать очередь для выполнения команд (т.е. MTLCommandQueue)
• Обернуть указатели на массивы данных в структурированный буфер; если данные являются исполняемым кодом, то состояние конвейера, в противном случае это обычный буфер. Графические процессоры Apple используют объединенное пространство памяти, что означает, что нам не нужно копировать данные в физическую память, специфическую для GPU
• Зафиксировать буфер команд для выполнения и либо дождаться результатов, либо установить обработчик событий по окончании
• Извлечь байты из буфера ответа и отформатировать их локально с помощью программного кода CPU

Программирование графического процессора выполняется асинхронно.

Язык программирования Metal Shading

Язык программирования Metal Shading является производным от C++14, который может использоваться для составления пользовательской логики (называемой “шейдерами”) для выполнения на совместимых с Metal графических процессорах. В целом, если это возможно, вероятнее всего вам будет лучше использовать фреймворк MPS (обсуждается позже) для эквивалентного функционала при возможности – он обычно оптимизирован для распространенных классов использования, выравненных на графический процессор, таких как умножение матриц или нейронные сети.

Отладка кода MSL довольно сложна. Вы можете использовать Отладчик шейдеров через Xcode, но если вы хотите проверить или вывести промежуточные значения без использования Xcode, вам придется записать данные в буфер отладки ответа и разобрать примитивы в вашей обертке на C++ или Objective-C.

Функции MSL доступны как общедоступные интерфейсы через обозначение kernel. Фреймворк Metal передает идентификаторы для текущего контекста вызывающего потока или группы потоков, которые могут использоваться для обеспечения неперекрывающей записи. Потоки могут быть представлены системой трехмерных идентификаторов; размеры этого пространства потока настраиваются в коде оболочки.

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

kernel void matrix_multiply_naive(  device const MatrixParams *params,  constant float *A,  constant float *B,  device float *C,  // Указывает уникальную позицию потока в рамках всей сетки выполняемых потоков. Тип uint2 является 2D координатой, где поля x и y представляют его индексы по каждой оси.  // Этот параметр не передается непосредственно из вызывающего кода,  // но предоставляется фреймворком Metal  uint2 gid [[thread_position_in_grid]]) {  if (gid.x >= params->a_rows || gid.y >= params->b_cols) {    return; // Этот поток находится вне диапазона размерности матрицы, делать ничего  }  float sum = 0.0;  int k;  // Unrolling цикла; улучшение производительности на значительный уровень  for (k = 0; k <= params->a_cols - 4; k += 4) {    sum += A[gid.x * params->a_cols + k]        * B[k * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 1]        * B[(k + 1) * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 2]        * B[(k + 2) * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 3]        * B[(k + 3) * params->b_cols + gid.y];  }  // Обработка оставшихся элементов  for (; k < params->a_cols; ++k) {    sum += A[gid.x * params->a_cols + k] * B[k * params->b_cols + gid.y];  }  C[gid.x * params->b_cols + gid.y] = sum;}

Я также реализовал наивную транспонированную функцию на MSL для сравнения. Учитывая транспонированную матрицу, эта функция легко адаптируется к вышеприведенной логике, где внутренний цикл запускается по строкам B, а не по столбцам:

// Unrolling цикла; улучшение производительности на значительный уровеньfor (k = 0; k <= params->a_cols - 4; k += 4) {  sum += A[gid.x * params->a_cols + k]         * B[gid.y * params->b_cols + k]; // Обратите внимание, что это gid.y * cols плюс k  sum += A[gid.x * params->a_cols + k + 1]     * B[gid.y * params->b_cols + k + 1];  sum += A[gid.x * params->a_cols + k + 2]     * B[gid.y * params->b_cols + k + 2];  sum += A[gid.x * params->a_cols + k + 3]     * B[gid.y * params->b_cols + k + 3];}// Обработка оставшихся элементовfor (; k < params->a_cols; ++k) {  sum += A[gid.x * params->a_cols + k] * B[gid.y * params->b_cols + k];}

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

Привязки Objective-C

Фреймворк Metal предоставляет возможность компилировать библиотеку из исходного кода Metal. После загрузки содержимого файла, код привязки ищет ядреные функции по имени и инициализирует новый MTLComputePipelineState, представляющий скомпилированный код функции.

id<MTLDevice> device = MTLCreateSystemDefaultDevice();// Компиляция и инициализация новой библиотеки по указанному пути исходного кода.MTLCompileOptions *compileOptions = [MTLCompileOptions new];compileOptions.languageVersion = MTLLanguageVersion3_0;// Обертка строки пути вводаNSString *ss = [NSString stringWithUTF8String:source_path];// Инициализация новой библиотеки, содержащей скомпилированные шейдерные функцииid<MTLLibrary> lib = [device newLibraryWithSource:ss  options:compileOptions  error:&error];// Создание представления наивной умножающей публичной шейдерной функции в // созданной выше библиотеке Metalid<MTLFunction> naiveFunction =    [lib newFunctionWithName:@"matrix_multiply_naive"];// Создание нового состояния вычислительного конвейераid<MTLComputePipelineState> pipelineStateNaive = [device newComputePipelineStateWithFunction:naiveFunction  error:&error];

Чтобы вызвать нативный код Metal, необходимо установить конфигурацию потока и инициализировать буферы GPU.

[computeEncoder setComputePipelineState:pipelineStateNaive];MTLSize threadsPerGrid = MTLSizeMake(params->a_cols, params->a_rows, 1);// Рассчитать размер группы потоков.// https://developer.apple.com/documentation/metal/calculating_threadgroup_and_grid_sizes?language=objcNSUInteger w = pipelineStateNaive.threadExecutionWidth;NSUInteger h = pipelineStateNaive.maxTotalThreadsPerThreadgroup / w;MTLSize threadsPerThreadgroup = MTLSizeMake(w, h, 1);// Закодировать входные данные ядра функции[computeEncoder setBytes:params length:16 atIndex:0];[computeEncoder setBuffer:bufferA offset:0 atIndex:1];[computeEncoder setBuffer:bufferB offset:0 atIndex:2];[computeEncoder setBuffer:bufferC offset:0 atIndex:3];// Закодировать команду вычисления.[computeEncoder dispatchThreads:threadsPerGrid   threadsPerThreadgroup:threadsPerThreadgroup];// Завершить вычислительный проход.[computeEncoder endEncoding];// Выполнить команду.[commandBuffer commit];

Это много, поэтому я проясню отношения здесь:

Фреймворк Metal Performance Shaders

Фреймворк MPS – это высокопроизводительная библиотека, предоставляемая Apple для использования с ее семейством графических процессоров Metal. Он предлагает функциональность от задач обработки изображений до поддержки нейронных сетей.

API в основном доступны через Swift или Objective-C, хотя также доступна библиотека Metal-cpp для использования.

API MPSMatrixMultiplication довольно прост в использовании. Как и в предыдущем коде MSL, команды MPS все еще должны быть закодированы в MTLCommandBuffer и асинхронно подтверждены для выполнения.

// Определение "описаний" матриц, учитывающих размерность и размер байтовMPSMatrixDescriptor *descriptorA = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:a_cols  rowBytes:a_cols * sizeof(float)  dataType:MPSDataTypeFloat32];MPSMatrixDescriptor *descriptorB = [MPSMatrixDescriptor matrixDescriptorWithDimensions:b_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// Матрица выводаMPSMatrixDescriptor *descriptorC = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// Инициализировать матричные представления с использованием приведенных выше описаний и буферов матрицMPSMatrix *matrixA = [[MPSMatrix alloc] initWithBuffer:bufferA descriptor:descriptorA];MPSMatrix *matrixB = [[MPSMatrix alloc] initWithBuffer:bufferB descriptor:descriptorB];MPSMatrix *matrixC = [[MPSMatrix alloc] initWithBuffer:bufferC descriptor:descriptorC];// Создать экземпляр умножения матрицMPSMatrixMultiplication *matrixMultiplication = [[MPSMatrixMultiplication alloc] initWithDevice:device  resultRows:a_rows  resultColumns:b_cols  interiorColumns:a_cols];// Закодировать команду умножения в буфер команд для GPUid<MTLCommandBuffer> commandBuffer = [commandQueue commandBuffer];[matrixMultiplication encodeToCommandBuffer:commandBuffer  leftMatrix:matrixA  rightMatrix:matrixB  resultMatrix:matrixC];

Go и cgo

Мне не особенно нравится работать с Objective-C, и целью этой программы является запуск кода на GPU, созданного из программы на Go.

Cgo – это функция языка Go, которая позволяет компилятору Go понимать директивы компилятора, содержащиеся в комментариях, относящихся к нативному коду на языке C. Она поддерживает версию программного интерфейса внешних функций.

Конфигурация директивы немного хрупкая, но любые комментарии, непосредственно предшествующие строке import “C” (называемой «преамбулой»), будут интерпретироваться как импорты заголовков или аргументы компиляции при компиляции связанного кода на C. Например:

/*#cgo LDFLAGS: -framework Foundation -framework CoreGraphics -framework Metal -framework MetalPerformanceShaders -L/opt/homebrew/opt/openblas/lib -lopenblas#include <stdlib.h>#include "metal.h"*/import "C"

• Передает флаги линковки линкеру через командную строку LDFLAGS
• Компилирует код на C с помощью стандартного заголовка stdlib.h
• Компилирует код на C с помощью локального заголовка проекта metal.h

Мне потребовалось некоторое время для настройки правильного набора флагов линковки на MacOS.

• Foundation: базовые библиотеки
• CoreGraphics: необходимо на MacOS для взаимодействия с GPU
• Metal: библиотеки и поддержка языка для Metal, включая MSL
• MetalPerformanceShaders: библиотеки для MPS, упомянутые выше

Оказалось, что Apple включает в свой фреймворк Accelerate реализацию BLAS, поэтому, помимо установки OpenBLAS через brew, необходимо также указать местоположение библиотеки при линковке:

-L/opt/homebrew/opt/openblas/lib -lopenblas

Директива go:embed позволяет программам на Go включать файлы на этапе компиляции, что полезно в данном случае, когда мы хотим передать содержимое исходного файла на языке MSL (mm.metal) в фреймворк Metal для компиляции, как уже обсуждалось выше.

//go:embed mm.metalvar source string// Компилирует исходный код шейдера и инициализирует конвейеры. Параметр metalSource// содержит содержимое встроенного файла языка Metal Shading Language.func Compile (metalSource string) { // Обернуть строку в строку C src := C.CString(metalSource) // Освободить вышеприведенную строку после инициализации командной очереди defer C.free(unsafe.Pointer(src)) // Компилировать исходный код, инициализировать конвейеры и командную очередь C.initializePipelineAndCommandQueue(src)}

Упомянутые выше ссылки на C используются для взаимодействия с C API через cgo, например:

// Вызывает initializeMTLBuffers из привязок Obj-CC.initializeMTLBuffers( a_data,                  // Входной непрозрачный указатель для A b_data,                  // Входной непрозрачный указатель для B C.int(4),                // Преобразует 4 в целочисленный тип C.int(a.Size()),          C.int(b.Size()),          C.int(a.Rows * b.Cols))params := MatrixParams{ a_rows: int32(a.Rows), a_cols: int32(a.Cols), b_rows: int32(b.Rows), b_cols: int32(b.Cols),}// Возвращает небезопасный указатель на эту структуру MatrixParams, приведенный к // собственному представлению на языке C, определенному в общем заголовочном файлеreturn (*C.MatrixParams)(unsafe.Pointer(&params));

Обратите внимание, что это означает, что C является зарезервированным ключевым словом и не может использоваться в качестве имени переменной.

Варианты реализации на Go и OpenBLAS

Мне хотелось сравнить производительность умножения матриц на основе GPU с библиотеками на более высоком уровне, такими как библиотека Gonum, а также интуитивно понятными, написанными вручную (и сравнительно неэффективными) реализациями.

Я реализовал несколько различных алгоритмов на Go, включая этот параллельный алгоритм наивного транспонирования, который наивно делит работу по умножению между N горутинами:

func (a Matrix[T]) TransposeMultParallel(b *Matrix[T]) *Matrix[T] {
    if a.Cols != b.Rows {
        panic("матрицы имеют неправильный размер для умножения")
    }
    c_data := make([]T, a.Rows*b.Cols)
    t := b.Transpose()
    var wg sync.WaitGroup
    for i := 0; i < a.Rows; i++ {
        wg.Add(1) // Добавить счетчик в WaitGroup для нового goroutine
        go func(i int) { // Запустить goroutine
            defer wg.Done() // Уменьшить счетчик по завершении goroutine
            ptr := i * b.Cols
            for j := 0; j < b.Cols; j++ {
                var sum T = 0.0
                for k := 0; k < a.Cols; k++ {
                    sum += a.At(i, k) * t.At(j, k)
                }
                c_data[ptr+j] = sum
            }
        }(i)
    }
    wg.Wait() // Ждать завершения всех goroutine
    return InitMatrixWithData(a.Rows, b.Cols, c_data)
}

Gonum BLAS – это чистая библиотека на Go, которая реализует интерфейсы BLAS. Однако она также может быть настроена для перенаправления алгебраических операций на реализацию BLAS на нативном коде, такую как OpenBLAS через netlib.

Я показал выше, как cgo может быть настроен для правильной связи с установкой OpenBLAS на MacOS. В коде приложения предпочтительная реализация BLAS может быть установлена напрямую. Из кода бенчмарка:

// Конвертируйте примитивные массивы в плотные матрицы gonumgonum_a := mat.NewDense(a_rows, a_cols, a64_data)gonum_b := mat.NewDense(b_rows, b_cols, b64_data)gonum_c := mat.NewDense(a_rows, b_cols, nil)gonum_d := mat.NewDense(a_rows, b_cols, nil)// Настройте Gonum на использование реализации Gonum-по умолчанию на Go blas64.Use(gonum.Implementation{})// Выполните умножение с использованием Gonum BLAS impl start = time.Now()gonum_c.Mul(gonum_a, gonum_b)bdata.TimeGonumNative(start)// Настройте Gonum на использование netlib, который пересылает операции на // реализацию C-code BLAS (в нашем случае OpenBLAS)blas64.Use(netlib.Implementation{})// Выполнение умножения с использованием реализации OpenBLAS через Gonum APIstart = time.Now()gonum_d.Mul(gonum_a, gonum_b)bdata.TimeGonumOpenBLAS(start)

Результаты

Мой код бенчмаркирования запускает несколько испытаний каждой из следующих реализаций умножения матриц, и сообщает среднее время, затраченное на умножение двух квадратных матриц постепенно возрастающей размерности:

- Наивное умножение в Go- Умножение с транспонированием в Go- Горутина-параллельное умножение с транспонированием в Go- Умножение Gonum на основе pure Go- Умножение OpenBLAS через Gonum, написанное на C- Хенд-имплементированное наивное умножение в MSL на GPU- Хенд-имплементированное умножение с транспонированием в MSL, на GPU- Фреймворк Metal Performance Shaders, вызывается из Objective-C, на GPU

Вывод бенчмаркинга выглядит следующим образом (числа с плавающей запятой в ms):

2023-12-01 11:12:51.644 go-mm[75818:22427382] Использование устройства Apple M2 по умолчаниюэлементы наивное транспонирование параллельное транспонирование metal-наивное metal-транспонирование mps gonum openblas160000 196.00 201.00 42.00 8.00 9.67 0.33 4.67 6.00250000 381.33 387.67 80.67 11.00 11.67 0.00 8.33 21.00360000 801.00 789.33 159.33 19.00 16.33 0.00 14.33 4.67490000 1228.00 1075.00 411.00 23.67 24.33 1.00 26.67 16.33...

Некоторая быстрая построение графиков через matplotlib

Как можно было ожидать, мои самостоятельно написанные реализации на Go неуправляемы. На самом деле, другие подходы настолько быстры, что на графике невозможно их различить. Вот скользящая гистограмма использования GPU во время этого запуска

Вы можете видеть, что GPU не особенно занят, потому что большая часть времени затрачивается на операции процессора. Вот еще один запуск, исключающий самые медленные три метода умножения:

Приближаясь к 16M элементам (4k x 4k), Gonum начинает деградировать. Здесь явно видно, что операции, основанные на GPU, и операции OpenBLAS превосходят чистые реализации на Go. Если смотреть только на подходы, основанные на GPU:

Есть пару интересных замечаний:

• Библиотека Metal Performance Shaders удивительно быстра
• Между наивным подходом и транспонированным наивным подходом нет реальной разницы в производительности

Что касается второго пункта: это отличается от характеристик производительности пары реализаций на Go выше. Оказывается, что выгодные шаблоны кэш-доступа для ЦП не работают так же для ГПУ и того, как их SIMD-группы (или варпы) обращаются к памяти. Сравните использование ГПУ здесь:

Теперь рассмотрим только OpenBLAS и MPS — два самых быстрых подхода:

Примерно при 35M элементах реализация OpenBLAS начинает деградировать, тогда как MPS стабильно продолжает работу. Разница здесь довольно заметна, при этом последний завершает те же самые операции умножения матриц с 35M элементами за менее чем 15% от времени. Разумно предположить, что разница будет продолжать расти с увеличением матричной величины.

Конечно, между этими двумя подходами, вероятно, есть алгоритмические различия, поэтому это не справедливое сравнение ЦП и ГПУ. Если построить график различий в производительности между моими двумя самостоятельно написанными реализациями, он будет выглядеть так:

Это значит, что наивная реализация на основе MSL выполняет умножение 5M элементов всего за 1% времени, затрачиваемого на мою реализацию на языке Go, и со временем этот коэффициент, кажется, становится все больше в пользу GPU.