У нас стартует третий набор на курс «Разработчик Python», а значит, что впереди и открытый урок, которые у нас частично замещают староформатные дни открытых дверей и где можно ознакомится с интересным материалом от наших преподавателей, и то, что мы нашли очередной интересный материальчик. На этот раз по ускорению «змеиного» кода :)
Поехали.
PyPy способен ускорить код в 2 раза, что радует очень многих людей. Хочу поделиться короткой, личной историей, доказывающей, что PyPy способен на большее.
ДИСКЛЕЙМЕР: это не чудодейственное средство на все случаи жизни, да, сработало конкретно в этом случае, но может оказаться не таким эффективным во многих других. Однако метод все равно интересный. Более того, шаги, описанные здесь, я применял во время разработки в том же порядке, что делает статью жизненным примером оптимизации PyPy.
Я экспериментировал с эволюционными алгоритмами несколько месяцев назад: план был амбициозным — автоматически развить логику, способную контролировать (симулированный) квадрокоптер, то есть PID-регулятор (спойлер: не летает).
Идея заключается в том, что при наличии популяции случайных существ, в каждом поколении наиболее приспособленные выживают и размножаются с небольшими, случайными изменениями.
Тем не менее, в рамках этого поста, первоначальная задача не так важна, поэтому перейдем непосредственно к коду. Для управления квадрокоптером существо использует метод run_step, запускаемый в каждом delta_t (код полностью):
class Creature(object):
INPUTS = 2 # z_setpoint, current z position
OUTPUTS = 1 # PWM for all 4 motors
STATE_VARS = 1
...
def run_step(self, inputs):
# state: [state_vars ... inputs]
# out_values: [state_vars, ... outputs]
self.state[self.STATE_VARS:] = inputs
out_values = np.dot(self.matrix, self.state) + self.constant
self.state[:self.STATE_VARS] = out_values[:self.STATE_VARS]
outputs = out_values[self.STATE_VARS:]
return outputs
- inputs — numpy массив, в котором находятся финальная точка и текущая позиция по оси Z;
- outputs — numpy массив, в котором находится тяга, передающаяся моторам. Для начала, все 4 мотора ограничены одинаковой тягой, поэтому квадрокоптер перемещается вверх-вниз по оси Z;
self.stateсодержит произвольные значения неизвестно размера, которые передаются из одного шага в другой;self.matrixиself.constantсодержат саму логику. При передачи в них “правильных” значений, теоретически, мы могли бы получить идеально настроенный PID-регулятор. Они случайным образом мутируют между поколениями.
run_step вызывается при 100 Гц (в виртуальном временном интервале симуляции). Поколение состоит из 500 существ, каждое из которых мы тестируем в течение 12 виртуальных секунд. Таким образом, в каждом поколении содержится 600,000 выполнений run_step.
Сначала я попробовал запустить код на CPython; и вот результат:
$ python -m ev.main
Generation 1: ... [population = 500] [12.06 secs]
Generation 2: ... [population = 500] [6.13 secs]
Generation 3: ... [population = 500] [6.11 secs]
Generation 4: ... [population = 500] [6.09 secs]
Generation 5: ... [population = 500] [6.18 secs]
Generation 6: ... [population = 500] [6.26 secs]
То есть ~6.15 секунд/поколение, за исключением первого.
Затем, я попробовал PyPy 5.9:
$ pypy -m ev.main
Generation 1: ... [population = 500] [63.90 secs]
Generation 2: ... [population = 500] [33.92 secs]
Generation 3: ... [population = 500] [34.21 secs]
Generation 4: ... [population = 500] [33.75 secs]
Ой! Получилось в ~5.5 раз медленнее, чем с CPython. Отчасти это было ожидаемо: nympy основан на cpyext, известном своей медлительностью. (На самом деле, мы работаем над этим — на бранче cpyext-avoid-roundtrip результат уже лучше, чем с CPython, но поговорим об этом в другом посте).
Попробуем избежать cpyext. Первым очевидным шагом станет использование numpypy вместо numpy (вот хак, позволяющий использовать только часть micronumpy). Проверим, улучшилась ли скорость:
$ pypy -m ev.main # using numpypy
Generation 1: ... [population = 500] [5.60 secs]
Generation 2: ... [population = 500] [2.90 secs]
Generation 3: ... [population = 500] [2.78 secs]
Generation 4: ... [population = 500] [2.69 secs]
Generation 5: ... [population = 500] [2.72 secs]
Generation 6: ... [population = 500] [2.73 secs]
В среднем, ~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 в этом случае даже не пытается так поступить.
Тем не менее, мы знаем, что матрица маленькая и всегда одного размера. Поэтому развернем цикл вручную:
class SpecializedCreature(Creature):
def __init__(self, *args, **kwargs):
Creature.__init__(self, *args, **kwargs)
# store the data in a plain Python list
self.data = list(self.matrix.ravel()) + list(self.constant)
self.data_state = [0.0]
assert self.matrix.shape == (2, 3)
assert len(self.data) == 8
def run_step(self, inputs):
# state: [state_vars ... inputs]
# out_values: [state_vars, ... outputs]
k0, k1, k2, q0, q1, q2, c0, c1 = self.data
s0 = self.data_state[0]
z_sp, z = inputs
#
# compute the output
out0 = s0*k0 + z_sp*k1 + z*k2 + c0
out1 = s0*q0 + z_sp*q1 + z*q2 + c1
#
self.data_state[0] = out0
outputs = [out1]
return outputs
В коде дополнительно есть проверка работоспособности, чтобы убедиться, что значение равно выданному через Creature.run_step.
Проверим, как оно работает. Сначала с CPython:
$ python -m ev.main
Generation 1: ... [population = 500] [7.61 secs]
Generation 2: ... [population = 500] [3.96 secs]
Generation 3: ... [population = 500] [3.79 secs]
Generation 4: ... [population = 500] [3.74 secs]
Generation 5: ... [population = 500] [3.84 secs]
Generation 6: ... [population = 500] [3.69 secs]
Выглядит неплохо — на 60% быстрее, чем оригинальная реализация CPython+numpy. Проверим PyPy:
Generation 1: ... [population = 500] [0.39 secs]
Generation 2: ... [population = 500] [0.10 secs]
Generation 3: ... [population = 500] [0.11 secs]
Generation 4: ... [population = 500] [0.09 secs]
Generation 5: ... [population = 500] [0.08 secs]
Generation 6: ... [population = 500] [0.12 secs]
Generation 7: ... [population = 500] [0.09 secs]
Generation 8: ... [population = 500] [0.08 secs]
Generation 9: ... [population = 500] [0.08 secs]
Generation 10: ... [population = 500] [0.08 secs]
Generation 11: ... [population = 500] [0.08 secs]
Generation 12: ... [population = 500] [0.07 secs]
Generation 13: ... [population = 500] [0.07 secs]
Generation 14: ... [population = 500] [0.08 secs]
Generation 15: ... [population = 500] [0.07 secs]
Да, это не ошибка. Спустя пару поколений, время стабилизируется в районе ~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 — пример небольшой, но реалистичный.
Как воспроизвести результат
$ git clone https://github.com/antocuni/evolvingcopter
$ cd evolvingcopter
$ {python,pypy} -m ev.main --no-specialized --no-numpypy
$ {python,pypy} -m ev.main --no-specialized
$ {python,pypy} -m ev.main
THE END
Если что, то ждём вопросов тут или на открытом уроке.
Комментариев нет:
Отправить комментарий