haqreu 16 фев в 21:37

Смогу ли я уложить оптимизирующий компилятор в тысячу строк питона? Прогон первый: mem2reg

Средний

22 мин

3.7K

Алгоритмы*Программирование*Python*Компиляторы*Математика*

Туториал

Год назад мне пришлось взять на себя курс лекций по теории компиляторов. Вы встречались некомпетентными преподавателями? Это я, здравствуйте! Прежде чем учить других, я всё-таки решил заглянуть в учебник сам, и это вылилось в серию статей "компилятор за выходные" (да, я помню, что за мной должок с описанием лексера/парсера). В итоге я уложил компилятор со мной придуманного си-подобного языка на GNU ассемблер в шестьсот строк кода, причём без внешних зависимостей, включая парсинг.

Всё бы хорошо, вроде работает, но кажется, самое веселье осталось за бортом. Мой компилятор, по факту, это простой pretty print вокруг синтаксического дерева, подумаешь. А как работают оптимизирующие компиляторы? И поставил я себе задачу попробовать уложить игрушечный, но всё же рабочий оптимизирующий компилятор в тысячу строк кода. Как думаете, получится?

Итак, тема сегодняшнего разговора - вынос переменных из памяти в регистры, оно же оптимизационный проход mem2reg, см. кпдв.

Что у нас уже есть?

Оглавление цикла

А есть у нас компилятор, который умеет компилировать программы на си-подобном языке в ассемблерный код. Давайте я приведу простенький пример программы на моём языке:

main() {
    int collatz(int n) {
        int its;

        its = 0;
        while (n != 1) {
            its = its + 1;
            if (n % 2 == 0) {
                n = n / 2;
            } else {
                n = 3*n + 1;
            }
        }
        return its;
    }
    println collatz(27);
}

К слову, я потратил сильно больше времени на написание тестовых программ, нежели на непосредственно компилятор :) Загляните в репозиторий, там есть несколько демок, рисующих анимации в консоли (если что, там и си-версии демок есть на всякий случай).

Язык очень простой, в нём нет ни массивов, ни указателей. Только два типа данных - целочисленный и булевский. Но just for fun там в моём языке есть вложенные функции. В принципе, как можно видеть по демкам, вполне достаточно для веселья.

Принцип работы крайне прост. Берётся файл с исходником на wend, и при помощи лексера превращается в поток лексем, который разбирается парсером, у которого есть в наличии грамматика языка. Парсер строит синтаксическое дерево, которое пропускается через семантический анализатор со всякими проверками типов. Заодно семантический анализатор добавляет дополнительную информацию в узлы дерева (вроде связывания инструкций использования переменных с их объявлением). Ну а затем генератор ассемблера проходит по этому украшенному дереву, выдавая конечный код.

Вот полностью код главного файла компилятора, видно, что я просто последовательно запускаю вышеупомянутые этапы:

import io, sys
from lexer import WendLexer
from parser import WendParser
from analyzer import decorate
from transasm import transasm

if len(sys.argv)!=2:
    sys.exit('Usage: compiler.py path/source.wend')
try:
    f = open(sys.argv[1], 'r')
    tokens = WendLexer().tokenize(f.read())
    ast = WendParser().parse(tokens)
    decorate(ast)
    print(transasm(ast))
except Exception as e:
    print(e)

Всё это хозяйство у меня заняло шесть сотен строк питона:

ssloy@periwinkle:~/tinycompiler$ cat *py|wc -l
608

Итого, если я хочу уложиться в тысячу, у меня осталось четыре сотни, чтобы утоптать оптимизацию, хотя бы самые базовые вещи. Как думаете, получится? :)

Поехали! Промежуточное представление

Для своей работы компилятор должен иметь какую-то структуру данных, которая позволяет представлять код. На данный момент в качестве такого промежуточного (между wend и ассемблером) представления я использую синтаксическое дерево. В принципе, какие-то оптимизации возможны и с такой структурой, но всё же есть лучшие альтернативы. Вариантов масса (и я в них совершенно не разбираюсь), но тут в комментариях к моим статьям упомянули любопытный цикл статей Максима Казанцева ( @xortator), а он специалист по LLVM, в котором для представления кода используется промежуточный язык LLVM IR.

Почитал я эти статьи, и решил, что LLVM IR как нельзя лучше подходит для моего компилятора. Мой проект не имеет внешних зависимостей, и я не собираюсь полагаться на LLVM, но при этом иметь возможность сохранить на диск моё промежуточное представление и непосредственно её запустить при помощи внешнего инструмента - это бесценно!

Давайте посмотрим, что сделает clang с вот таким файлом:

int max(int a, int b) {
    int result = a;
    if (b > a) {
        result = b;
    }
    return result;
}

Запускаем clang -cc1 -O0 -emit-llvm -disable-O0-optnone max.c -o max.ll и получаем файл следующего содержания:

define dso_local i32 @max(i32 noundef %a, i32 noundef %b) #0 {
entry:
  %a.addr = alloca i32, align 4
  %b.addr = alloca i32, align 4
  %result = alloca i32, align 4
  store i32 %a, ptr %a.addr, align 4
  store i32 %b, ptr %b.addr, align 4
  %0 = load i32, ptr %a.addr, align 4
  store i32 %0, ptr %result, align 4
  %1 = load i32, ptr %b.addr, align 4
  %2 = load i32, ptr %a.addr, align 4
  %cmp = icmp sgt i32 %1, %2
  br i1 %cmp, label %if.then, label %if.end

