habrahabr

Я ускорил генерацию blurhash в 3̶6̶ 8̶7̶ 128 раз

  • понедельник, 14 октября 2024 г. в 00:00:08
https://habr.com/ru/articles/850114/

Старую собаку новым трюкам не обучишь, вот и я взялся за старое. Blurhash — это компактный способ представления размытой превьюшки изображения в виде ASCII-строки. Разработан финской компанией Wolt (аналог Delivery Club). Давно хотелось внедрить такое к себе в API, чтобы любой клиент мог более плавно и изящно делать загрузку контент на своем сайте. Но сколько я на него смотрел — всегда не давала покоя скорость работы, уж больно медленно и «в лоб» он был написан. Но вот время пришло наконец-то разобраться, что же он так медленно работает.

Основной репозиторий — https://github.com/woltapp/blurhash, в нём лежит референсная имплементация аж на четырех языках (Си, Swift, Kotlin, TS). Но я начал своё путешествие с библиотеки для Пайтона.

0. Интерфейс Cи <=> Пайтон

Для начала, давайте разберемся как вообще устроена библиотека. В источниках видно, что взаимодействие c библиотекой реализована через интерфейс cffi, но реализовано немного нетипичным образом: вместо реального интерфейса к какой‑либо системной библиотеке, в проект почти дословно скопированы файлы decode.c и encode.c и других зависимостей нет. Такой подход имеет право на жизнь для небольших библиотек с простым интерфейсом, коей в данном случае и является blurhash.

Сишная функция create_hash_from_pixels принимает:

  • Количество компонентов хэша (другими словами его разрешение) x_components и y_components

  • Размер передаваемого ей изображения width и height вместе с самим изображением rgb.

  • Инкремент в байтах для перехода на следующую строку bytes_per_row — на случай если строка занимает больше места, чем нужно для её размещения. Хозяйке на заметку: такой параметр при необходимости позволяет за бесплатно получить «кроп» изображения. Для этого нужно лишь сместить начальный указатель на первый пиксель кропнутого изображения и указать bytes_per_row соответствующий не кропнутому.

  • Строка для результата работы destination.

Со стороны Пайтона вызов происходит так:

def encode(image, x_components, y_components):
    if not isinstance(image, Image.Image):
        image = Image.open(image)
    if image.mode != 'RGB':
        image = image.convert('RGB')
    red_band = image.getdata(band=0)
    green_band = image.getdata(band=1)
    blue_band = image.getdata(band=2)
    rgb_data = list(chain.from_iterable(zip(red_band, green_band, blue_band)))
    width, height = image.size
    image.close()

    rgb = _ffi.new('uint8_t[]', rgb_data)
    bytes_per_row = _ffi.cast('size_t', width * 3)
    destination = _ffi.new("char[]", 2 + 4 + (9 * 9 - 1) * 2 + 1)

    result = _lib.create_hash_from_pixels(
        x_components, y_components,
        width, height, rgb, bytes_per_row,
        destination,
    )

Тут уже можно схватиться за волосы: все пиксели изображения итерируются, потом из них создается список, который уже помещается в сишный массив rgb. Давайте замерим, как это работает:

from blurhash import encode, Image
im = Image.open('/Code/imgs/bologna2k.jpg').resize((360, 240), Image.BICUBIC)
assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
%timeit encode(im, 6, 4)

Я вставил assert, чтобы убедиться, что ничего не сломаю в процессе оптимизаций. Уже тут нас ждет неприятный сюрприз:

ValueError: Operation on closed image

Что это значит? А это тот самый image.close() на 11-й строке. Ну, ок, пока что обойдем это, создавая копию изображения перед каждым вызовом функции:

   ...: assert encode(im.copy(), 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im.copy(), 6, 4)
126 ms ± 995 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

126 мс на крошечное изображение разрешением 360x240 пикселей — это, конечно, мощно. Хорошо, что на самом деле нам не нужно получать никакие отдельные каналы изображения, чтобы снова их соединить. Чтобы получить представление картинки в виде куска памяти в Pillow есть метод .tobytes():

