На Хабре есть две статьи, автор которых пишет виртуальную машину для исполнения простого байткода, а потом применяет различные оптимизации для ускорения этой виртуальной машины. Кроме того, есть и компилятор простого С-подобного языка в этот самый байткод. Ознакмившись со статьями и этим компилятором, я подумал, что будет интересно изучить, как написать виртуальную машину этого языка, которая сможет делать JIT-компиляцию байткода с помощью библиотеки libjit. Опыт этого я и описываю в настоящей статье. В интернете есть статьи, описывающие испльзование этой библитеки, но все, что я видел, описывают генерацию машинного кода с помощью libоit для конкретных программ, а не произвольного байткода: есть официальный tutorial, серия статей и ещё серия сравнений на Хабре.

Весь мой код приведён в моём репозитории.

Компиляция в байт-код

Как я уже упоминал, PigletC умеет компилировать простой С-подобный язык в байткод в текстовом представлении. Мне сперва понадобится программа, которая будет переводить байткод из текстового формата в бинарный.

Некоторые детали:

• Каждую инструкцию для простоты будем кодировать 4-мя байтами (тип int). Это не совсем оптимально, так как инструкций всего 25, но так будет проще, потому что я хочу, чтобы аргументы инструкций тоже были 4-ёх байтными.

• Таким образом, выходной бинарный файл представляет из себя последовательность int-ов, каждый из которых является либо кодом инструкции, либо аргументом инструкции, либо специальным значением (см дальше). Аргументы инструкций идут только после тех инструкций, которые имеют аргумент.

• После инструкций перехода (JUMP, JUMP_IF_TRUE, JUMP_IF_FALSE) ставится индекс инструкции, на которую переход идёт - то есть порядковый номер числа, кодирующего эту инструкцию.

• Для удобства, я решил что на месте каждой метки в бинарном файле будут записываться числа 0xcafe и 0xbabe. Таким образом, при анализе байткода будет сразу понятно, что сюда могут вести инструкции перехода. Вообще-то, это необязательно: можно было бы просто сохранить все аргументы инструкций перехода.

В остальном ничего особенно интересного в коде ассемблера нет. Приведён он здесь.

Интерпретация байт-кода

Дальше напишем реализацию интерпретатора байт-кода, чтобы было с чем сравнивать исполнение JIT-компилированного кода.

Состояние виртуальной машины определяется следующими элементами:

• Стек. На стек можно класть 4-ёх байтных знаковые числа

• Память - лента 4-ёх байтных знаковых числе, к которым можно обращаться по индексу.

• Номер следующей инструкции для выполнения

Изначально для реализации стека и памяти можно использовать STL контейнеры stack и vector, соответственно (можно и для обоих использовать вектор, а можно для обоих использвать дек) - тогда размеры памяти и стека не будут ограничены какими-либо константами. Однако, потом мне пришлось от этого отказаться, поскольку с JIT-компиляцией кусков кода это бы плохо сочеталось (вообще-то, в принципе в JIT компилированный код можно было бы запихнуть вызовы методов стека/вектора/дека, но я почти уверен, что заинлайнтить это было бы невозможно, а без инлайна это работало бы неэффективно).

В общем-то, реализация интерпретатора байт-кода довольно тривиальна: это огромный switch, в котором для каждой инструкции написано, что она делает.

void store_to_memory(int addr, int value) {
        memory[addr] = value;
    }

int stack_pop() {
    return stack[--stack_size];
}