if.then:                                          ; preds = %entry
  %3 = load i32, ptr %b.addr, align 4
  store i32 %3, ptr %result, align 4
  br label %if.end

if.end:                                           ; preds = %if.then, %entry
  %4 = load i32, ptr %result, align 4
  ret i32 %4
}

Это и есть промежуточное представление кланга о нашей функции int max(int, int). Его можно отрисовать чуть более наглядно при помощи того же llvm. Запускаем вот такой скрипт, который отрисует все .ll файлы в текущей папке:

#!/bin/bash

for ll in `ls -1 *.ll`; do
    opt -passes=dot-cfg -disable-output -cfg-dot-filename-prefix=`basename $ll .ll` $ll
done

for f in `ls -1 *.dot`; do
    dot -Tpng $f -O
done

И получаем наш int max(int, int):

В принципе, это тот же самый файл с промежуточным кодом, но тут явно виден граф потока управления, который, собственно, и является структурой, представляющей код в недрах компилятора. Промежуточный код LLVM - это своего рода ассемблерный код, который можно разбить на базовые блоки с линейным кодом. Каждый базовый блок заканчивается инструкцией ветвления br (ну или обычным ret), и их можно нарисовать рёбрами между базовыми блоками. Видов инструкций крайне немного, и всё действительно очень похоже на ассемблер, только для виртуальной машины, а не для настоящей.

Обратите внимание, что кланг сгенерировал максимально дубовый код, абсолютно никаких оптимизаций. Для всех трёх переменных a, b и result он выделил место на стеке (инструкция alloca), и любое обращение к ним происходит явно при помощи инструкций load и store. Разумеется, это страшно неэффективно, но кланг ничуть это не беспокоит, поскольку улучшение кода - это задача LLVM, а не clang.

Tinycompiler -> tinyoptimizer

Ну и славно, tinycompiler генерирует дубовый ассемблерный код, а поскольку байткод LLVM IR от ассемблера не особо отличается, мне ничто не мешает просто сменить темплейты целевого языка с GNU assembly на LLVM IR. Поскольку llvm позволяет запускать напрямую файлы своего промежуточного языка, то я пока что выкину ассемблер, и добавлю его в самом конце, когда закончу играть с оптимизациями.

Я не имею ни малейшего представления, куда меня заведёт эта кривая дорожка, да и я хотел бы оставить tinycompiler минималистичным, так что я его форкнул в tinyoptimizer. К слову сказать, какого чёрта на гитхабе нельзя форкнуть свой собственный репозиторий? Мне пришлось создать клон, а не форк...

Вот тут коммит, в котором я сменил темплейт языка. По пути я сделал ещё небольшое изменение: поскольку ассемблер не имеет переменных, то я к ним обращался при помощи смещения на стеке, а у LLVM IR вполне себе есть нормальные идентификаторы и передача параметров функциям (посмотрите на пример, что я привёл ранее), почему бы ими не воспользоваться? Остаётся вопрос с доступом к переменным, которые были объявлены вне функции. В tinycompiler я к ним обращался при помощи дисплеев, но сейчас просто добавил к параметрам функций указатели на внешние переменные, и дело с концом.

Давайте приведу пример. Среди моих тестовых программ есть вот такая:

main() {

    // Sum of prime factors (with repetition) of a number n.
    // E.g., 40 factors as 2^3 * 5 and sopfr(40) = 2 * 3 + 5 = 11.
    // It is also known as the integer logarithm.
    // https://oeis.org/A001414

    int sopfr(int n) {
        int div;
        int sopfr_aux(int n) {
            int recursion;
            recursion = 0;
            if n % div == 0 {
                recursion = div;
                if n!=div {
                    recursion = recursion + sopfr_aux(n/div);
                }
            } else {
                div = div + 1;
                recursion = sopfr_aux(n);
            }
            return recursion;
        }

        div = 2;
        return sopfr_aux(n);
    }

    println sopfr(42);
}

Для этой программы на моём языке я генерирую промежуточный код LLVM, который был бы сгенерирован клангом из вот такого сишного эквивалента:

#include <stdio.h>

int sopfr_aux(int n, int *div) {
    int recursion = 0;
    if (n % *div == 0) {
        recursion = *div;
        if (n!=*div) {
            recursion = recursion + sopfr_aux(n / *div, div);
        }
    } else {
        *div = *div + 1;
        recursion = sopfr_aux(n, div);
    }
    return recursion;
}

int sopfr(int n) {
    int div = 2;
    return sopfr_aux(n, &div);
}

void main() {
    printf("%d\n", sopfr(42));
}

Я преобразовываю вложенные функции в обычные, и прокидываю указатели на необходимые переменные. Таким образом, sopfr_aux автоматически получает указатель на переменную div, которая живёт на фрейме стека функции sopfr. Разумеется, я привёл сишный эквиалент просто для читабельности, я из wend генерирую промежуточный код llvm напрямую.

