Байт-код

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

Узнав про KMath, я решил обобщить эту оптимизацию для множества математических структур и операторов, определенных на них.

KMath это библиотека для математики и компьютерной алгебры, активно использующая контекстно-ориентированное программирование в Kotlin. В KMath разделены математические сущности (числа, векторы, матрицы) и операции над ними они поставляются отдельным объектом, алгеброй, соответствующей типу операндов, Algebra<T>.

import scientifik.kmath.operations.*ComplexField {   (i pow 2) + Complex(1.0, 3.0)}

Таким образом, после написания генератора байт-кода, с учетом оптимизаций JVM можно получить быстрые расчеты для любого математического объекта достаточно определить операции над ними на Kotlin.

API

Для начала нужно было разработать API выражений и только после этого приступать к грамматике языка и синтаксическому дереву. Также появилась удачная идея определить алгебру над самими выражениями, чтобы представить более наглядный интерфейс.

База всего API выражений интерфейс Expression<T>, а простейший способ его реализовать прямо определить метод invoke от данных параметров или, например, вложенных выражений. Подобная реализация была интегрирована в корневой модуль как справочная, хоть и самая медленная.

interface Expression<T> {   operator fun invoke(arguments: Map<String, T>): T}

Более продвинутые реализации уже основаны именно на MST. К ним относятся:

• интерпретатор MST,
• генератор классов по MST.

API для парсинга выражений из строки в MST уже доступно в dev-ветке KMath как и более или менее окончательный генератор кода JVM.

Перейдем к MST. Сейчас в MST представлены четыре вида узлов:

• терминальные:
  
  • символы (то есть переменные)
  • числа;
• унарные операции;
• бинарные операции.

Первое, что с ним можно сделать, это обойти и посчитать результат по имеющимся данным. Передав в целевую алгебру ID операции, например +, и аргументы, например 1.0 и 1.0, мы можем надеяться на результат, если эта операция определена. В противном случае при вычислении выражение упадет с исключением.



Для работы с MST, помимо языка выражения, еще есть алгебра например MstField:

import scientifik.kmath.ast.*import scientifik.kmath.operations.*import scientifik.kmath.expressions.*RealField.mstInField { number(1.0) + number(1.0) }() // 2.0

Результатом метода выше является реализация Expression<T>, при вызове вызывающая обход дерева, полученного при вычислении функции, переданной в mstInField.

Генерация кода

Но это не все: при обходе мы можем как угодно применять параметры дерева и не волноваться о порядке действий и арности операций. Именно это и используется для генерации байт-кода.

Генерация кода в kmath-ast это параметризованная сборка класса JVM. Входные данные MST и целевая алгебра, на выходе инстанс Expression<T>.

Соответствующий класс, AsmBuilder, и еще несколько extension-функций предоставляют методы для императивной сборки байт-кода поверх ObjectWeb ASM. С их помощью обход MST и сборка класса выглядят чисто и занимают менее 40 строк кода.

Рассмотрим сгенерированный класс для выражения 2*x, приведен декомпилированный из байт-кода исходник на Java:

package scientifik.kmath.asm.generated;import java.util.Map;import scientifik.kmath.asm.internal.MapIntrinsics;import scientifik.kmath.expressions.Expression;import scientifik.kmath.operations.RealField;public final class AsmCompiledExpression_1073786867_0 implements Expression<Double> {   private final RealField algebra;   public final Double invoke(Map<String, ? extends Double> arguments) {       return (Double)this.algebra.add(((Double)MapIntrinsics.getOrFail(arguments, "x")).doubleValue(), 2.0D);   }   public AsmCompiledExpression_1073786867_0(RealField algebra) {       this.algebra = algebra;   }}

Сначала здесь был сгенерирован метод invoke, в котором были последовательно расставлены операнды (т.к. они находятся глубже в дереве), потом вызов add. После invoke был записан соответствующий бридж-метод. Далее было записано поле algebra и конструктор. В некоторых случаях, когда константы нельзя просто положить в пул констант класса, записывается еще поле constants, массив java.lang.Object.

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

Оптимизация вызовов к Algebra

Чтобы вызвать операцию от алгебры, нужно передать ее ID и аргументы:

RealField { binaryOperation("+", 1.0, 1.0) } // 2.0

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

RealField { add(1.0, 1.0) } // 2.0