-    red_band = image.getdata(band=0)
-    green_band = image.getdata(band=1)
-    blue_band = image.getdata(band=2)
-    rgb_data = list(chain.from_iterable(zip(red_band, green_band, blue_band)))
+    rgb_data = image.tobytes()
   ...: assert encode(im.copy(), 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im.copy(), 6, 4)
109 ms ± 487 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Вот уже 15% прироста. Заодно избавляемся от лишних импортов из itertools и six (что?).

Но что же там с image.close()? На самом деле он не бесполезен, т.к. функция может принимать как готовое изображение, так и открывать его из файла. И кстати тут присутствует баг, что если открытое изображение не было в режиме RGB, то закроется не открытое изображение, а конвертированное (строка номер 5). Однако закрывать готовое изображение явно не требуется. Исправить это можно, задав контекст:

from contextlib import nullcontext

def encode(image, x_components, y_components):
    if isinstance(image, Image.Image):
        image_context = nullcontext()
    else:
        image = Image.open(image)
        image_context = image
    with image_context:
        if image.mode != 'RGB':
            image = image.convert('RGB')
        rgb_data = image.tobytes()
        width, height = image.size

    rgb = _ffi.new('uint8_t[]', rgb_data)
    ...

Остальной код остался прежним. Но теперь можно убрать явный вызов .copy() перед вызовом encode()

In [2]: from blurhash import encode, Image
   ...: im = Image.open('/Code/imgs/bologna2k.jpg').resize((360, 240), Image.BICUBIC)
   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
108 ms ± 30 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Большого ускорения это не дало, зато посмотрите, насколько стабильнее стали результаты без лишнего копирования в памяти.

1. Кешируем всё что можно

Давайте наконец-то взглянем на encode.c. Тут вся работа (и 99.9% времени выполнения) происходит в функции multiplyBasisFunction. Она вызывается для каждого пикселя результирующего изображения (для каждого значения xComponents и yComponents).

static struct RGB multiplyBasisFunction(int xComponent, int yComponent, int width, int height, uint8_t *rgb, size_t bytesPerRow) {
	struct RGB result = { 0, 0, 0 };
	float normalisation = (xComponent == 0 && yComponent == 0) ? 1 : 2;

	for(int y = 0; y < height; y++) {
		for(int x = 0; x < width; x++) {
			float basis = cosf(M_PI * xComponent * x / width) * cosf(M_PI * yComponent * y / height);
			result.r += basis * sRGBToLinear(rgb[3 * x + 0 + y * bytesPerRow]);
			result.g += basis * sRGBToLinear(rgb[3 * x + 1 + y * bytesPerRow]);
			result.b += basis * sRGBToLinear(rgb[3 * x + 2 + y * bytesPerRow]);
		}
	}

	float scale = normalisation / (width * height);

	result.r *= scale;
	result.g *= scale;
	result.b *= scale;

	return result;
}

Что тут происходит? Высчитывается какой-то коэффициент basis и умножается на каждый пиксель. Ба, да это же косинусное преобразование! То есть по сути каждый blurhash — это миниатюрный JPEG без заголовков и с фиксированной таблицей квантования. Интересный подход. Ещё из интересного, вычисления выполняются в линейном цветовом пространстве (функция sRGBToLinear). И это первый кандидат на серьезную оптимизацию тут. Дело в том, что пиксели изображения могут содержать только 256 разных значений. А вот самих пикселей даже в нашем примере 360 * 240 * 3 = 259 200, в тысячу раз больше. То есть намного выгоднее было бы заранее составить таблицу всех возможных ответов функции sRGBToLinear и в последствии уже обращаться к ней. Но можно пойти ещё дальше: дело в том, что значения эти не изменятся, сколько раз не вызывай sRGBToLinear. А значит, что можно посчитать эту таблицу при первом вызове encode и закешировать навечно. В памяти такая таблица займет всего 256 * 4 = 1024 байта.