Тема сегодняшней статьи: mem2reg

Наконец-то мы добрались до темы этой статьи. Сегодня мы будем разбираться с прогоном mem2reg. Давайте снова рассмотрим функцию нахождения максума двух чисел. Кланг её преобразует в промежуточный код, который показан справа на рисунке. А теперь давайте отложим кланг и попросим llvm его чуточку улучшить. Не до конца, лишь один явно указанный проход, который занимается выносом переменных из памяти в регистры.

Запускаем opt -passes=mem2reg max.ll -S, и получаем преобразованный код, показанный на картинке слева. Что любопытно, все инструкции alloca, load и store убраны начисто, а вместо них появилась странная инструкция phi. Что же это такое?

Давайте посмотрим, как её определяют сами разработчики llvm :)

Промежуточный код LLVM основан на форме статического единственного присваивания (SSA). Основная идея построения SSA-формы заключается в выделении уникальных имен всем результатам присваивания. Практически все программы содержат ветвления и узлы слияния. В узлах слияния необходимо добавлять специальную форму присваивания, называемую $\phi$ -функцией.

$\phi$ -функция всегда исполняется на входе в базовый блок, и она имеет следующий вид: $x \gets \phi(x_1, x_2, \dots, x_n)$ где — это количество предшествующих узлов в графе потока управления. Семантика $\phi$ -функций проста: выбирается значение источника , соответствующее блоку , из которого был осуществлен переход управления, и присваивается переменной. Если управление переходит в блок от его -го предшественника, то получает значение . Ещё раз, все $\phi$ -функции выполняются до выполнения обычных инструкций в этом узле. Если хотите, можете почитать соответствующую статью @xortator.

То есть, на входе в блок if.end регистр %result.0 получает значение %b, если поток управления проходит через блок if.then, и значение %a, если не проходит. Получился такой своебразный тернарный условный оператор вместо интенсивной работы со стеком.

Зачем нужны $\phi$ -функции? Посмотрите на приведённые примеры и обратите внимание, что в SSA-форме каждый регистр получает значение только один раз. Каждый раз, когда мы в SSA пишем инструкцию вида %a = ..., это как если бы на Си нам разрешали бы использовать оператор присваивания только для инициализации постоянных переменных типа const int a = .... Это не означает, что во время исполнения программы a не будет содержать разных значений. Каждый базовый блок можно рассматривать как маленькую функцию, в которой мы объявляем локальную константу. Вызывая функции с разными аргументами, мы получим разные константы a. Это же означает, что мы можем работать локально с блоком и не должны беспокоиться о том, что где-то вниз по графу регистр будет переопределён. Определение только одно, и от него мы и пляшем.

На всякий случай, как можно представить цикл, в котором счётчик явно должен изменяться? А давайте посмотрим!

Я запустил clang и потом оптимизационный проход mem2reg на тривиальном коде счётчика. Кланг честно положил переменную i на стек, и досчитал до десяти. LLVM же при выносе переменных из памяти в регистры убрал начисто alloca/load/store, и добавил одну $\phi$ -функцию на входе в блок while.cond. Если при исполнении программы мы заходим в этот блок из блока entry, тогда %i.0 получит нулевое значение. Если же туда зайти из блока while.body, то %i.0 получит значение %inc. Таким образом, все виртуальные регистры честно имеют только одну определяющую инструкцию, но их непосредственное значение зависит от того, каким путём мы до неё дошли.

Существует несколько способов для этого графа сделать вывод, что ret всегда возвращает 10, и что можно вообще выкинуть всё это барахло, заинлайнив вызов функции int loop() константой 10, но об этом мы поговорим когда-нибудь в другой раз.

В принципе, эти рассуждения можно пытаться делать и на других промежуточных представлениях кода, но SSA-формы на данный момент обрели наибольшую популярность. И оптимизационный проход mem2reg, который выносит переменные из памяти в регистры SSA, играет в оптимизации ключевую роль. Разумеется, поскольку никакой ассемблер не имеет аналогов $\phi$ -функций, рано или поздно нам придётся вызвать обратный проход reg2mem, заменив $\phi$ -функции на alloca/store/load. Но между mem2reg и reg2mem будут другие оптимизационные проходы, в том числе и вывод информации о константах, а также многое-многое другое.

Давайте же уже писать код, или ненавидимые мной числа Фибоначчи

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

int fib(int n) {
    int a = 0;
    int b = 1;
    int c = 0;
    int i = 1;
    if (n == 0) return 0;
    while (i < n) {
        c = b;
        b = a + b;
        a = c;
        i = i + 1;
    }
    return b;
}

Если мы попросим кланг сгенерировать промежуточный код для неё, то получим такой граф потока управления:

Давайте на некоторое время абстрагируемся от моего компилятора, и напишем независимый код. Какую структуру данных мы можем использовать для представления такого графа потока управления в памяти? Я набросал вот такой файл ir.py:

class Instruction:
    def __init__(self, string, *args):
        self.string, self.op = string, args # string and parameters (typically three)

    def __repr__(self):
        return self.string.format(op = self.op)

