Узнав про 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.
