Самый медленный способ ускорить программу на Go

Есть что-то прекрасное в программировании на ассемблере. Оно может быть очень медленным и полным ошибок, по сравнению с программированием на языке, таким как Go, но иногда — это хорошая идея или, по крайней мере, очень весёлое занятие.

Зачем тратить время на программирование на ассемблере, когда есть отличные языки программирования высокого уровня? Даже с сегодняшними компиляторами все ещё есть несколько случаев, когда захотите написать код на ассемблере. Таковыми являются криптография, оптимизация производительности или доступ к вещам, которые обычно недоступны в языке. Самое интересное, конечно же, оптимизация производительности.

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

Пишем ассемблерный код в Go

Лучший способ начать — написать простейшую функцию. Например, функция add складывает два int64.

package main

import "fmt"

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

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

Запуск: go build -o add-go && ./add-go

Для реализации этой функции на ассемблере создайте отдельный файл 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:

#include "textflag.h"

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

Для запуска примера поместите два этих файла в одну директорию и выполните команду go build -o 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.

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

Разберём код уже знакомой функции add:

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 загружает переменную x по адресу памяти +0(FP) в регистр BX. Затем она загружает из памяти y по адресу +8(FP) в регистр BP, складывает BP и BX, сохраняя результат в BX, и, наконец, копирует BX по адресу +16(FP) и возвращается из функции. Вызывающая функция, которая помещает все аргументы в стек, будет читать возвращаемое значение, оттуда где мы его оставили.

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

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

Допустим, у вас есть куча векторов, и вы хотите их умножить на матрицу преобразования. Возможно, векторы являются точками, и вы хотите переместить их в пространстве (перевод на Хабре — прим. пер.). Мы будем использовать векторы с матрицей преобразования размером 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)
    }
}

Выполнение занимает 140 мс для 128 МБ данных. Какая реализация может быть быстрее? Эталоном будет копирование памяти, которое занимает около 14 мс.

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

#include "textflag.h"

// func multiply(data []V4, m M4)
//
// компоновка памяти стека относительно FP
//  +0 слайс data, ptr
//  +8 слайс data, len
// +16 слайс data, 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
  // указатель на data
  MOVQ $0, AX
  // ранний возврат, если нулевая длина
  CMPQ AX, SI
  JE END
  // загрузка матрицы в 128-битные xmm-регистры (https://en.wikipedia.org/wiki/Streaming_SIMD_Extensions#Registers)
  // загрузка [m[0], m[1], m[2], m[3]] в xmm0
  MOVUPS m+24(FP), X0
  // загрузка [m[4], m[5], m[6], m[7]] в xmm1
  MOVUPS m+40(FP), X1
  // загрузка [m[8], m[9], m[10], m[11]] в xmm2
  MOVUPS m+56(FP), X2
  // загрузка [m[12], m[13], m[14], m[15]] в xmm3
  MOVUPS m+72(FP), X3
LOOP:
  // загрузка каждого компонента вектора в регистры xmm
  // загрузка data[i][0] (x) в xmm4
  MOVSS    0(CX), X4
  // загрузка data[i][1] (y) в xmm5
  MOVSS    4(CX), X5
  // загрузка data[i][2] (z) в xmm6
  MOVSS    8(CX), X6
  // загрузка data[i][3] (w) в xmm7
  MOVSS    12(CX), X7
  // копирование каждого компонента матрицы в регистры
  // [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:
  // так как используем невременные (Non-Temporal) SSE-инструкции (MOVNTPS)
  // убедимся, что все записи видны перед выходом из функции (с помощью SFENCE)
  SFENCE
  RET

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

Время работы для разных программ, включая инлайн-версию Go и ассемблерную реализацию без SIMD (со ссылками на исходный код):

Программа	Время, мс	Ускорение
Оригинальная, zip	140	1x
Инлайн-версия, zip	69	2x
Ассемблерная, zip	43	3x
Ассемблерная с SIMD, zip	17	8x
Копирование памяти, zip	15	9x

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

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

Автор разработал ассемблерные части в основном на C и 64-битном ассемблере с использованием XCode, а затем портировал их в формат Go. У XCode хороший отладчик, который позволяет проверять значения регистров процессора во время работы программы. Если включить ассемблерный файл .s в проект XCode, IDE соберёт его и слинкует его с нужным исполняемым файлом.

Автор использовал справочник по набору инструкций Intel x64 и руководство Intel Intrinsics, чтобы выяснить, какие инструкции нужно использовать. Преобразование на язык ассемблера Go не всегда простое, но многие 64-битные ассембленые инструкции указаны в x86/anames.go, а если нет, они могут быть закодированы напрямую с двоичным представлением.

---

Примечание переводчика

В оригинале статьи в ассемблерные файлы не включён заголовок #include "textflag.h", из-за чего при компиляции выдаётся ошибка illegal or missing addressing mode for symbol NOSPLIT.
Поэтому выложил на GitHub Gist исправленные версии. Для запуска распаковываем архив, выполняем команды: chmod +x run.sh && ./run.sh.

Запускать код с ассемблером можно лишь собрав бинарник с помощью go build, иначе компилятор ругнётся на пустое тело функции.

К сожалению, в интернете действительно мало информации по ассемблеру в Go. Советую почитать статью на Хабре про архитектуру ассемблера Go.

Ещё один способ использовать ассемблерные вставки в Go

Как известно, Go поддерживает использование кода, написанного на C. Поэтому, ничего не мешает сделать так:

sqrt.h:

double sqrt(double x) {
  __asm__ ("fsqrt" : "+t" (x));
  return x;
}

sqrt.go:

package main

/*
#include "sqrt.h"
*/
import "C"
import "fmt"

func main() {
  fmt.Printf("%f\n", C.sqrt(C.double(16)))
}

И запустить:

$ go run sqrt.go 
4.000000

Ассемблер — это весело!