float *sRGBToLinear_cache = NULL;

static void init_sRGBToLinear_cache() {
	if (sRGBToLinear_cache != NULL) {
		return;
	}
	sRGBToLinear_cache = (float *)malloc(sizeof(float) * 256);
	for (int x = 0; x < 256; x++) {
		sRGBToLinear_cache[x] = sRGBToLinear(x);
	}
}

static struct RGB multiplyBasisFunction(int xComponent, int yComponent, int width, int height, uint8_t *rgb, size_t bytesPerRow) {
	struct RGB result = { 0, 0, 0 };

    init_sRGBToLinear_cache();

	for(int y = 0; y < height; y++) {
		for(int x = 0; x < width; x++) {
			float basis = cosf(M_PI * xComponent * x / width) * cosf(M_PI * yComponent * y / height);
			result.r += basis * sRGBToLinear_cache[rgb[3 * x + 0 + y * bytesPerRow]];
			result.g += basis * sRGBToLinear_cache[rgb[3 * x + 1 + y * bytesPerRow]];
			result.b += basis * sRGBToLinear_cache[rgb[3 * x + 2 + y * bytesPerRow]];
		}
	}
    ...
}

Проверяем:

   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
14.2 ms ± 6.37 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Воу-воу-воу! Мы только начали, а уже ускорили весь код в 8.87 раз! Впрочем, это только начало.

Взгляните, как вычисляется basis. Дело в том, что cosf — довольно тяжелая функция, она выполняется ни один такт (а скорее всего, даже не десять). Но значения косинусов зависят только от счетчиков, которые раз за разом проходят одни и те же значения. К счастью, компилятор достаточно умен, чтобы не считать cos(y) для каждого пикселя. Но вот с cos(x) он сам не в состоянии справиться. Давайте ему поможем:

static struct RGB multiplyBasisFunction(int xComponent, int yComponent, int width, int height, uint8_t *rgb, size_t bytesPerRow) {
	struct RGB result = { 0, 0, 0 };
	float *cosx = (float *)malloc(sizeof(float) * width);

	for(int x = 0; x < width; x++) {
		cosx[x] = cosf(M_PI * xComponent * x / width);
	}

	for(int y = 0; y < height; y++) {
		float cosy = cosf(M_PI * yComponent * y / height);
		uint8_t *src = rgb + y * bytesPerRow;
		for(int x = 0; x < width; x++) {
			float basis = cosy * cosx[x];
			result.r += basis * sRGBToLinear_cache[src[3 * x + 0]];
			result.g += basis * sRGBToLinear_cache[src[3 * x + 1]];
			result.b += basis * sRGBToLinear_cache[src[3 * x + 2]];
		}
	}

