Мой первый супероптимизатор

Настало время для очередного бесполезного проекта.

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

Смысл такой: сгенерировать все возможные пермутации инструкций кода и протестировать каждую полученную программу на равнозначность исходной. Вот, по сути, и всё. Несложно представить, что область возможных программ очень быстро разрастается, да и протестировать две программы на равнозначность тоже нелегко. Но, если эта задача была выполнима на компьютерах в 1987 году, то мой ноутбук определённо с ней справится.

Примечание: обсуждение этой статьи есть на Hacker News.

Я попробую создать примитивный супероптимизатор.

Чтобы сделать его реалистичным, я придумаю собственный язык ассемблера, содержащий всего несколько инструкций. Для проверки равнозначности двух программ мне потребуется написать эмулятор, выполняющий эти программы по порядку. Затем оптимизатор, исходя из заданного набора ограничений, будет генерировать на этом языке все возможные варианты программ и возвращать самую короткую из них.

Я разбил этот проект на несколько частей:

• проектирование простого языка ассемблера;
• написание эмулятора, выполняющего программу ассемблера и возвращающего её конечное состояние;
• написание функции, тестирующей равнозначность двух программ на основе их состояния до/после;
• написание языка, способного выполнять преобразование между понятной человеку и внутренней формой программы ассемблера;
• написание оптимизатора, генерирующего каждую программу длиной от 1 до n инструкций со всеми возможными операндами от 0 до k и выводящего первые x ячеек памяти.

Исходный код проекта лежит на GitHub.

▍ Язык ассемблера

Хочу ограничить его всего несколькими инструкциями, которые окажутся достаточно интересными, чтобы супероптимизатор мог находить неочевидные для меня оптимизации.
Напомню, что моя цель — не создание полноценного языка ассемблера, а простое изучение супероптимизаторов.

Для начала я произвольно выбрал такие инструкции:

• LOAD val, загружающую непосредственное значение в область памяти 0;
• SWAP mem, mem, меняющую местами значения в двух областях памяти;
• XOR mem, mem, выполняющую побитовое ИЛИ для значений в областях памяти и сохраняющую результат в первой области;
• INC mem, инкрементирующую значение области памяти.

Вы увидите, что все они довольно скучные. Среди них даже нет инструкции для потока управления. Хотя позднее можно будет без проблем добавить и другие.

▍ Эмулятор

Написать эмулятор на удивление легко. Состояние памяти представляет список чисел, а программа — список инструкций с операндами. Мы выполняем все инструкции и в качестве результата возвращаем конечное состояние памяти.

class CPU:
    def __init__(self, max_mem_cells):
        self.max_mem_cells = max_mem_cells
        self.state = [0] * max_mem_cells
        self.ops = {'LOAD': self.load, 'SWAP': self.swap, 'XOR': self.xor, 'INC': self.inc}

    def execute(self, program):
        state = self.state.copy()
        for instruction in program:
            op = instruction[0]
            args = list(instruction[1:])
            args.insert(0, state)
            state = op(*args)
        return state

Инструкция — это ссылка на функцию, выполняющую операцию:

 def load(self, state, val):
        state[0] = val
        return state
    
    def swap(self, state, mem1, mem2):
        state[mem1], state[mem2] = state[mem2], state[mem1]
        return state
    
    def xor(self, state, mem1, mem2):
        state[mem1] = state[mem1] ^ state[mem2]
        return state
    
    def inc(self, state, mem):
        state[mem] += 1
        return state

Это код эмулятора.

Чтобы упростить его использование, я также реализовал функцию, получающую читаемую форму кода ассемблера и преобразующую её в программу, а также обратную функцию, получающую программу и преобразующую её в читаемый код. Ничего особо интересного в них нет, но посмотреть эти функции можно здесь: assembler.py.

▍ Оптимизатор

А теперь самое весёлое — генерация всех возможных программ.

class Superoptimizer:
    def generate_programs(self, cpu, max_length, max_mem, max_val):
        for length in range(1, max_length + 1):
            for prog in product(cpu.ops.values(), repeat=length):
                arg_sets = []
                for op in prog:
                    if op == cpu.load:
                        arg_sets.append([tuple([val]) for val in range(max_val + 1)])
                    elif op == cpu.swap or op == cpu.xor: 
                        arg_sets.append(product(range(max_mem), repeat=2))
                    elif op == cpu.inc:
                        arg_sets.append([tuple([val]) for val in range(max_mem)])
                for arg_set in product(*arg_sets):
                    program = [(op, *args) for op, args in zip(prog, arg_set)] 
                    yield program

Пусть эта функция вас не пугает. Она генерирует все возможные варианты программы на основе нескольких переменных: длины программы, количества инструкций и числа возможных операндов (размер памяти или максимальное значение). Я превратил её в генератор, так как первая версия съедала всю память моего ноутбука.