class Phi:
    def __init__(self, reg, t, choice=None):
        self.reg, self.t, self.choice = reg, t, choice or {}

    def __repr__(self):
        choice = ', '.join([ '[{percent}{value}, %{block}]'.format(percent = '%' if type(value) is str else '', value=value, block=block) for block, value in self.choice.items() ])
        return f'%{self.reg} = phi {self.t} {choice}'

class BasicBlock:
    def __init__(self, label):
        self.label = label
        self.phi_functions, self.instructions  = [], []
        self.successors, self.predecessors = set(), set()

    def find_and_replace(self, find, replace):
        for i in self.instructions:
            i.op = list( replace if s==find else s for s in i.op )

    def __repr__(self):
        return f'{self.label}:\n' + \
                ''.join([ f'\t{phi}\n' for phi in self.phi_functions ]) + \
                ''.join([ f'\t{i}\n'   for i   in self.instructions  ])

class ControlFlowGraph:                  # a control flow graph is made of basic blocks and
    def __init__(self, header, footer):  # header + footer strings to form a valid LLVM IR program
        self.header, self.footer, self.blocks, self.last = header, footer, {}, None

    def add_block(self, label):
        self.blocks[label] = BasicBlock(label)
        self.last = self.blocks[label]   # the last block added to the CFG, heavily used by the IR builder

    def __repr__(self):
        return f'{self.header}\n' + \
               ''.join([ f'{block}' for block in self.blocks.values() ]) + \
               f'{self.footer}\n'

    def find_and_replace(self, find, replace):
        for b in self.blocks.values():
            b.find_and_replace(find, replace)

    def compute_adjacency(self):
        for b1 in self.blocks.values():
            if any('unreachable' in i.string or 'ret ' in i.string for i in b1.instructions):
                continue # even if there is a br instruction later on in the block, there are no outgoing edges
            for succ in b1.instructions[-1].op[int('br i1' in b1.instructions[-1].string):]:
                b2 = self.blocks[succ]
                b1.successors.add(b2)
                b2.predecessors.add(b1)

Это 55 строчек, при помощи которых я буду манипулировать промежуточным кодом. У меня есть граф потока управления ControlFlowGraph, в котором есть просто список базовых блоков BasicBlock. Каждый базовый блок - это просто список $\phi$ -функций и обычных инструкций, а также список ссылок на предков и потомков данного блока, который изначально пуст, а потом заполняется (compute_adjacency) , исходя из инструкций ветвления. Каждая инструкция - это довольно-таки произвольная строка и список имён регистров и констант, в ней участвующих.

Я взял промежуточный код функции int fib(int), сгенерированный клангом, и руками вколотил его в мою программу:

from ir import *
cfg = ControlFlowGraph('define i32 @fib(i32 %n) {', '}')
cfg.add_block('entry')
cfg.last.instructions = [
    Instruction('%{op[0]} = alloca i32', 'retval'),
    Instruction('%{op[0]} = alloca i32', 'n.addr'),
    Instruction('%{op[0]} = alloca i32', 'a'),
    Instruction('%{op[0]} = alloca i32', 'b'),
    Instruction('%{op[0]} = alloca i32', 'i'),
    Instruction('%{op[0]} = alloca i32', 'c'),
    Instruction('store i32 %{op[0]}, ptr %{op[1]}', 'n', 'n.addr'),
    Instruction('store i32 {op[0]}, ptr %{op[1]}', 0, 'a'),
    Instruction('store i32 {op[0]}, ptr %{op[1]}', 1, 'b'),
    Instruction('store i32 {op[0]}, ptr %{op[1]}', 1, 'i'),
    Instruction('store i32 {op[0]}, ptr %{op[1]}', 0, 'c'),
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '0', 'n.addr'),
    Instruction('%{op[0]} = icmp eq i32 %{op[1]}, {op[2]}', 'cmp', '0', 0),
    Instruction('br i1 %{op[0]}, label %{op[1]}, label %{op[2]}', 'cmp', 'if.then', 'if.end')
]

cfg.add_block('if.then')
cfg.last.instructions = [
    Instruction('store i32 {op[0]}, ptr %{op[1]}', 0, 'retval'),
    Instruction('br label %{op[0]}', 'return')
]

cfg.add_block('if.end')
cfg.last.instructions = [
    Instruction('br label %{op[0]}', 'while.cond')
]

cfg.add_block('while.cond')
cfg.last.instructions = [
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '1', 'i'),
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '2', 'n.addr'),
    Instruction('%{op[0]} = icmp slt i32 %{op[1]}, %{op[2]}', 'cmp1', '1', '2'),
    Instruction('br i1 %{op[0]}, label %{op[1]}, label %{op[2]}', 'cmp1', 'while.body', 'while.end')
]