Никакой особенной конвенции о маппинге ID операций к методам в реализациях Algebra<T> нет, поэтому пришлось вставить костыль, в котором вручную прописано, что +, например, соответствует методу add. Также есть поддержка для благоприятных ситуаций, когда для ID операции можно найти метод с таким же именем, нужным количеством аргументов и их типами.

private val methodNameAdapters: Map<Pair<String, Int>, String> by lazy {    hashMapOf(        "+" to 2 to "add",        "*" to 2 to "multiply",        "/" to 2 to "divide",...

private fun <T> AsmBuilder<T>.findSpecific(context: Algebra<T>, name: String, parameterTypes: Array<MstType>): Method? =    context.javaClass.methods.find { method ->...        nameValid && arityValid && notBridgeInPrimitive && paramsValid    }

Другая серьезная проблема это боксинг. Если мы посмотрим на Java-сигнатуры методов, которые получаются после компиляции того же RealField, увидим два метода:

public Double add(double a, double b)// $FF: synthetic method// $FF: bridge methodpublic Object add(Object var1, Object var2)

Конечно, легче не мучаться с боксингом и анбоксингом и вызвать бридж-метод: он тут появился из-за type erasure, чтобы правильно реализовывать метод add(T, T): T, тип T в дескрипторе которого на самом деле стерся до java.lang.Object.

Прямой вызов add от двух double тоже не идеален, потому что боксит возвращаемое значение (по этому поводу есть обсуждение в YouTrack Kotlin (KT-29460), но лучше вызывать именно его, чтобы в лучшем случае сэкономить два приведения типа входных объектов к java.lang.Number и их анбоксинга в double.

На решение этой проблемы ушло больше всего времени. Сложность здесь заключается не в создании вызовов к примитивному методу, а в том, что нужно сочетать на стеке и примитивные типы (как double), и их обертки (java.lang.Double, например), а в нужных местах вставлять боксинг (например, java.lang.Double.valueOf) и анбоксинг (doubleValue) здесь категорически не хватало работы с типами инструкций в байт-коде.

У меня были идеи навесить свою типизированную абстракцию поверх байт-кода. Для этого мне пришлось глубже разобраться в API ObjectWeb ASM. В итоге я обратился к бэкенду Kotlin/JVM, подробно изучил класс StackValue (типизированный фрагмент байт-кода, который в итоге приводит к получению какого-то значения на стеке операндов), разобрался с утилитой Type, которая позволяет удобно оперировать с типами, доступными в байт-коде (примитивы, объекты, массивы), и переписал генератор с ее использованием. Проблема, нужно ли боксить или анбоксить значение на стеке, решилась сама собой путем добавления ArrayDeque, хранящего типы, которые ожидаются следующим вызовом.

  internal fun loadNumeric(value: Number) {        if (expectationStack.peek() == NUMBER_TYPE) {            loadNumberConstant(value, true)            expectationStack.pop()            typeStack.push(NUMBER_TYPE)        } else ...?.number(value)?.let { loadTConstant(it) }            ?: error(...)    }

Выводы

В итоге мне удалось сделать генератор кода с помощью ObjectWeb ASM для вычисления выражений MST в KMath. Прирост производительности по сравнению с простым обходом MST зависит от количества узлов, так как байт-код линеен и в итоге не тратит время на выбор узла и рекурсию. Например, для выражения с 10 узлами разница во времени между вычислением с помощью сгенерированного класса и интерпретатора составляет от 19 до 30%.

Изучив проблемы, с которыми я столкнулся, я сделал следующие выводы:

• нужно сразу изучить возможности и утилиты ASM они сильно упрощают разработку и делают код читаемым (Type, InstructionAdapter, GeneratorAdapter);
• нет смысла тратить время на подсчет MaxStack самостоятельно, если это не критично для производительности, есть ClassWriter COMPUTE_MAXS и COMPUTE_FRAMES;
• очень полезно изучать бэкенды компиляторов реальных языков;
• следует разбираться в синтаксисе дескрипторов и, в частности, сигнатур, чтобы не ошибаться при использовании дженериков;
• и наконец, далеко не во всех случаях нужно лезть настолько глубоко есть более удобные способы работать с классами в рантайме, например, ByteBuddy и cglib.

Cпасибо, что прочитали.

Авторы статьи:

Ярослав Сергеевич Постовалов, МБОУ Лицей 130 им. академика Лаврентьева, участник лаборатории математического моделирования под руководством Войтишека Антона Вацлавовича

Татьяна Абрамова, исследователь лаборатории методов ядерно-физических экспериментов в JetBrains Research.