SMID оптимизация в Go

Перевод статьи "SIMD Optimization in Go".

Из этой статьи вы узнаете много прекрасного о программировании на ассемблере.

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

Программирование на ассемблере может быть очень медленным и сложным в сравнении с программированием на более высокоуровневых языках, таких как Go. Но иногда в этом есть смысл, к тому же, это бывает весело. Например замечательная игра на ассемблере TIS-100 с ее 14 страничной инструкцией доставляет удовольствие без всяких реальных компьютеров с его 3,883 страничным мануалом и целым арсеналом инструментов сборки.

Зачем тратить свое время на ассемблер, если есть значительно более удобные языки программирования? Ваш отец сделал этот новый язык собственноручно, но для вас он недостаточно хорош? Он потратил годы и перепробовал кучу подходов для решения проблем, встретившихся на этом пути, а после кристализировал все это в навороченных компиляторах, которые не позволят вам повторить его ошибки.

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

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

Пишем ассемблер в Go

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

package main

import "fmt"

func add(x, y int64) int64 {
    return x + y
}

func main() {
    fmt.Println(add(2, 3))
}

Для реализации аналогичного кода на ассемблере необходимо создать файл с названием add_amd64.s в котором, собственно, мы будем писать наш ассемблерный код. Так как я использую мак, то я буду писать реализацию для архитектуры AMD64.

// add.go
package main

import "fmt"

func add(x, y int64) int64

func main() {
    fmt.Println(add(2, 3))
}

// add_amd64.s
TEXT ·add(SB),NOSPLIT,$0
    MOVQ x+0(FP), BX
    MOVQ y+8(FP), BP
    ADDQ BP, BX
    MOVQ BX, ret+16(FP)
    RET

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

mkdir -p src
unzip add.zip -d src/add
GOPATH=$PWD go build add
./add

Синтаксис ассемблера... обескураживает. В официальном руководстве Go есть небольшой, несколько древний, гайд по ассемблеру в Plan 9, он дает некоторое представление о том как работает Go ассемблер. Но самая лучшая документация, конечно же, это существующий код на Go ассемблере и скомпилированные в ассемблер Go функции, которые вы можете получить с помощью команды:

go tool compile -S <go file>

Самое важно, что следует запомнить в первую очередь, это объявление функций и использование стека.

Наше магическое заклинание начинается с TEXT ·add(SB),NOSPLIT,$0. Да, обратите внимание, что используется специальный юникодный символ ·, который играет роль разделителя для названия пакета и функции. Так как мы находимся в пакете main, его название не нужно указывать, а функция в нашем примере называется add. Директива NOSPLIT означает что не нужно указывать размер аргументов следующим параметров. Константа $0 стоит как раз на том месте, где нужно было указать размер аргументов, но так как мы уже используем NOSPLIT, то просто оставим $0. К сожалению, не совсем понятно, почему после названия функции стоит (SB), но без этого ничего не работает.

Аргументы в функцию передаются из стека в соответствующем порядке, начиная с 0(FP), что означает нулевой байт смещения от указателя FP, затем идут все остальные параметры до самого возвращаемого значения. Для функции func add(x, y int64) int64 это выглядит примерно так:

И ассемблерный код для справки:

TEXT ·add(SB),NOSPLIT,$0
    MOVQ x+0(FP), BX
    MOVQ y+8(FP), BP
    ADDQ BP, BX
    MOVQ BX, ret+16(FP)
    RET

Ассемблерная версия функции add загружает переменную a из памяти по адресу +0(FP) в регистр BX. Затем грузит b из памяти по адресу +8(FP) в регистр BP, добавляет BP к BX сохраняя результаты в BX и, в конце концов, копирует BX в память по адресу +16(FP) и выходит из функции. Вызывающая функция, которая размещает аргументы на стеке, в итоге читает возвращаемое значение по конкретному адресу в памяти.

Оптимизация функций с помощью ассемблера

Конечно, нам не нужно писать на ассемблере программу для складывания двух чисел. Но что мы вообще будем с ним делать?

Давайте представим, что у вас есть пачка векторов и вам нужно умножить их на матрицу преобразований. Возможно, векторы это точки и вы хотите переместить их в 3D пространстве. Мы будем использовать 4D векторы и матрицу преобразования 4x4, векторы будут упакованы в массив. Это лучше чем использование массива с объектами которые представляют вектор, так как линейное сканирование упакованных данных в памяти позволяет лучше использовать кеш.

type V4 [4]float32
type M4 [16]float32

func M4MultiplyV4(m M4, v V4) V4 {
    return V4{
        v[0]*m[0] + v[1]*m[4] + v[2]*m[8] + v[3]*m[12],
        v[0]*m[1] + v[1]*m[5] + v[2]*m[9] + v[3]*m[13],
        v[0]*m[2] + v[1]*m[6] + v[2]*m[10] + v[3]*m[14],
        v[0]*m[3] + v[1]*m[7] + v[2]*m[11] + v[3]*m[15],
    }
}