cfg.add_block('while.body')
cfg.last.instructions = [
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '3', 'b'),
    Instruction('store i32 %{op[0]}, ptr %{op[1]}', '3', 'c'),
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '4', 'a'),
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '5', 'b'),
    Instruction('%{op[0]} = add i32 %{op[1]}, %{op[2]}', 'add', '4', '5'),
    Instruction('store i32 %{op[0]}, ptr %{op[1]}', 'add', 'b'),
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '6', 'c'),
    Instruction('store i32 %{op[0]}, ptr %{op[1]}', '6', 'a'),
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '7', 'i'),
    Instruction('%{op[0]} = add i32 %{op[1]}, {op[2]}', 'add2', '7', 1),
    Instruction('store i32 %{op[0]}, ptr %{op[1]}', 'add2', 'i'),

    Instruction('br label %{op[0]}', 'while.cond')
]
cfg.add_block('while.end')
cfg.last.instructions = [
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '8', 'b'),
    Instruction('store i32 %{op[0]}, ptr %{op[1]}', '8', 'retval'),
    Instruction('br label %{op[0]}', 'return')
]

cfg.add_block('return')
cfg.last.instructions = [
    Instruction('%{op[0]} = load i32, ptr %{op[1]}', '9', 'retval'),
    Instruction('ret i32 %{op[0]}', '9')
]
cfg.compute_adjacency()
print(cfg)

Вывод этого питоновского скрипта совпадает с промежуточным кодом, сгенерированным клангом (ну, с точностью до служебных атрибутов, которые мне не нужны), и точно так же компилируется/отрисовывается при помощи llvm.

А вот теперь самое интересное

Мы отложили компилятор, и нарисовали 55+69 строчек питона, позволяющих манипулировать одним примером, вычисляющим числа Фибоначчи. В нём используются только локальные переменные, и нигде нет операции взятия адреса от них, поэтому мы можем смело выкинуть вообще все alloca.

А затем нам нужно убрать все store/ load, заменив их на $\phi$ -функции в нужных местах. А где они, нужные места?

Скрытый текст

В принципе, есть довольно дубовый метод: на входе в каждый блок (ну, за исключением стартового) вставить по $\phi$ -функции на каждую переменную. Да-да, включая блоки, у которых только один предок!

Это называется максимальной SSA-формой, и такой подход обычно критикуют за то, что он создаёт слишком много избыточных $\phi$ -функций, что замедляет компиляцию (мне всё равно) и затрудняет оптимизацию. С другой стороны, я сильно подозреваю, что обычный DCE (dead code elimination) проход уберёт мусор начисто. Когда я дойду до DCE в tinyoptimizer, надо будет попробовать посчитать максимальную SSA форму, не исключено, что это будет проще и понятнее.

Не забываем про «premature optimization is the root of all evil»...

Давайте начнём с довольно очевидного рассуждения (тут и далее я работаю с последним примером): вот у нас в блоке if.then есть инструкция store i32 0, ptr %retval, записывающая значение 0 в память по адресу %retval. А ещё в блоке return есть инструкция %9 = load i32, ptr %retval, которая очевидным образом читает память с того же адреса %retval в регистр %9.

Но ведь мы можем попасть в блок return, минуя блок if.then! Для каждой инструкции store нам понадобится вставить $\phi$ -функцию в следующем случае:

Если существует путь из точки входа в граф, проходящий через блок со store, и достигает блока с load,
А также если существует другой путь из точки входа в граф, достигающий блока с тем жеload, но при этом не проходящий через наш store.

Вполне очевидно, что где-то эти два пути сходятся вместе, прежде чем дойти до load , вот в этом месте слияния нам и понадобится $\phi$ -функция для нашей переменной %retval.

Тот базовый блок, в котором сходятся наши два пути через граф, является одним из элементов границы доминирования блока, содержащего инструкцию store. И вот тут нам понадобится немножко технического вокабуляра.

Границы доминирования

Определение первое: блок A доминирует над блоком B тогда и только тогда, когда любой путь из начала графа в блок B проходит через узел A.

Вполне очевидным образом каждый базовый блок доминирует сам над собой, и, например, есть только два блока, которые доминируют над блоком return, поскольку только entry и сам return фигурируют во всех возможных путях, они и являются доминаторами для return.

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

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

{
    entry:      {entry},
    if.then:    {if.then, entry},
    if.end:     {if.end, entry},
    while.cond: {if.end, while.cond, entry},
    while.body: {while.body, if.end, while.cond, entry},
    while.end:  {while.end, if.end, while.cond, entry},
    return:     {return, entry}
}

Определение второе: блок A строго доминирует над блоком B, если A доминирует над B, но при этом ему не равен. Соответственно, вот множества строгих доминаторов для каждого из базовых блоков нашей программы:

{
    entry:      {},
    if.then:    {entry},
    if.end:     {entry},
    while.cond: {if.end, entry},
    while.body: {if.end, while.cond, entry},
    while.end:  {if.end, while.cond, entry},
    return:     {entry}
}

Определение третье: A является непосредственным доминатором B, если A строго доминирует над B, но не доминирует строго любой другой строгий доминатор B.

Чё? :)

Давайте на пальцах. У каждого узла есть набор строго его доминирующих, да? Ну, за исключением стартового узла. Так вот, возьмём ближайший из них и назовём непосредственным доминатором. Вот список непосредственных доминаторов для каждого из наших базовых блоков:

{
    entry:      None,
    if.then:    entry,
    if.end:     entry,
    while.cond: if.end,
    while.end:  while.cond,
    return:     entry,
    while.body: while.cond
}

Эту же информацию можно отрисовать в форме дерева, которое называют деревом доминирования:

Обратите внимания, что рёбра дерева доминирования далеко не всегда являются рёбрами исходного графа.

