Путь хакера — решение задачи с codechef на питон. С машинным кодом

Материал из DISCOPAL
Перейти к: навигация, поиск

В процессе написания 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\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA' \
             b'\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00' \
             b'\x00A\x8d~\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8tk\x8d\x04?1\xdb' \
             b'E1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H' \
             b'\x0f\xa3\xd8sq9\xddu51\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b\x04' \
             b'\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f\x1fD\x00\x00M\x03#\xeb\xea' \
             b'\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3' \
             b'\xc2s\x04I\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95\xc2E\x85\xed\x0f' \
             b'\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D' \
             b'\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L' \
             b'\x8b\\$ D\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]
    ...
    ...

Нам осталось сдампить байты нашей функции в отдельный файл (-O binary = вывести сырые байты, -j .text = работать с секцией .text):

$ objcopy -O binary -j .text inline inline_bytes

И записать получившиеся байты внутрь кода в питоне, например так можно получить строку байт чтобы скопировать:

# with open('inline_bytes', 'rb') as f:
#     print(f.read())
 
func_bytes = b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA' \
             b'\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00' \
             b'\x00A\x8d~\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8tk\x8d\x04?1\xdb' \
             b'E1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H' \
             b'\x0f\xa3\xd8sq9\xddu51\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b\x04' \
             b'\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f\x1fD\x00\x00M\x03#\xeb\xea' \
             b'\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3' \
             b'\xc2s\x04I\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95\xc2E\x85\xed\x0f' \
             b'\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D' \
             b'\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L' \
             b'\x8b\\$ D\x0f\xb6T$/\xe9M\xff\xff\xff'

[ Хронологический вид ]Комментарии

(нет элементов)

Войдите, чтобы комментировать.