func multiply(data []V4, m M4) {
    for i, v := range data {
        data[i] = M4MultiplyV4(m, v)
    }
}

На моем 2012 Macbook Pro с использованием Go 1.5.3 этот код выполняется за 140ms при размере данных в 128MB. Насколько быстрой может быть реализация? Для копирования памяти, кажется, хорошим показателем 14ms.

Ниже приведена функция написанная на ассемблере с использованием SIMD инструкций для умножения. Она позволяет умножать 4 числа типа float32 параллельно:

// func multiply(data []V4, m M4)
//
// memory layout of the stack relative to FP
//  +0 data slice ptr
//  +8 data slice len
// +16 data slice cap
// +24 m[0]  | m[1]
// +32 m[2]  | m[3]
// +40 m[4]  | m[5]
// +48 m[6]  | m[7]
// +56 m[8]  | m[9]
// +64 m[10] | m[11]
// +72 m[12] | m[13]
// +80 m[14] | m[15]

TEXT ·multiply(SB),NOSPLIT,$0
  // data ptr
  MOVQ data+0(FP), CX
  // data len
  MOVQ data+8(FP), SI
  // index into data
  MOVQ $0, AX
  // return early if zero length
  CMPQ AX, SI
  JE END
  // load the matrix into 128-bit wide xmm registers
  // load [m[0], m[1], m[2], m[3]] into xmm0
  MOVUPS m+24(FP), X0
  // load [m[4], m[5], m[6], m[7]] into xmm1
  MOVUPS m+40(FP), X1
  // load [m[8], m[9], m[10], m[11]] into xmm2
  MOVUPS m+56(FP), X2
  // load [m[12], m[13], m[14], m[15]] into xmm3
  MOVUPS m+72(FP), X3
LOOP:
  // load each component of the vector into xmm registers
  // load data[i][0] (x) into xmm4
  MOVSS    0(CX), X4
  // load data[i][1] (y) into xmm5
  MOVSS    4(CX), X5
  // load data[i][2] (z) into xmm6
  MOVSS    8(CX), X6
  // load data[i][3] (w) into xmm7
  MOVSS    12(CX), X7
  // copy each component of the matrix across each register
  // [0, 0, 0, x] => [x, x, x, x]
  SHUFPS $0, X4, X4
  // [0, 0, 0, y] => [y, y, y, y]
  SHUFPS $0, X5, X5
  // [0, 0, 0, z] => [z, z, z, z]
  SHUFPS $0, X6, X6
  // [0, 0, 0, w] => [w, w, w, w]
  SHUFPS $0, X7, X7
  // xmm4 = [m[0], m[1], m[2], m[3]] .* [x, x, x, x]
  MULPS X0, X4
  // xmm5 = [m[4], m[5], m[6], m[7]] .* [y, y, y, y]
  MULPS X1, X5
  // xmm6 = [m[8], m[9], m[10], m[11]] .* [z, z, z, z]
  MULPS X2, X6
  // xmm7 = [m[12], m[13], m[14], m[15]] .* [w, w, w, w]
  MULPS X3, X7
  // xmm4 = xmm4 + xmm5
  ADDPS X5, X4
  // xmm4 = xmm4 + xmm6
  ADDPS X6, X4
  // xmm4 = xmm4 + xmm7
  ADDPS X7, X4
  // data[i] = xmm4
  MOVNTPS X4, 0(CX)
  // data++
  ADDQ $16, CX
  // i++
  INCQ AX
  // if i >= len(data) break
  CMPQ AX, SI
  JLT LOOP
END:
  // since we use a non-temporal write (MOVNTPS)
  // make sure all writes are visible before we leave the function
  SFENCE
  RET

Выполнение функции занимает 18ms, а это уже очень близко к скорости копирования памяти. Еще большую оптимизацию можно получить при использовании GPU вместо CPU. Вот время выполнения разных версий реализации, включая заинлайненную Go версию и реализацию без использования SMID:

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

Замечания по реализации

Я разрабатывал ассемблерную часть на C и x64 ассемблере используя XCode и затем портировал ассемблер в Go формат. В XCode есть замечательный дебагер, который позволяет просматривать регистры во время выполнения кода. Если добавить .s файл в XCode то он будет компилироваться и линковаться с вашим бинарником.

Я использовал документацию по набору инструкций к Intel x64 и Intel Intrinsics Guide что бы разобраться, какие инструкции использовать. Конвертирование в Go ассемблер это не так уж и просто, но большинство инструкций есть в x86/anames.go, а все остальные можно кодировать непосредственно в бинарное представление.