Как сделать ваш код в 80 раз быстрее
У нас стартует третий набор на курс «Разработчик Python», а значит, что впереди и открытый урок, которые у нас частично замещают староформатные дни открытых дверей и где можно ознакомиться с интересным материалом от наших преподавателей, и то, что мы нашли очередной интересный материальчик. На этот раз по ускорению «змеиного» кода.
PyPy способен ускорить код в 2 раза, что радует очень многих людей. Хочу поделиться короткой, личной историей, доказывающей, что PyPy способен на большее.
ДИСКЛЕЙМЕР: это не чудодейственное средство на все случаи жизни, да, сработало конкретно в этом случае, но может оказаться не таким эффективным во многих других. Однако метод все равно интересный. Более того, шаги, описанные здесь, я применял во время разработки в том же порядке, что делает статью жизненным примером оптимизации PyPy.
Я экспериментировал с эволюционными алгоритмами несколько месяцев назад: план был амбициозным — автоматически развить логику, способную контролировать (симулированный) квадрокоптер, то есть PID-регулятор (спойлер: не летает).
Идея заключается в том, что при наличии популяции случайных существ, в каждом поколении наиболее приспособленные выживают и размножаются с небольшими, случайными изменениями.
Тем не менее, в рамках этого поста, первоначальная задача не так важна, поэтому перейдем непосредственно к коду. Для управления квадрокоптером существо использует метод run_step , запускаемый в каждом delta_t (код полностью):
- inputs — numpy массив, в котором находятся финальная точка и текущая позиция по оси Z;
- outputs — numpy массив, в котором находится тяга, передающаяся моторам. Для начала, все 4 мотора ограничены одинаковой тягой, поэтому квадрокоптер перемещается вверх-вниз по оси Z;
- self.state содержит произвольные значения неизвестно размера, которые передаются из одного шага в другой;
- self.matrix и self.constant содержат саму логику. При передачи в них “правильных” значений, теоретически, мы могли бы получить идеально настроенный PID-регулятор. Они случайным образом мутируют между поколениями.
Сначала я попробовал запустить код на CPython; и вот результат:
6.15 секунд/поколение, за исключением первого.
Затем, я попробовал PyPy 5.9:
Ой! Получилось в
5.5 раз медленнее, чем с CPython. Отчасти это было ожидаемо: nympy основан на cpyext, известном своей медлительностью. (На самом деле, мы работаем над этим — на бранче cpyext-avoid-roundtrip результат уже лучше, чем с CPython, но поговорим об этом в другом посте).
Попробуем избежать cpyext. Первым очевидным шагом станет использование numpypy вместо numpy (вот хак, позволяющий использовать только часть micronumpy). Проверим, улучшилась ли скорость:
2.7 секунды: в 12 раз быстрее, чем PyPy+numpy и более чем в 2 раза быстрее, чем оригинальный CPython. Уже сейчас многие бы побежали рассказывать, как хорош PyPy.
Но попробуем улучшить результат. Как обычно, проанализируем, что именно требует больше всего времени. Вот ссылка на vmprof профиль. Мы тратим много времени внутри numpypy и выделяем тонны временных массивов для хранения результатов различных операций.
Кроме того, посмотрим на трассировку jit и найдем функцию run: это цикл, в котором мы проводим большую часть времени, он состоит из 1796 операций. Операции для строки np.dot(. ) + self.constant находятся между строками 1217 и 1456. Ниже приведен отрывок, вызывающий np.dot(. ) . Большая часть операторов ничего не стоят, но в строке 1232 мы видим вызов функции RPython descr_dot; по реализации видим, что она создает новый W_NDimArray для хранения результата, а значит придется делать malloc() :
Интересная реализация части + self.constant — вызов W_NDimArray.descr_add был встроен JIT, поэтому нам проще понять, что происходит. В частности, мы видим вызов к __0_alloc_with_del____ , выделяющий W_NDimArray для результата, и aw_malloc, выделяющий сам массив. Затем идет длинный список из 149 простых операций, которые задают поля итогового массива, создают итератор и, наконец, вызывают acall_assembler — это фактическая логика для выполнения добавления, которую JIT’или отдельно. call_assembler одна из операций для выполнения JIT-to-JIT вызовов:
Все это не очень оптимально: в нашем случае известно, что размер self.matrix всегда равен (3, 2) — а значит мы выполняем огромное количество работы, включая 2 вызова malloc() для временных массивов, просто чтобы вызвать две функции, которые в общей сумме делают 6 умножений и 6 сложений. Отметим, что это не вина JIT: CPython+numpy приходится делать то же самое, но в скрытых вызовах C.
Решить эту проблему вероятно поможет хорошо известная оптимизация компилятора — разворот цикла. С точки зрения компилятора, разворот цикла всегда сопровождается риском — если матрица слишком большая, вы, в конечном итоге, останетесь с большой, бесформенной массой кода, которая будет бесполезна, если размер массива постоянно меняется. Это одна из основных причин, почему PyPy JIT в этом случае даже не пытается так поступить.
Тем не менее, мы знаем, что матрица маленькая и всегда одного размера. Поэтому развернем цикл вручную:
В коде дополнительно есть проверка работоспособности, чтобы убедиться, что значение равно выданному через Creature.run_step .
Проверим, как оно работает. Сначала с CPython:
Выглядит неплохо — на 60% быстрее, чем оригинальная реализация CPython+numpy. Проверим PyPy:
Да, это не ошибка. Спустя пару поколений, время стабилизируется в районе
0.07-0.08 секунды за поколение. Что примерно в 80 (восемьдесят) раз быстрее, чем реализация CPython+numpy, и в 35-40 раз быстрее наивной PyPy+numpy.
Еще раз посмотрим на трассировку: в ней больше нет дорогих вызовов и точно никаких временных malloc()’ов. Корень логики находится в строках 386-416, где мы видим выполнение быстрых C-level умножений и сложений: float_mul и float_add переводятся прямо в команды mulsd и addsdx86.
Как я и говорил, это очень конкретный пример, и такой метод не универсален: к сожалению, не стоит ждать 80-кратного ускорения произвольного кода. Однако это ясно показывает потенциал PyPy в высокоскоростных вычислениях. И самое важное, это не пробный бенчмарк, спроектированный специально для хорошей работы PyPy — пример небольшой, но реалистичный.