Определение четвёртое: границей (фронтиром) доминирования узла A называется такое множество узлов B, что:

A доминирует хотя бы одного из их прешественников B.
A не доминирует сам узел B.

Другими словами, граница доминирования узла A состоит из узлов, где доминирование A перестаёт доминировать: в эти узлы можно попасть как через A, так и альтернативными путями, не проходя через A.

Если узел в вашем графе потока управления имеет менее двух предшественников, то он не будет находиться в границе доминирования какого-либо узла, так как он не может быть точкой слияния конкурирующих определений. Концепция «границы доминирования» заключается в том, что это именно те узлы, где доминирование некоторого узла заканчивается. Можно также сказать, что к фронтиру относятся все узлы, в которые идут рёбра из поддерева доминирования.

Граница доминирования показывает места в CFG, где разные потоки управления сходятся. Узлы фронтира — это места, где могут встретиться разные версии одной и той же переменной, и поэтому там нужно вставлять $\phi$ -функции.

Сверим часы, вот для каждого узла набор узлов его фронтира:

{
    entry:      {}
    if.then:    {return}
    return:     {}
    while.body: {while.cond}
    if.end:     {return}
    while.end:  {return}
    while.cond: {return, while.cond}
}

Вернёмся к нашей операции store внутри блока if.then. Блок return находится на границе доминирования if.then, так что в нём нужно вставить $\phi$ -функцию. Вроде всё сходится. Кстати, обратите внимания, что while.cond находится в собственном фронтире. Это совершенно нормальная ситуация, беспокоиться не нужно. Но нужно заметить, что у нас есть load переменной %i в блоке while.cond, и к ней мы можем прийти как из блока if.end, так и из блока while.body. Так что там тоже понадобятся $\phi$ -функции.

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

from ir import *

class Optimizer:
    def __init__(self, cfg):
        self.cfg, self.entry = cfg, cfg.blocks['entry']
        self.reachable, self.postorder = set(), []  # for efficiency reasons, we visit blocks in a way that ensures
        self.dfs(self.entry)                        # processing of each block after its predecessors have been processed

    def dfs(self, block):
        self.reachable.add(block)
        for b in block.successors:
            if b not in self.reachable: self.dfs(b)
        self.postorder.append(block)

    def dominators(self):                       # the iterative dominator algorithm [Cooper, Harvey and Kennedy, 2006]
        dom = { b : set(self.cfg.blocks.values()) for b in self.cfg.blocks.values() }
        dom[ self.entry ] = { self.entry }
        changed = True
        while changed:
            changed = False
            for b in self.postorder[-2::-1]:    # for all blocks (except the source) in reverse postorder
                new_dom = set.intersection(*[ dom[p] for p in b.predecessors ]) | { b }
                changed = dom[b] != new_dom
                dom[b]  = new_dom
        return dom                              # dom[b] contains every block that dominates b

    def immediate_dominators(self):
        dom = self.dominators()
        idom = { self.entry : None }                                # idom[b] contains exactly one block, the immediate dominator of b
        for b in self.postorder[-2::-1]:                            # reverse or not we do not care here, but do not visit the source block
            idom[b] = max(dom[b] - {b}, key=lambda x: len(dom[x]))  # immediate dominator is the one with the maximum number of dominators (except the block itself)
        return idom

    def dominance_frontiers(self):
        idom = self.immediate_dominators()
        df = { b : set() for b in self.reachable }
        for b in filter(lambda x: len(x.predecessors)>1, self.reachable): # iterate through junction points among reachable blocks
            for p in b.predecessors:
                runner = p
                while runner != idom[b]:
                    df[runner].add(b)
                    runner = idom[runner]
        return df

Я тут не изобретал ничего нового, взял самые дубовые (и заодно самые медленные), которые легко можно найти в той же самой википедии. Всё считается чётко по определениям, которые я привёл. Сначала считаются множества доминирования, потом из них выводится дерево доминирования, а из него выводятся фронтиры. Если кто может предложить код покороче, я с интересом слушаю, а то сорок строчек - это сразу одна десятая остававшегося мне бюджета на проект :)

Непосредственно mem2reg

Вооружившись терминологией, мы можем наконец чётко определить алгоритм размещения $\phi$ -функций в нашем графе потока управления. Трепещите же!

для каждой переменной v
  добавить в очередь все блоки, содержащего store v
  пока очередь непуста
    достать один блок b1
    для каждого блока b2 на фронтире доминирования b1
      если мы еще не вставляли фи-функцию для переменной v, то
        вставить её
        добавить блок b2 в очередь

Всё крайне просто, но есть один важный нюанс, про который нужно не забыть: вставка $\phi$ -функции в какой-то мере сродни добавлению операции store: добавив $\phi$ -функцию в блок b2, мы обязаны и его поставить в очередь на обработку.

Обратите внимание, что на этом этапе мы добавляем $\phi$ -функции, но не заполняем её аргументы. Этим мы займёмся при удалении load и store. Давайте посмотрим на полный код прохода mem2reg в моём исполнении:

class mem2reg(Optimizer):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.allocas = self.remove_promotable_allocas()
        self.place_phi()
        stack = [{}] # stack of maps (variable -> value); necessary for storing context while branching
        self.remove_store_load(self.entry, None, set(), stack)

    def remove_promotable_allocas(self):
        allocas = { i.op[0] : i.string.split()[-1] for i in self.entry.instructions if 'alloca' in i.string } # find all allocas in the entry block
        for v in allocas.copy().keys():             # for every variable
            for b in self.reachable:
                for i in b.instructions:            # for every instruction in the CFG
                    if f'* %{v}' in i.string:       # if we find a funcall with a reference to the variable
                        allocas.pop(v, None)        # then the variable is not promotable
        for i in self.entry.instructions[:]:
            if 'alloca' in i.string and i.op[0] in allocas:
                self.entry.instructions.remove(i)   # remove promotable allocas
        return allocas # variable name -> type dictionary

    def place_phi(self):                                        # Fig 9.9b in [Cooper and Torczon, Engineering a Compiler Broché]
        df = self.dominance_frontiers()
        phi_places = { v:set() for v in self.allocas.keys() }   # map (variable -> set of basic blocks)
        for v in self.allocas.keys():
            blocks_with_store = { b for b in self.reachable for i in b.instructions if 'store' in i.string and i.op[1]==v }
            blocks_to_consider = blocks_with_store.copy()
            while blocks_to_consider:
                block = blocks_to_consider.pop()
                for frontier in df[block]:
                    if frontier in phi_places[v]: continue
                    phi_places[v].add(frontier)
                    blocks_to_consider.add(frontier)
        for v, blocks in phi_places.items():                    # insert phi nodes (for the moment without choices)
            for b in blocks:
                b.phi_functions.append(Phi(v + '_' + b.label, self.allocas[v]))

    def remove_store_load(self, block, prev, visited, stack):
        def find_variable(v, stack):                # walk the stack back until
            for frame in reversed(stack):           # the current variable instance is found
                if v in frame: return frame[v]      # N.B. we suppose that the variables were initialized explicitly

        for phi in block.phi_functions:             # place phi node choice for the current path
            v = phi.reg[:-len(block.label)-1]       # phi saves the choice into a register named foo_bar, where foo is the name of the variable, and bar is the name of the basic block
            val = find_variable(v, stack)
            phi.choice[prev.label] = val
            stack[-1][v] = phi.reg

        if block in visited: return                 # we need to revisit blocks with phi functions as many times as we have incoming edges,
        visited.add(block)                          # therefore the visited check is made after the choice placement

        for i in block.instructions[:]:             # iterate through a copy since we modify the list
            if 'load' in i.string and i.op[1] in self.allocas:
                val = find_variable(i.op[1], stack)
                block.instructions.remove(i)
                block.find_and_replace(i.op[0], val)
            elif 'store' in i.string and i.op[1] in self.allocas:
                stack[-1][i.op[1]] = i.op[0]
                block.instructions.remove(i)
            elif 'br i1' in i.string:
                stack.append({})
                self.remove_store_load(self.cfg.blocks[i.op[1]], block, visited, stack)
                stack.pop()
                stack.append({})
                self.remove_store_load(self.cfg.blocks[i.op[2]], block, visited, stack)
                stack.pop()
            elif 'br label' in i.string:
                self.remove_store_load(self.cfg.blocks[i.op[0]],  block, visited, stack)

Вызывается он тривиально:

from optimizer import *
mem2reg(cfg)
print(cfg)

Класс mem2reg наследуется от класса Optimizer, так что у меня есть полный доступ к функциям доминирования. Всю работу делает конструктор. Для начала он находит, какие переменные в cfg поддаются этой оптимизации при помощи. remove_promotable_allocas()- удаляет alloca локальных переменных, от которых не берётся адрес, то есть, убирает со стека все локальные переменные, недоступные для изменения другими функциями. Затем place_phi() вставляет $\phi$ -функции по вышеописанному алгоритму с фронтирами. Остаётся заполнить их аргументы, что и делает remove_store_load().

Давайте разберёмся с тем, как она работает, всё на том же примере с вычислением чисел Фибоначчи. После удаления alloca (в данном случае вообще все удалены) и расстановки $\phi$ -функций, мы получаем следующий промежуточный код:

define i32 @fib(i32 %n) {
entry:
        store i32 %n, ptr %n.addr
        store i32 0, ptr %a
        store i32 1, ptr %b
        store i32 1, ptr %i
        store i32 0, ptr %c
        %0 = load i32, ptr %n.addr
        %cmp = icmp eq i32 %0, 0
        br i1 %cmp, label %if.then, label %if.end
if.then:
        store i32 0, ptr %retval
        br label %return
if.end:
        br label %while.cond
while.cond:
        %a_while.cond = phi i32 [?], [?]
        %b_while.cond = phi i32 [?], [?]
        %i_while.cond = phi i32 [?], [?]
        %c_while.cond = phi i32 [?], [?]
        %1 = load i32, ptr %i
        %2 = load i32, ptr %n.addr
        %cmp1 = icmp slt i32 %1, %2
        br i1 %cmp1, label %while.body, label %while.end
while.body:
        %3 = load i32, ptr %b
        store i32 %3, ptr %c
        %4 = load i32, ptr %a
        %5 = load i32, ptr %b
        %add = add i32 %4, %5
        store i32 %add, ptr %b
        %6 = load i32, ptr %c
        store i32 %6, ptr %a
        %7 = load i32, ptr %i
        %add2 = add i32 %7, 1
        store i32 %add2, ptr %i
        br label %while.cond
while.end:
        %8 = load i32, ptr %b
        store i32 %8, ptr %retval
        br label %return
return:
        %retval_return = phi i32 [?], [?]
        %a_return = phi i32 [?], [?]
        %b_return = phi i32 [?], [?]
        %i_return = phi i32 [?], [?]
        %c_return = phi i32 [?], [?]
        %9 = load i32, ptr %retval
        ret i32 %9
}