void run() {
        size_t ip = 0;
        while (ip < instructions_number) {
            ip++;
            switch (instructions[ip - 1]) {
                case OP_JUMP: {
                    auto arg = instructions[ip];
                    ip++;
                    ip = arg;
                    break;
                }
                case OP_JUMP_IF_TRUE: {
                    auto arg = instructions[ip];
                    ip++;
                    if (stack_pop()) {
                        ip = arg;
                    }
                    break;
                }
                case OP_JUMP_IF_FALSE: {
                    auto arg = instructions[ip];
                    ip++;
                    if (!stack_pop()) {
                        ip = arg;
                    }
                    break;
                }
                case OP_LOADADDI: {
                    auto arg = instructions[ip];
                    ip++;
                    stack[stack_size - 1] += memory[arg];
                    break;
                }
                case OP_LOADI: {
                    auto arg = instructions[ip];
                    ip++;
                    stack[stack_size++] = memory[arg];
                    break;
                }
                case OP_PUSHI: {
                    auto arg = instructions[ip];
                    ip++;
                    stack[stack_size++] = arg;
                    break;
                }
                case OP_DISCARD: {
                    stack_size--;
                    break;
                }
                case OP_STOREI: {
                    auto addr = instructions[ip];
                    ip++;
                    store_to_memory(addr, stack_pop());
                    break;
                }
                case OP_LOAD: {
                    auto addr = stack_pop();
                    stack[stack_size++] = memory[addr];
                    break;
                }
                case OP_STORE: {
                    auto val = stack_pop();
                    auto addr = stack_pop();
                    store_to_memory(addr, val);
                    break;
                }
                case OP_ADDI: {
                    auto arg = instructions[ip];
                    ip++;
                    stack[stack_size - 1] += arg;
                    break;
                }
                case OP_DUP: {
                    stack[stack_size] = stack[stack_size - 1];
                    stack_size++;
                    break;
                }
                case OP_SUB: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] -= arg;
                    break;
                }
                case OP_ADD: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] += arg;
                    break;
                }
                case OP_DIV: {
                    auto arg = stack_pop();
                    if (arg == 0) {
                        cerr << "ZERO DIVISION\n";
                        return;
                    }
                    stack[stack_size - 1] /= arg;
                    break;
                }
                case OP_MUL: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] *= arg;
                    break;
                }
                case OP_ABORT: {
                    cerr << "OP_ABORT called\n";
                    return;
                }
                case OP_DONE: {
                    cout << "program DONE\n";
                    return;
                }
                case OP_PRINT: {
                    cout << stack_pop() << "\n";
                    break;
                }
                case OP_POP_RES: {
                    stack_pop();
                    break;
                }
                case OP_GREATER_OR_EQUALI: {
                    auto arg = instructions[ip];
                    ip++;
                    stack[stack_size - 1] = stack[stack_size - 1] >= arg;
                    break;
                }
                case OP_GREATER_OR_EQUAL: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] = stack[stack_size - 1] >= arg;
                    break;
                }
                case OP_GREATER: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] = stack[stack_size - 1] > arg;
                    break;
                }
                case OP_LESS: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] = stack[stack_size - 1] < arg;
                    break;
                }
                case OP_LESS_OR_EQUAL: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] = stack[stack_size - 1] <= arg;
                    break;
                }
                case OP_EQUAL: {
                    auto arg = stack_pop();
                    stack[stack_size - 1] = stack[stack_size - 1] == arg;
                    break;
                }
                case 0xcafe: {
                    // skip 0xbabe
                    ip++;
                }
            }
        }
  }

Переходим к JIT компиляции

Поскольку я пишу на С++, то буду использовать версию libjit для этого языка, с классами и наследованием.

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

Для того, чтобы создать компилируемую функцию, надо создать класс, наследуемый от jit_function, определить конструктор и переопределить методы build и create_signature. Вообще-то, в PigletC программа (по крайней мере пока) может иметь только одну функцию, а байткод не имеет инструкций для вызова функции и возврата из неё (и поскольку инструкции перехода могут переходить только по адресам, известным на момент компиляции, то без дополнительной поддержки на уровне байткода никак не обойтись), но тем не менее функцию create_signature можно написать универсальную:

jit_type_t create_signature() override {
        auto ret_type = (is_void ? jit_type_void : jit_type_int);
        jit_type_t params[num_args];
        for (int i = 0; i < num_args; i++) {
            params[i] = jit_type_int;
        }
        return jit_type_create_signature(jit_abi_cdecl, ret_type, params, num_args, 1);
}