	free(cosx);
    ...
In [1]: from blurhash import encode, Image
   ...: im = Image.open('/Code/imgs/bologna2k.jpg').resize((360, 240), Image.BICUBIC)
   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
3.45 ms ± 3.25 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

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

2. Делаем всё за один проход

Для чего нужна функция multiplyBasisFunction? Она считает ровно один коэффициент косинусного преобразования. Делает она это считая basis для каждого пикселя исходного изображения. Но вот в чем штука: для одинаковых xComponent и yComponent наборы посчитанных косинусов будет одинаковыми. А мы их все равно считаем каждый раз. Следовательно, можно вычислять их заранее и передавать в функцию уже готовый массив.

Сама функция при этом станет только проще:

static struct RGB multiplyBasisFunction(
	int xComponent, int yComponent, int width, int height, uint8_t *rgb, size_t bytesPerRow,
	float *cosx, float *cosy
) {
	struct RGB result = { 0, 0, 0 };

	for(int y = 0; y < height; y++) {
		uint8_t *src = rgb + y * bytesPerRow;
		for(int x = 0; x < width; x++) {
			float basis = cosy[y] * cosx[x];
			result.r += basis * sRGBToLinear_cache[src[3 * x + 0]];
			result.g += basis * sRGBToLinear_cache[src[3 * x + 1]];
			result.b += basis * sRGBToLinear_cache[src[3 * x + 2]];
        }
    }
    ...
}

А вот вызывающая её blurHashForPixels заметно преобразится:

const char *blurHashForPixels(int xComponents, int yComponents, int width, int height, uint8_t *rgb, size_t bytesPerRow, char *destination) {
	if(xComponents < 1 || xComponents > 9) return NULL;
	if(yComponents < 1 || yComponents > 9) return NULL;

	float factors[9 * 9][3];

	init_sRGBToLinear_cache();

	float *cosx = (float *)malloc(sizeof(float) * width * xComponents);
	if (! cosx) return NULL;
	float *cosy = (float *)malloc(sizeof(float) * height);
	if (! cosy) {
		free(cosx);
		return NULL;
	}
	for(int x = 0; x < xComponents; x++) {
		for(int i = 0; i < width; i++) {
			cosx[x * width + i] = cosf(M_PI * x * i / width);
		}
	}
	for(int y = 0; y < yComponents; y++) {
		for(int i = 0; i < height; i++) {
			cosy[i] = cosf(M_PI * y * i / height);
		}
		for(int x = 0; x < xComponents; x++) {
			struct RGB factor = multiplyBasisFunction(x, y, width, height, rgb, bytesPerRow,
				cosx + x * width, cosy);
			factors[y * xComponents + x][0] = factor.r;
			factors[y * xComponents + x][1] = factor.g;
			factors[y * xComponents + x][2] = factor.b;
		}
	}
	free(cosx);
	free(cosy);

Здесь мы заранее подготавливаем все коэффициенты.

   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
3.4 ms ± 3.65 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

К сожалению прироста от этого практически никакого, я и так уже очень серьезно сократил количество вычислений косинусов. Значит, нужно придумать что-то ещё.

Немного подумав, можно заметить, что не только косинусы между разными вызовами multiplyBasisFunction не меняются, но и пиксели изображения. В текущем коде мы проходим по изображению и извлекаем значения из sRGBToLinear_cache 24 раза (6 * 4) для каждого канала. Какой в этом смысл? Правильно — никакого. Значит пришла пора для самого крупного изменения в коде и логике. Нужно сделать все вычисления за один проход.

Здесь возникает дилемма: если сделать в лоб и поместить обход xComponent и yComponent внутрь каждого пикселя, то получится 4 уровня вложенности циклов, причем внутренние циклы будут иметь очень небольшую дистанцию. Я не проверял, как это будет работать, но из опыта знаю, что код может сильно замедлиться из-за постоянных ошибок предсказателя ветвлений процессора.

Я решил попробовать другой подход: схлопнуть два внутренних цикла в один и не различать отдельно xComponent и yComponent, а сделать один плоский массив. Но тогда возникает вопрос: а как нам на каждом шаге узнать значения этих xComponent и yComponent, чтобы получить правильные индексы для косинусов? Целочисленное деление и взятие модуля от счетчика — тоже не самая быстрая операция. Я решил, что намного проще будет просто продублировать нужные коэффициенты столько раз, сколько у нас xComponents и yComponents. То есть каждый массив cosX и cosY будет иметь размерность не width * xComponents, а width * factorsCount, где factorsCount = xComponents * yComponents. Это не увеличит количество вычислений, а лишь незначительно увеличит использование памяти. В общем итоговый вид функции multiplyBasisFunction стал таким:

static void multiplyBasisFunction(
	struct RGB *factors, int factorsCount, int width, int height, uint8_t *rgb, size_t bytesPerRow,
	float *cosX, float *cosY
) {
	for(int y = 0; y < height; y++) {
		uint8_t *src = rgb + y * bytesPerRow;
		float *cosYLocal = cosY + y * factorsCount;
		for(int x = 0; x < width; x++) {
			float pixel[3];
			float *cosXLocal = cosX + x * factorsCount;
			pixel[0] = sRGBToLinear_cache[src[3 * x + 0]];
			pixel[1] = sRGBToLinear_cache[src[3 * x + 1]];
			pixel[2] = sRGBToLinear_cache[src[3 * x + 2]];
			for (int i = 0; i < factorsCount; i++) {
				float basis = cosYLocal[i] * cosXLocal[i];
				factors[i].r += basis * pixel[0];
				factors[i].g += basis * pixel[1];
				factors[i].b += basis * pixel[2];
			}
		}
	}

	for (int i = 0; i < factorsCount; i++) {
		float normalisation = (i == 0) ? 1 : 2;
		float scale = normalisation / (width * height);
		factors[i].r *= scale;
		factors[i].g *= scale;
		factors[i].b *= scale;
	}
}
   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
1.44 ms ± 950 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Новый рекорд 87.5 раз!

3. SIMD (SSE + NEON)

Наконец-то моя любимая тема. Вот сейчас начнется жара — думал я.

Я всегда писал SIMD только под x86, но тут у меня была цель наконец-то попробовать сделать полноценную поддержку и ARM (все предыдущие тесты я запускал под эмуляцией Rosetta на M1 Pro).

Изменений для SSE потребовалось минимальное количество. И кстати, эта версия действительно SSE (самой первой версии), то есть запустится на самом древнем процессоре, вроде бы начиная с Pentium III.

static void multiplyBasisFunction(
	float factors[][4], int factorsCount, int width, int height, uint8_t *rgb, size_t bytesPerRow,
	float *cosX, float *cosY
) {
	for(int y = 0; y < height; y++) {
		uint8_t *src = rgb + y * bytesPerRow;
		float *cosYLocal = cosY + y * factorsCount;
		int x = 0;
		for(; x < width; x++) {
			float *cosXLocal = cosX + x * factorsCount;
#ifdef __SSE__
			__m128 pixel = _mm_set_ps(0, sRGBToLinear_cache[src[3 * (x+0) + 2]],
				sRGBToLinear_cache[src[3 * (x+0) + 1]], sRGBToLinear_cache[src[3 * (x+0) + 0]]);
			for (int i = 0; i < factorsCount; i++) {
				__m128 basis = _mm_set1_ps(cosYLocal[i] * cosXLocal[i + 0 * factorsCount]);
				__m128 factor = _mm_loadu_ps(factors[i]);
				factor = _mm_add_ps(factor, _mm_mul_ps(basis, pixel));
				_mm_storeu_ps(factors[i], factor);
			}
#else
			float pixel[4];
			pixel[0] = sRGBToLinear_cache[src[3 * x + 0]];
			pixel[1] = sRGBToLinear_cache[src[3 * x + 1]];
			pixel[2] = sRGBToLinear_cache[src[3 * x + 2]];
			for (int i = 0; i < factorsCount; i++) {
				float basis = cosYLocal[i] * cosXLocal[i];
				factors[i][0] += basis * pixel[0];
				factors[i][1] += basis * pixel[1];
				factors[i][2] += basis * pixel[2];
			}
#endif
		}
	}
    ...
}

Я получил свой заслуженный прирост:

   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
973 μs ± 1.13 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Это, кстати, те самые обещанные 128 раз! Наивный же код, без использования интринсиков NEON и без Rosetta, выполнялся на M1 pro c ещё большей скоростью (эти измерения не идут в общий зачет):

   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
808 μs ± 454 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Эквивалент кода под NEON получился вполне прямолинейным:

#elif defined(__aarch64__)
			float32x4_t pixel = {sRGBToLinear_cache[src[3 * (x+0) + 0]],
				sRGBToLinear_cache[src[3 * (x+0) + 1]], sRGBToLinear_cache[src[3 * (x+0) + 2]]};
			for (int i = 0; i < factorsCount; i++) {
				float32x4_t basis = vdupq_n_f32(cosYLocal[i] * cosXLocal[i + 0 * factorsCount]);
				float32x4_t factor = vld1q_f32(factors[i]);
				factor = vmlaq_f32(factor, basis, pixel);
				vst1q_f32(factors[i], factor);
			}
#else

Радостный я запустил бенчмарк:

   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
962 μs ± 24.5 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Опачки. Опачки. Код замедлился до 0.84 от исходной скорости :-/ Даже несмотря на операцию сдвоенного умножения и сложения vmlaq_f32.

На самом деле к этому моменту я уже давно перешел в основной репозиторий blurhash и тестировал код на чистом Си, это в статье я продолжаю пользоваться blurhash-python. Поэтому мне не составила труда сдампить объектный файл encode.o ванильной версии и посмотреть, какие оптимизации применял компилятор, оказавшийся умнее меня. Оказалось, что компилятор не просто разворачивает внутренний цикл по factorsCount, он ещё и читал коэффициенты инструкциями ld4, которые загружали гораздо больше данных. (Кстати, тут я снова нарвался на баг в Докере с замедлением этих инструкций, если включена Rosetta).

Чтож, попытка оказалась провальной. Можно было бы конечно разворачивать циклы в своем коде тоже, подсматривая у компилятора. Но вместо того, чтобы пытаться соперничать с ним, я решил пойти другим путём, я решил посмотреть что будет, если попытаться ему помочь.

4. Развязываем компилятору руки

Первое, что можно было сделать — выровнять структуры данных по 4 элемента, чтобы компилятор мог свободно распоряжаться ими как одним вектором. Для этого я отказался от struct RGB и перешел на просто массив из 4 элементов.

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

#ifdef __SSE__
			__m128 pixel = _mm_set_ps(0, sRGBToLinear_cache[src[3 * (x+0) + 2]],
				sRGBToLinear_cache[src[3 * (x+0) + 1]], sRGBToLinear_cache[src[3 * (x+0) + 0]]);
			for (int i = 0; i < factorsCount; i++) {
				__m128 basis = _mm_set1_ps(cosYLocal[i] * cosXLocal[i + 0 * factorsCount]);
				__m128 factor = _mm_loadu_ps(factors[i]);
				factor = _mm_add_ps(factor, _mm_mul_ps(basis, pixel));
				_mm_storeu_ps(factors[i], factor);
			}
#else

Для каждого элемента factorsCount, на каждой итерации внутреннего цикла я загружаю и сохраняю factors[i]. Загружаю и сохраняю, загружаю и сохраняю. А как иначе? Нам же нужно их как-то суммировать. Это так, но что интересно — для следующего пикселя мы снова будем загружать и сохранять эти же значения. Идеально было бы держать эти значения в регистрах, но так сделать не получится, потому что во-первых, число factorsCount переменное. Во-вторых, оно может быть довольно велико, до 9 * 9, никаких регистров не хватит. А вот что можно было бы сделать — загружать и сохранять пореже. Не для каждого пикселя, а для каждого четвертого, например. Это сделать не просто, а очень просто:

static void multiplyBasisFunction(
	float factors[][4], int factorsCount, int width, int height, uint8_t *rgb, size_t bytesPerRow,
	float *cosX, float *cosY
) {
	for(int y = 0; y < height; y++) {
		uint8_t *src = rgb + y * bytesPerRow;
		float *cosYLocal = cosY + y * factorsCount;
		int x = 0;
		for(; x < width - 3; x += 4) {
			float *cosXLocal = cosX + x * factorsCount;
			float pixel0[4] = {sRGBToLinear_cache[src[3 * (x+0) + 0]], sRGBToLinear_cache[src[3 * (x+0) + 1]], sRGBToLinear_cache[src[3 * (x+0) + 2]]};
			float pixel1[4] = {sRGBToLinear_cache[src[3 * (x+1) + 0]], sRGBToLinear_cache[src[3 * (x+1) + 1]], sRGBToLinear_cache[src[3 * (x+1) + 2]]};
			float pixel2[4] = {sRGBToLinear_cache[src[3 * (x+2) + 0]], sRGBToLinear_cache[src[3 * (x+2) + 1]], sRGBToLinear_cache[src[3 * (x+2) + 2]]};
			float pixel3[4] = {sRGBToLinear_cache[src[3 * (x+3) + 0]], sRGBToLinear_cache[src[3 * (x+3) + 1]], sRGBToLinear_cache[src[3 * (x+3) + 2]]};
			for (int i = 0; i < factorsCount; i++) {
				float basis0 = cosYLocal[i] * cosXLocal[i + 0 * factorsCount];
				float basis1 = cosYLocal[i] * cosXLocal[i + 1 * factorsCount];
				float basis2 = cosYLocal[i] * cosXLocal[i + 2 * factorsCount];
				float basis3 = cosYLocal[i] * cosXLocal[i + 3 * factorsCount];
				factors[i][0] += basis0 * pixel0[0] + basis1 * pixel1[0] + basis2 * pixel2[0] + basis3 * pixel3[0];
				factors[i][1] += basis0 * pixel0[1] + basis1 * pixel1[1] + basis2 * pixel2[1] + basis3 * pixel3[1];
				factors[i][2] += basis0 * pixel0[2] + basis1 * pixel1[2] + basis2 * pixel2[2] + basis3 * pixel3[2];
			}
		}
		for(; x < width; x++) {
			float pixel[4];
			float *cosXLocal = cosX + x * factorsCount;
			pixel[0] = sRGBToLinear_cache[src[3 * x + 0]];
			pixel[1] = sRGBToLinear_cache[src[3 * x + 1]];
			pixel[2] = sRGBToLinear_cache[src[3 * x + 2]];
			for (int i = 0; i < factorsCount; i++) {
				float basis = cosYLocal[i] * cosXLocal[i];
				factors[i][0] += basis * pixel[0];
				factors[i][1] += basis * pixel[1];
				factors[i][2] += basis * pixel[2];
			}
		}
	}

В результате компилятор теперь может крутить и оптимизировать этот код как считает нужным и что самое интересное, прекрасно с этим справляется. Результат для rosetta x86 такой же как от ручного использования интринсиков SSE:

   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
975 μs ± 1.81 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Результаты ARM:

   ...: assert encode(im, 6, 4) == 'W7EUfLs:M{8yM{wJ$%MepHrE%LV[NaVtBUXSsBNbVuOrM|NaMeXm'
   ...: %timeit encode(im, 6, 4)
589 µs ± 2.27 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Что на 37% быстрее, чем без последней оптимизации.

Итого

Если честно, очень хотелось добить заветную соточку. Есть ли тут пути ещё ускорить? Я думаю, что есть. Во-первых, можно попробовать перейти на целочисленные вычисления. В разные моменты я пытался это сделать, и это давало кое-какой результат, но код сильно усложняется и когда открывал очередную оптимизацию в основной ветке, было лень каждый раз переделывать на целые числа. Кроме того, я думаю что может получиться все же держать рабочую часть массива factors в регистрах, сделав несколько проходов по изображению. Я попытался, но судя по скорости, компилятор не оценил моей задумки. Нужно точно переходить на интринсики, чтобы контролировать это.

Сейчас у меня открыты два пулреквеста, в основную библиотеку и в пайтоновскую:
https://github.com/woltapp/blurhash/pull/256
https://github.com/woltapp/blurhash-python/pull/25
От авторов пока что ничего не слышно. Если у кого-то есть желание, можете пройтись по остальным реализациям и применить эти или похожие оптимизации. Кроме того, кроме энкодера, есть ещё декодер, который я даже не смотрел. Уверен, что его тоже можно сильно ускорить.

Оптимизация

x86 GCC 12.2.0

ARM Clang 14.0.3

Без

126 ms

55.5 ms

Интерфейс Пайтона

108

1.17x

47.7

1.16x

Кеш значений

3.45

36x

2.65

21x

Один проход

1.44

87x

1.01

55x

Помощь компилятору

0.98

128x

0.59

94x