Я не привожу рисунка, поскольку llvm отказывается его рисовать :) Оно и понятно, код сломан: в первом же базовом блоке первый же регистр %a не определён, да и $\phi$ -функции поставлены хоть и верно, но без аргументов не являются валидной инструкцией.

Давайте пройдём по нашему графу в глубину, начнём, очевидно, со входной точки в программу. Итак, у нас есть блок entry:

entry:
        store i32 %n, ptr %n.addr
        store i32 0, ptr %a
        store i32 1, ptr %b
        store i32 1, ptr %i
        store i32 0, ptr %c
        %0 = load i32, ptr %n.addr
        %cmp = icmp eq i32 %0, 0
        br i1 %cmp, label %if.then, label %if.end

Пройдёмся поочерёдно по всем инструкциям, и удалим все store и load, но при этом посмотрим на их аргументы. Первая же инструкция сохраняет %n по адресу %n.addr, это означает, что любой load по этому адресу в этом блоке может быть заменён на %n, и в частности, %0 тоже заменяется на %n. Так же запоминаем значения для %a, %b, %i, %c, и убирем их store. После обработки entry принимает следующий вид:

entry:
        %cmp = icmp eq i32 %n, 0
        br i1 %cmp, label %if.then, label %if.end

У entry два потомка - if.then и if.end. Поскольку у нас обход графа в глубину, сначала идём в if.then. Он имеет следующий вид:

if.then:
        store i32 0, ptr %retval
        br label %return

Удаляем store, запоминаем, что %retval - это 0, идём в return. Он имеет следующий вид:

return:
        %retval_return = phi i32 [?], [?]
        %a_return = phi i32 [?], [?]
        %b_return = phi i32 [?], [?]
        %i_return = phi i32 [?], [?]
        %c_return = phi i32 [?], [?]
        %9 = load i32, ptr %retval
        ret i32 %9

Пришло время начать заполнять $\phi$ -функции. Мы пришли из ветки if.then, поэтому вспоминаем, что %retval в этой ветке это 0, %a это 0, %b это 1, %i это 1, %c это 0. Заполняем соответствующий аргумент $\phi$ -функций, удаляем load, заменяя %9 на %retval_return. После обработки блок принимает следующий вид:

return:
        %retval_return = phi i32 [0, %if.then], [?]
        %a_return = phi i32 [0, %if.then], [?]
        %b_return = phi i32 [1, %if.then], [?]
        %i_return = phi i32 [1, %if.then], [?]
        %c_return = phi i32 [0, %if.then], [?]
        ret i32 %retval_return

Это терминальный блок, так что обход в глубину возвращается к обработке if.end (там нечего делать), а затем while.cond и так далее.

В итоге мой код сгенерирует вот такой граф потока управления:

Давайте его сравним с тем, что сгенерировал llvm из кланговского вывода:

Как видно, llvm оказался чуточку умнее: он заметил, что a,b,i,c не используются в блоке return, и не стал их пробрасывать, а также понял, что c переписывается в while.body, поэтому не стал его пробрасывать через while.cond. В остальном же результат просто идентичен, так что я вполне доволен моим игрушечным компилятором.

Ну что, осталось склеить это дело с компилятором. Тесты проходят, ура!

ssloy@periwinkle:~/tinyoptimizer$ make test
Testing test-programs/nontrivial/mandelbrot.wend... ok
Testing test-programs/nontrivial/bitwise.wend... ok
Testing test-programs/nontrivial/trig-hp12c.wend... ok
Testing test-programs/nontrivial/sqrt.wend... ok
Testing test-programs/simple/fixed-point.wend... ok
Testing test-programs/simple/eight-queens.wend... ok
Testing test-programs/simple/mutual-recursion.wend... ok
Testing test-programs/simple/popcount.wend... ok
Testing test-programs/simple/collatz.wend... ok
Testing test-programs/simple/sopfr.wend... ok
Testing test-programs/elementary/helloworld.wend... ok
Testing test-programs/elementary/arithmetic.wend... ok
Testing test-programs/elementary/scope.wend... ok
Testing test-programs/elementary/overload.wend... ok
Testing test-programs/elementary/int-overflow.wend... ok

Код компилятора живёт тут, а поиграть с mem2reg вручную, без остального компилятора можно в ветке mem2reg. Критика, советы и предложения приветствуются!

Теги:

Хабы:

Если эта публикация вас вдохновила и вы хотите поддержать автора — не стесняйтесь нажать на кнопку