Здесь считаем, что функция может применять только аргументы типа int, а возвращать либо int, либо void.

Теперь осталось всего ничего - собрать функцию на основе байткода. Тут у меня много времени ушло на размышления о том, что делать со стеком. При выполнении кода интерпретатором стек - это массив/вектор/стек, определённый в виртуальной машине. При JIT-компиляции же мне думалось, что будет эффективным чтобы как стек виртуальной машины использовался настоящий стек (то есть тот, на который под x86 указывает регистр %rsp). Однако, я не нашёл в libjit функций, которые позволили бы такое сделать: во-первых, видимо, библитека должна быть кроссплатформенной, а стек на разных архитектурах может быть разный, а во-вторых, libjit будет испльзовать стек как-то самостоятельно, и поэтому не может давать у нему прямой доступ.

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

Подготовка к компиляции функции будет выглядеть так:

auto memory = (jit_type_create_pointer(jit_type_int, 1));
memory = this->new_constant(memory_outer);
auto stack = new_value(jit_type_create_pointer(jit_type_int, 1));
stack = this->new_constant(stack_outer);
jit_value stack_size_ptr = new_value(jit_type_create_pointer(jit_type_int, 1));
stack_size_ptr = this->new_constant(stack_size_ptr_outer);
auto push = [this, &stack, &stack_size_ptr] (jit_value&& value) {
    push_on_stack(stack, stack_size_ptr, std::move(value));
};
auto pop = [this, &stack, &stack_size_ptr] () {
    return pop_from_stack(stack, stack_size_ptr);
};
int& ip = *ip_ptr;

И где-то после этого кода:

void push_on_stack(jit_value& stack, jit_value& stack_size_ptr, jit_value&& arg) {
    insn_store_elem(stack, insn_load_elem(stack_size_ptr, new_constant(0), jit_type_int), arg);
    insn_store_elem(stack_size_ptr, new_constant(0),
                        insn_load_elem(stack_size_ptr, new_constant(0), jit_type_int) + new_constant(1));
}

jit_value pop_from_stack(jit_value& stack, jit_value& stack_size_ptr) {
    insn_store_elem(stack_size_ptr, new_constant(0),
                    insn_load_elem(stack_size_ptr, new_constant(0), jit_type_int) - new_constant(1));
    return insn_load_elem(stack, insn_load_elem(stack_size_ptr, new_constant(0), jit_type_int), jit_type_int);
}

jit_value peek_stack(jit_value& stack, jit_value& stack_size_ptr) {
    return insn_load_elem(stack, insn_load_elem(stack_size_ptr, new_constant(0), jit_type_int) - new_constant(1),
                          jit_type_int);
}

Все значения (константы) и переменные, которые JIT-компилированная функция может использовать, должны быть объявлены как переменные типа jit_value. Строчка за строчкой, мы делаем следующее:

1. С помощью jit_type_create_pointer создаём тип "указатель на число". Первый аргумент jit_type_int - тип, на который указываем, число "1" - на сколько увеличивается счётчик количества использований типа. Наконец, тип передаём в new_value - и создаём в JIT - функции переменную типа "указатель на число" для хранения массива с памятью виртуальной машины

2. Записываем в переменную для адреса памяти собственно адрес памяти

3. проделываем пункт 1. для стека

4. проделываем пункт 2. для стека

5. проделываем пункт 1. для указателя на размер стека

6. проделываем пункт 2. для указателя на размер стека

Далее для удобства надо написать функции, которые будут "класть значения на стек" и "брать значения со стека". В кавычках потому, что в реальности они будут добавлять в промежуточное представление нашей функции инструкции, которые будут делать соответствующие операции.

Общий смысл такой: мы читаем значение по указателю на размер стека, записываем по этому указателю значение на 1 больше или меньше. Далее читаем (или пишем) по адресу на нужную ячейку стека значение.

Наконец, внутри функции сборки функции объвляем две вспомогательные лямбды, чтобы операции со стеком были короче.

