Путь хакера — решение задачи с codechef на питон. С машинным кодом
В процессе написания Blog:Advanced Algorithms/Python-оптимизация алгоритма динамического программирования из codechef, когда удалось получить питон-решение проходящее таймлимиты на PyPy, возник вопрос-челлендж, можно ли все-таки пройти этот квест на обычном питоне («CPython»).
Казалось, что это невозможно, ведь CPython грубо получает таймлимит больше в 2.5 раза, чем для PyPy (но это совсем небольшая фора), и для PyPy мы еле-еле вписались, а по нашим оценкам, наш же алгоритм на CPython работал на порядок медленней. Возможно ожидалась какая-то магия с использованием модулей numpy или collections, чтобы по сути, дергать питоном функции, написанные на си или даже высокообогащенном оружейном фортране.
Другие методы оптимизации — отдельные скомпилированные модули на С, компиляция Cython или Nuitka, тут явно невыполнимы — на вход принимается только один питон-файл.
Но оказалось, есть еще и путь хакера.
import copy import ctypes import mmap import sys DIGITS = 19 DP_SIZE = DIGITS * 2 * 2 def get_input(): lines = sys.stdin.readlines() data = ' '.join(lines).strip() for token in data.split(): yield token func_bytes = (b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1' b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L' b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~' b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t' b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83' b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51' b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b' b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f' b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00' b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I' b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95' b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$' b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe' b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D' b'\x0f\xb6T$/\xe9M\xff\xff\xff') assert len(func_bytes) <= mmap.PAGESIZE func_buffer = mmap.mmap(-1, mmap.PAGESIZE, prot=mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC) func_type = ctypes.CFUNCTYPE(ctypes.c_longlong, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.POINTER(ctypes.c_longlong), ctypes.POINTER(ctypes.c_longlong)) func_pointer = ctypes.c_void_p.from_buffer(func_buffer) func = func_type(ctypes.addressof(func_pointer)) func_buffer.write(func_bytes) token_generator = get_input() next(token_generator) dp = (ctypes.c_longlong * DP_SIZE)(*(-1 for _ in range(DP_SIZE))) p10 = (ctypes.c_longlong * DIGITS)(*(10 ** i for i in range(DIGITS))) for input_str in token_generator: N = len(input_str) input_digits = (ctypes.c_int * N)(*(int(ch) for ch in input_str)) print(func(0, 0, 0, input_digits, N, copy.copy(dp), p10))
Идея — будем действовать как JIT. Мы можем переписать кусок кода на си, взять получившийся машинный код, внутри питона его записать в динамически выделенную память и запустить.
Мы можем выделить память с разрешением на чтение, запись и исполнение. Запись понятно для чего, а разрешение на исполнение означает, что процессор может сделать jmp или call в эту область и не крешнуться (разрешение на чтение казалось бы не нужно, но как ни странно, на x86-64 для исполнения также нужно разрешение на чтение):
func_buffer = mmap.mmap(-1, mmap.PAGESIZE, prot=mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC)
Дальше мы пользуемся библиотекой ctypes в питоне — она позволяет из питона вызывать сишные функции. Для этого нам нужно указать какие аргументы наша функция будет принимать и что возвращать (чтобы питоновский код знал как нужно передавать аргументы перед тем как сделать call в наш код, а также как распарсить ответ):
func_type = ctypes.CFUNCTYPE(ctypes.c_longlong, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.POINTER(ctypes.c_longlong), ctypes.POINTER(ctypes.c_longlong))
Вот как это соответствует нашей сишной функции:
int64_t calculate(int i, int taken, int ls, int *input_digits, int N, int64_t *dp, int64_t *p10)
Дальше нам нужно создать объект сишной функции и записать байты (о них дальше) в выделенную память:
func_pointer = ctypes.c_void_p.from_buffer(func_buffer) func = func_type(ctypes.addressof(func_pointer)) func_buffer.write(func_bytes)
И после всего этого мы можем вызывать эту функцию как обычно, только нужно передавать аргументы нужных типов:
dp = (ctypes.c_longlong * DP_SIZE)(*(-1 for _ in range(DP_SIZE))) p10 = (ctypes.c_longlong * DIGITS)(*(10 ** i for i in range(DIGITS))) for input_str in token_generator: N = len(input_str) input_digits = (ctypes.c_int * N)(*(int(ch) for ch in input_str)) print(func(0, 0, 0, input_digits, N, copy.copy(dp), p10))
Теперь о том как получить нужный машинный код и как вообще выбрать что компилировать.
Наши возможности в том что мы можем запустить довольно ограничены — мы не можем запустить что-либо что использует какие-либо библиотеки, потому что мы загружаем не «честный ELF», в котором можно подгружать библиотеки, делать релокации и многое другое, а только набор машинных инструкций. Это значит что мы не можем использовать ничего из libc — то есть без ввода/вывода, динамического выделения памяти, ничего из стандартной библиотеки плюсов — то есть без каких-либо структур данных сложнее статического массива.
Поэтому идеальный кандидат — функция calculate из оригинального кода на питоне. Она как раз использует только массивы фиксированного размера и простые операции:
def calculate(i, taken, ls): """ i — индекс цифры taken — встречали ли простые цифры раньше ls — в текущем индексе можно делать полный перебор """ global input_digits global dp global N limit = 9 if not ls: limit = input_digits[i] res = 0 for x in range(limit + 1): newtaken = taken or (x == 2) or (x == 3) or (x == 5) or (x == 7) newi = i + 1 if newi == N: ret = newtaken else: newls = ls or (x != limit) if newls and newtaken: ret = p10[N - newi] else: cached = dp[newtaken][newls][newi] if cached != -1: ret = cached else: ret = calculate(newi, newtaken, newls) res += ret dp[taken][ls][i] = res return res
API ctypes позволяет нам запускать сишные функции — поэтому ее и напишем:
// file: inline.c # include <stdint.h> int64_t calculate(int i, int taken, int ls, int *input_digits, int N, int64_t *dp, int64_t *p10) { int limit = 9; if (!ls) { limit = input_digits[i]; } int64_t res = 0; for (int x = 0; x <= limit; ++x) { int new_taken = taken || (x == 2) || (x == 3) || (x == 5) || (x == 7); int new_i = i + 1; int64_t ret; if (new_i == N) { ret = new_taken; } else { int new_ls = ls || (x != limit); if (new_ls && new_taken) { ret = p10[N - new_i]; } else { int64_t cached = dp[4 * new_i + 2 * new_taken + new_ls]; if (cached != -1) { ret = cached; } else { ret = calculate(new_i, new_taken, new_ls, input_digits, N, dp, p10); } } } res += ret; } dp[4 * i + 2 * taken + ls] = res; return res; }
(глобальные переменные мы создаем на стороне питона и передаем как аргументы, также мы преобразовали многомерный массив dp в одномерный, чтобы не париться с указателями)
Дальше нам нужно это скомпилировать, и нужно передать некоторые нестандарные флаги. Обычно исполнение ELF начинается не с main, а с функции _start, которая вызывает функцию __libc_start_main из libc, которая подгружает переменные окружения, argv и многое другое, а только потом запускает main. Нам ничего такого не нужно, поэтому мы это отключаем с помощью флага -nostartfiles, и к тому же без этого компилятор ожидал бы от нас функцию main, которую мы писать не хотим. Также можем добавить -nolibc чтобы не было соблазна вызывать функции оттуда, хотя на результат это не повлияет. И конечно -O3 чтобы получить все оптимизации что есть.
$ gcc inline.c -nostartfiles -nolibc -o inline -O3
Можем проверить что в секции .text у нас лежит одна единственная функция (-d — вывод дизассемблера секции с кодом, -M intel — чтобы получить вывод в намного более приятном интеловском синтаксисе, а не в AT&T):
$ objdump -d -M intel inline inline: file format elf64-x86-64 Disassembly of section .text: 0000000000001000 <calculate>: 1000: 41 57 push r15 1002: 41 89 f7 mov r15d,esi 1005: 41 56 push r14 1007: 41 89 fe mov r14d,edi 100a: 41 55 push r13 100c: 41 89 d5 mov r13d,edx 100f: 41 54 push r12 1011: 55 push rbp 1012: 53 push rbx 1013: 48 83 ec 38 sub rsp,0x38 1017: 85 d2 test edx,edx 1019: 75 35 jne 1050 <calculate+0x50> 101b: 48 63 c7 movsxd rax,edi 101e: 45 31 e4 xor r12d,r12d 1021: 8b 2c 81 mov ebp,DWORD PTR [rcx+rax*4] 1024: 85 ed test ebp,ebp 1026: 79 2d jns 1055 <calculate+0x55> 1028: 43 8d 04 77 lea eax,[r15+r14*2] 102c: 41 8d 44 45 00 lea eax,[r13+rax*2+0x0] ... ...
Также стоит заметить что наш код не использует никаких глобальных переменных, а значит нам нужна только секция .text из ELF. И еще все инструкции jmp и call используют относительные оффсеты для прыжков, поэтому наш код может быть загружен по любому адресу памяти, что нам и нужно, так как мы не можем полностью контроллировать какой адрес вернет mmap.
Нам осталось сдампить байты нашей функции в отдельный файл (-O binary = вывести сырые байты, -j .text = работать с секцией .text):
$ objcopy -O binary -j .text inline inline_bytes
И записать получившиеся байты внутрь кода в питоне, например так можно получить строку байт чтобы скопировать:
from pprint import pprint with open('inline_bytes', 'rb') as f: pprint(f.read())
Какой будет вывод:
(b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1' b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L' b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~' b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t' b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83' b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51' b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b' b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f' b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00' b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I' b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95' b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$' b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe' b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D' b'\x0f\xb6T$/\xe9M\xff\xff\xff')
И это нужно скопировать, чтобы можно было в codechef послать только один файл:
func_bytes = (b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1' b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L' b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~' b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t' b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83' b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51' b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b' b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f' b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00' b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I' b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95' b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$' b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe' b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D' b'\x0f\xb6T$/\xe9M\xff\xff\xff')
Вот и все!
[ Хронологический вид ]Комментарии
Войдите, чтобы комментировать.