Какая здесь вычислительная сложность? Ужасная.

Для тестирования каждой программы мы будем использовать следующий код:

 def search(self, max_length, max_mem, max_val, target_state):
        cpu = CPU(max_mem)
        for program in self.generate_programs(cpu, max_length, max_mem, max_val):
            state = cpu.execute(program)
            if state == target_state:
                state = tuple(state) 
                if state not in self.program_cache or len(program) < len(self.program_cache[state]):
                    self.program_cache[state] = program
        return self.program_cache.get(tuple(target_state), None)

Функция search получает в качестве параметров ограничения и целевое состояние. Она оценивает две программы как равнозначные, если состояние памяти обеих оказывается идентичным. Позже это ограничение можно будет ослабить. Она сохраняет кратчайшую программу с корректным состоянием и по завершении поиска возвращает её.

▍ Попробуем

Пришло время протестировать наш супероптимизатор. Если передать ему программу, сможет ли он найти кратчайший её вариант с тем же конечным результатом?

Моей первой программой стал выдуманный язык ассемблера:

LOAD 3
SWAP 0, 1
LOAD 3
SWAP 0, 2
LOAD 3
SWAP 0, 3
LOAD 3

Состояние завершившейся программы представлено как [3, 3, 3, 3]. Она заполняет конечную область памяти тройками.

Каков будет её оптимальный код, согласно супероптимизатору?

LOAD 3
XOR 1, 0
XOR 2, 0
XOR 3, 0

Работает!

Супероптимизатор вместо LOAD и SWAP использует инструкцию XOR для повтора значений и их сохранения в подходящей области. Очень круто, я об этом не подумал.

А теперь ещё один тест.

Также можно передать супероптимизатору только желаемое конечное состояние без написания программы, его создающей. То есть я могу попросить его найти кратчайшую программу, которая приведёт к состоянию [0, 2, 1]. Вот её оптимальный вариант:

LOAD 2
SWAP 0, 1
INC 2

Для 2 он задействовал паттерн из LOAD и SWAP, но для 1 использовал INC, так как для этого нужно на одну инструкцию меньше. Снова победа!

▍ Производительность

Производительность здесь никак не назовёшь хорошей. Поиск конечного состояния [1, 0, 1, 0, 1, 0] при максимум 6 инструкциях и значениях в пределах 2 привёл к тому, что супероптимизатор за 30+ минут сгенерировал 1,580,000,000 программ, после чего я его просто закрыл.

▍ Дальнейшая работа

Здесь можно многое доработать:

• Начальное состояние. Эмулятор предполагает, что стартовое состояние всегда одинаковое: 0 для всех областей памяти. Но в реальности нужно допустить различные стартовые состояния, чтобы программа могла получать ввод инструкции fib(n), вычисляющей n номер в последовательности Фибоначчи. Это можно сделать, добавив в эмулятор возможность получения параметра начального состояния, устанавливающего первое подмножество значений в памяти.
• Равнозначность программ. Для более полноценной оценки равнозначности программ, включая их ввод и вывод, например, для отражения, что fib(n) равнозначна opt_fib(n), также нужно ввести примечание вывода. В качестве него может выступить ещё один параметр, определяющий вывод как первые x ячеек памяти на момент завершения программы. Помимо этого, нам нужна функция, показывающая, что равнозначность относится к нескольким входным критериям. Будет, пожалуй, достаточно тестировать для каждого ввода предустановленный набор целых чисел, а не их все.
• Отсечение лишнего. Супероптимизатор перебирает невероятное множество программ, многие из которых не имеют смысла. К примеру, я достиг пика в тысячи сгенерированных программ, которые циклически выполняли XOR нулей и перезаписывали неиспользуемые значения. Это усложнит функцию генерации, но зато значительно сократит пространство поиска.
• Дополнительные инструкции. Располагая ограниченным набором инструкций, супероптимизатор не сможет находить более интересные паттерны. Первым делом я бы добавил какую-нибудь условную инструкцию, чтобы она могла воздействовать на ввод. Например, CMP mem, mem, которая устанавливает область памяти 0 на конкретное значение в зависимости от того, больше/меньше/равно ли значение первого операнда значению второго. Далее я бы добавил всё необходимое для реализации классических функций вроде fib(n) и max(a, b), возможно, ADD и SUB. Классические компиляторы отлично справляются с заменой условных выражений арифметическими, так что это будет хорошим тестом для супероптимизаторов. Будьте осторожны с эмуляцией переходов, поскольку в этом случае не будете знать, завершится ли когда-нибудь программа.

Исходный код проекта ищите на GitHub.

Автор: Дмитрий Брайт

Источник