Теперь обрабтка инструкции байткода пишется довольно просто. Вот несколько примеров:

case OP_STORE: {
    auto val = pop();
    auto index = pop();
    insn_store_elem(memory, index, val);
    break;
}
case OP_LOAD: {
    auto index = pop();
    push(insn_load_elem(memory, index, jit_type_int));
    break;
}
case OP_PRINT: {
    auto tmp = pop().raw();
    this->insn_call_native("print_int", (void*)(&print_int), signature_helper(jit_type_void, jit_type_int, end_params),
        &tmp, 1, 0);
    break;
}

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

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

Виртуальная машина (интерпретатор) теперь будет проверять, является ли следующая инструкция переходом. Если да, выполняет. Иначе вызывает JIT-компиляцию, которая компилирует функцию пока не встретит инструкцию, которую не поддерживает. Далее скомпилированную функцию будем сохранять в кеш и вызывать.

Естественно, пока это всё работает очень медленно - куски, которые мы компилируем, довольно небольшие, и интепретатору приходиться работать очень часто. Пока что это только демонстрация возможностей libjit.

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

Теперь при JIT-компиляции мы заведём словарь с отображением адресов инструкций на метки. Когда встречаем 0xcafe - ставим метку соответствующую на этом месте. Когда встречаем инструкцию перехода - указываем, что надо выполнить переход по метке, которую достаём из словаря по адресу перехода. Код:

case 0xcafe: {
    // skipping 0xbabe
    ip++;
    // ip is now pointing to the instruction after the label
    if (labels.count(ip) == 0) {
        labels[ip] = jit_label_undefined;
    }
    insn_label(labels[ip]);
    break;
}
case OP_JUMP: {
    auto arg = instructions[ip];
    ip++;
    if (labels.count(arg) == 0) {
        labels[arg] = jit_label_undefined;
    }
    insn_branch_if(new_constant(true), labels[arg]);
    break;
}
case OP_JUMP_IF_TRUE: {
    auto arg = instructions[ip];
    ip++;
    if (labels.count(arg) == 0) {
        labels[arg] = jit_label_undefined;
    }
    insn_branch_if(pop(), labels[arg]);
    break;
}
case OP_JUMP_IF_FALSE: {
    auto arg = instructions[ip];
    ip++;
    if (labels.count(arg) == 0) {
        labels[arg] = jit_label_undefined;
    }
    insn_branch_if_not(pop(), labels[arg]);
    break;
}

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

В чём же дело? Проблема в том, что мы генерируем довольно сложный код, а оптимизатор libjit очень слаб - поэтому он особенно ничего с этим сделать не может.

Попробуем немного оптимизировать код, который мы генерируем, чтобы увидеть ускорение от JIT-компиляции.

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

jit_value stack_size = new_value(jit_type_int);
stack_size = new_constant(0);
auto push = [this, &stack, &stack_size](jit_value&& value) {
    insn_store_elem(stack, stack_size, value);
    stack_size = stack_size + new_constant(1);
};
auto pop = [this, &stack, &stack_size]() {
    stack_size = stack_size - new_constant(1);
    return insn_load_elem(stack, stack_size, jit_type_int);
};
auto peek = [this, &stack, &stack_size]() {
    return insn_load_elem(stack, stack_size - new_constant(1), jit_type_int);
};

Вот на этот раз ускорение значительное: интерпретация байткода работает за 1.2 секунды, а JIT-компилированный код работает за 0.4 секунды.

Хотя, с другой стороны, можно было бы ожидать и более существенного ускорения. Я предполагаю, что слабые оптимизации libjit в сочетании с тем, что байткод у нас от стековой, а не регистровой виртуальной машины, не позволяют коду действительно быть быстрым — ведь инструкции байткода постоянно используют стек, что сильно медленнее использования регистров. Отчасти эту проблему мы можем решить заменой некоторых последовательностей инструкций на более сложные инструкции, но в меньшем количестве. Эта оптимизация описывалась в статье про оптимизацию интерпретаторов байткода, на которую я ссылался в начале статьи.