Решение задачи с Route 256 на goalng

В этой статье разбирается решение задачи «Гистограммы» с контеста Route 256 от Ozon.

Ссылочки:
Assembler в Go: техники ускорения и оптимизации / Хабр
Руководство по ассемблеру Go / Хабр
Часть 1. Почему Go-ассемблер и векторизация могут быть полезны: идея для ускорения / Хабр

Условие задачи

Гистограммой является массив, каждый элемент которого указывает высоту столбика на соответствующей позиции. Две гистограммы считаются совпадающими, если при совмещении одной гистограммы с другой гистограммой, повёрнутой на угол 180°, получается ровный прямоугольник без наложений и пропусков.

Иллюстрация

Пример теста

1 2 4
1 3 4
1 4 3

Ответ: 1 пара

Тривиальное решение:

func profile(a []uint32) string { n := len(a) res := make([]uint32, n-1) for i := 1; i < n; i++ { res[i-1] = a[i] - a[i-1] }  b := unsafe.Slice((*byte)(unsafe.Pointer(&res[0])), (n-1)*4) return *(*string)(unsafe.Pointer(&b)) }  func reverseProfile(a []uint32) string { n := len(a) res := make([]uint32, n-1) for i := n - 1; i > 0; i-- { res[n-1-i] = a[i] - a[i-1] }  b := unsafe.Slice((*byte)(unsafe.Pointer(&res[0])), (n-1)*4) return *(*string)(unsafe.Pointer(&b)) }  func processHistograms(histos [][]uint32) int { profCount := make(map[string]int) for i := 0; i < len(histos); i++ { p := profile(histos[i]) profCount[p]++ }  res := 0 for i := 0; i < len(histos); i++ { rp := reverseProfile(histos[i]) res += profCount[rp] // если профиль и реверс совпадают, не считать пару (a,a) if profile(histos[i]) == rp { res-- } } return res / 2 }

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

Это решение проходит тесты с контеста, но я решил попробовать ускорить вычисление с помощью SIMD.

Ускорить func profile довольно просто.

// func calcDifferences(src []int32, dst []int32) // комменты от GPT-4.1 TEXT ·calcDifferences(SB), $0-48     MOVQ src_base+0(FP),    SI  // Загружаем базовый адрес исходного массива     MOVQ src_len+8(FP),     CX  // Загружаем длину исходного массива     MOVQ dst_base+24(FP),   DI  // Загружаем базовый адрес массива назначения          CMPQ CX, $1                 // Проверяем, если длина <= 1     JLE  done                   // Если да, завершаем работу          DECQ CX                     // Длина dst = длина src - 1     XORQ AX, AX                 // Инициализируем счетчик индекса нулем      // Проверяем, достаточно ли элементов для AVX обработки (минимум 8)     CMPQ CX, $8     JL   scalar_loop            // Если меньше 8, переходим к скалярной обработке      // Вычисляем количество элементов для AVX обработки (кратно 8)     MOVQ CX, DX     ANDQ $-8, DX                // Округляем вниз до ближайшего кратного 8  avx_loop:     VMOVDQU (SI)(AX*4), Y0      // Загружаем 8 элементов из src в YMM0     VMOVDQU 4(SI)(AX*4), Y1     // Загружаем следующие 8 элементов (со сдвигом +1)     VPSUBD Y0, Y1, Y2           // Вычисляем разницы (Y1 - Y0)     VMOVDQU Y2, (DI)(AX*4)      // Сохраняем 8 результатов в dst          ADDQ $8, AX                 // Увеличиваем индекс на 8     CMPQ AX, DX                 // Сравниваем с количеством обработанных элементов     JL   avx_loop               // Продолжаем цикл, если есть элементы для AVX обработки      // Обрабатываем оставшиеся элементы скалярно     CMPQ AX, CX     JGE  done                   // Если все элементы обработаны, завершаем  scalar_loop:     MOVL (SI)(AX*4), BX         // Загружаем src[i] в регистр BX     MOVL 4(SI)(AX*4), DX        // Загружаем src[i+1] в регистр DX     SUBL BX, DX                 // Вычисляем разницу: DX = src[i+1] - src[i]     MOVL DX, (DI)(AX*4)         // Сохраняем результат в dst[i]          INCQ AX                     // Увеличиваем индекс     CMPQ AX, CX                 // Сравниваем с длиной массива     JL   scalar_loop            // Продолжаем, если индекс < длины  done:     RET                         // Возврат из функции 

// Стандартная реализация на Go для проверки func CalcDifferencesGo(arr []uint32) []uint32 { if len(arr) <= 1 { return []uint32{} }  result := make([]uint32, len(arr)-1) for i := 0; i < len(arr)-1; i++ { result[i] = arr[i+1] - arr[i] } return result }  //go:noescape func calcDifferences(src []uint32, dst []uint32)  // CalculateGrowth - ускоренная реализация на SIMD func CalcDifferencesASM(arr []uint32) []uint32 { if len(arr) <= 1 { return []uint32{} }  if len(arr) <= 4 { return CalcDifferencesGo(arr) }  result := make([]uint32, len(arr)-1) calcDifferences(arr, result) return result }

Как это работает

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

«Хвост» данных, которые не войдут полностью в векторный регистр, обрабатываем скалярно.

Полученные результаты для длин гистограмм от 10 до 100000
Ускорение в ~2 раза (а хотелось бы в 4~8)

goos: windows goarch: amd64 pkg: qwe cpu: AMD Ryzen 5 8400F 6-Core Processor Benchmark_profile/Go-10-12               67875990                19.03 ns/op Benchmark_profile/SIMD-10-12             66618182                18.76 ns/op Benchmark_profile/Go-100-12              13702243                92.77 ns/op Benchmark_profile/SIMD-100-12            15909126                77.92 ns/op Benchmark_profile/Go-1000-12              1294166                954.4 ns/op Benchmark_profile/SIMD-1000-12            1835912                636.4 ns/op Benchmark_profile/Go-10000-12              144406                 7867 ns/op Benchmark_profile/SIMD-10000-12            277306                 4469 ns/op Benchmark_profile/Go-100000-12              19545                63121 ns/op Benchmark_profile/SIMD-100000-12            43426                28094 ns/op

Ускорение func reverseProfile чуть сложнее

Нам нужно не только делать векторное вычитание, но и перестановку элементов в векторе.

// func calcReverseDifferences(src []int32, dst []int32) // комменты от GPT-4.1 TEXT ·calcReverseDifferences(SB), $0-48     MOVQ src_base+0(FP), SI     // SI = &src[0]     MOVQ src_len+8(FP), CX      // CX = len(src)     MOVQ dst_base+24(FP), DI    // DI = &dst[0]          CMPQ CX, $1                 // Проверяем минимальную длину     JLE  done          DECQ CX                     // CX = len(dst) = len(src)-1     MOVQ CX, AX                 // AX будет индексом в dst (идем с конца)     XORQ DX, DX                 // DX = индекс в src (идем с начала)      // Проверяем, достаточно ли элементов для AVX     CMPQ CX, $8     JL   scalar_loop_start      // AVX-обработка (блоками по 8 элементов)     MOVQ CX, R8     ANDQ $-8, R8                // R8 = количество элементов, кратных 8     MOVQ R8, R9                 // R9 = счетчик обработанных элементов  avx_loop:     // Загружаем текущие 8 элементов     VMOVDQU (SI)(DX*4), Y0      // Y0 = src[DX..DX+7]     // Загружаем следующие 8 элементов (сдвиг на 1)     VMOVDQU 4(SI)(DX*4), Y1     // Y1 = src[DX+1..DX+8]          // Вычисляем разницы     VPSUBD Y0, Y1, Y2           // Y2 = разницы          // Реверсируем порядок элементов     VPERM2I128 $0x01, Y2, Y2, Y3  // Меняем местами 128-битные половины     VPSHUFD $0x1B, Y3, Y3       // Реверсируем порядок в каждой половине          // Сохраняем в обратном порядке     MOVQ AX, R10     SUBQ $8, R10                // R10 = AX-7 (начальная позиция для записи)     VMOVDQU Y3, (DI)(R10*4)     // Сохраняем 8 элементов          ADDQ $8, DX                 // Увеличиваем счетчик src     SUBQ $8, AX                 // Уменьшаем счетчик dst     CMPQ DX, R8                 // Проверяем завершение AVX-цикла     JL   avx_loop  scalar_loop_start:     DECQ AX scalar_loop:     CMPQ DX, CX                 // Проверяем завершение     JGE  done          MOVL (SI)(DX*4), BX         // Загружаем src[i]     MOVL 4(SI)(DX*4), BP        // Загружаем src[i+1]     SUBL BX, BP                 // Вычисляем разницу     MOVL BP, (DI)(AX*4)         // Сохраняем в dst          INCQ DX                     // Следующий элемент src     DECQ AX                     // Предыдущая позиция dst     JMP  scalar_loop  done:     RET 

// Стандартная реализация на Go для проверки func CalcReverseDifferencesGo(arr []uint32) []uint32 { if len(arr) <= 1 { return []uint32{} }  result := make([]uint32, len(arr)-1) for i := len(arr) - 1; i > 0; i-- { result[len(arr)-1-i] = arr[i] - arr[i-1] } return result }  //go:noescape func calcReverseDifferences(src []uint32, dst []uint32)  // CalculateGrowth - обертка, которая создает массив назначения func CalReversecDifferencesASM(arr []uint32) []uint32 { if len(arr) <= 1 { return []uint32{} }  if len(arr) <= 4 { return CalcReverseDifferencesGo(arr) }  result := make([]uint32, len(arr)-1) calcReverseDifferences(arr, result) return result }

Все также как и с calcDifferences но другой порядок работы с индексами

VPERM2I128 $0x01, Y2, Y2, Y3

Эта инструкция берет два 256-битных регистра (в данном случае оба — Y2) и переставляет их 128-битные половины.

VPSHUFD $0x1B, Y3, Y3

Эта инструкция переставляет 32-битные элементы внутри каждой 128-битной половины регистра Y3.

Иллюстрация

Полученные результаты ускорения идентичны calcDifferences

goos: windows goarch: amd64 pkg: qwe cpu: AMD Ryzen 5 8400F 6-Core Processor Benchmark_reverse_profile/Go-10-12                57465208             20.66 ns/op Benchmark_reverse_profile/ASM-10-12               61966186             19.28 ns/op Benchmark_reverse_profile/Go-100-12               11149742             107.8 ns/op Benchmark_reverse_profile/ASM-100-12              15162312             76.59 ns/op Benchmark_reverse_profile/Go-1000-12               1000000              1022 ns/op Benchmark_reverse_profile/ASM-1000-12              1756675             679.6 ns/op Benchmark_reverse_profile/Go-10000-12               129410              9519 ns/op Benchmark_reverse_profile/ASM-10000-12              208200              4928 ns/op Benchmark_reverse_profile/Go-100000-12               15706             76804 ns/op Benchmark_reverse_profile/ASM-100000-12              38287             30263 ns/op

Итого совокупный прирост для набора гистограмм с 200000 парами составил

2.80 раза (180.07% быстрее)

func TestHistogramProcessing(t *testing.T) { // Инициализация генератора случайных чисел seed := time.Now().UnixNano() rand.Seed(seed) t.Logf("Используем seed: %d", seed)  const lengthHistograms = 10_000 const numHistograms = 10_000 const numPair = 20 const numRandom = 10000  var Histogram [][]uint32  for i := 0; i < numHistograms; i++ { qwe, asd := generatePairedHistograms(lengthHistograms, numPair) Histogram = append(Histogram, asd...) Histogram = append(Histogram, qwe) }  // Создаем случайных гистограмм for i := 0; i < numRandom; i++ { Histogram = append(Histogram, generateRandomHistogram(lengthHistograms)) }  // Замер времени для processHistograms start1 := time.Now() result1Paired := processHistograms(Histogram) duration1 := time.Since(start1)  // Замер времени для processHistogramsAVX start2 := time.Now() result2Paired := processHistogramsAVX(Histogram) duration2 := time.Since(start2)  if result1Paired != result2Paired { t.Errorf("Результаты для парных гистограмм не совпадают: processHistograms=%d, processHistogramsAVX=%d", result1Paired, result2Paired) } else { fmt.Printf("Результаты для парных гистограмм совпадают: %d \n", result1Paired) t.Logf("Результаты для парных гистограмм совпадают: %d", result1Paired) }  // Выводим замеры времени t.Logf("Время выполнения processHistograms: %v", duration1) t.Logf("Время выполнения processHistogramsAVX: %v", duration2)  // Вычисляем во сколько раз одна функция быстрее другой var faster, slower time.Duration var fasterName, slowerName string if duration1 < duration2 { faster, slower = duration1, duration2 fasterName, slowerName = "processHistograms", "processHistogramsAVX" } else { faster, slower = duration2, duration1 fasterName, slowerName = "processHistogramsAVX", "processHistograms" }  speedup := float64(slower) / float64(faster) percentDiff := (speedup - 1) * 100  t.Logf("%s быстрее %s примерно в %.2f раза (%.2f%% быстрее)", fasterName, slowerName, speedup, percentDiff)  fmt.Printf("%s быстрее %s примерно в %.2f раза (%.2f%% быстрее)\n", fasterName, slowerName, speedup, percentDiff) }