Java і Data Science. Розглядаємо підходи до обчислень SIMD в Java

Мене звати Юрій Зайчик, я Senior Java Developer у Luxoft. Працюю у сфері автомобільної навігації, де ми маємо справу з великими обсягами геопросторових даних.

З часом Java довела, що є правильним вибором для корпоративних рішень, мобільних платформ і навіть смарт-карт :). Але чи може Java бути корисним інструментом для такої важливої ​​сфери, як Data Science?

У цій статті я хотів би поділитися з вами своїми думками щодо проектів, які можуть зробити Java ближчою до науки про дані. Тут я розгляну підходи до обчислень SIMD в Java.

Векторна математика та SIMD

Single Instruction Multiple data або SIMD — це обчислювальний підхід, який дозволяє виконувати скалярні операції над декількома елементами даних одночасно. Це означає, що елементи даних будуть оброблятися паралельно. Сучасні архітектури процесорів підтримують цей підхід і мають спеціальний набір векторних інструкцій.

У процесорах Intel він реалізований за допомогою технологій SSE (Intel Streaming SIMD Extensions) або AVX (Advanced Vector Extensions). Компанія AMD також підтримує інструкції SSE.

Зараз JVM не підтримує прямий доступ до векторних інструкцій. Але, на щастя, є така класна штука, як процес Java Enhancement Proposal. І в рамках цього процесу, як частина проєкту Project Panama, з’явився векторний модуль java.incubator.vector. Цей модуль має API, що передбачає використання векторних інструкцій, коли JVM переводить байт-код в команди процесора.

У таблиці нижче ви можете побачити еволюцію Java Vector API.

Отже, для яких завдань можна використовувати Java vector API? Автори Java vector API модуля визначають використання векторного API як інструменту для розробників для створення високопродуктивного коду з передбачуваною і надійною векторизацією в порівнянні з автовекторизатором HotSpot. Серед областей, де може бути корисне використання векторного API, автори згадують машинне навчання, криптографію, фінтех та інші.

Як я вже згадував, я займаюся автомобільною навігацією, тому розглянемо приклад використання векторного API для розрахунків просторової геометрії.

JTS (Java Topology Suite) — це стандартний інструмент для роботи з векторною геометрією на мові Java. Він надає інтерфейси для типових геометричних фігур, таких як POINT, LINESTRING, POLYGON тощо. Ці інтерфейси забезпечують типові операції, які можуть бути виконані над геометричними об’єктами, наприклад, обчислення довжини геометричного об’єкта. Для LINESTRING алгоритм обчислення довжини базується на теоремі Піфагора.

У JTS LineString визначається набором точок у декартовій системі. Кожна точка містить значення широти (X) і довготи (Y). Щоб обчислити довжину такої ламаної, потрібно обчислити довжину кожного відрізка, а потім підсумувати їх. Довжину кожного відрізка можна обчислити за теоремою Піфагора, вважаючи гіпотенузою лінію, що з’єднує сусідні точки. Цей підхід можна визначити наступною формулою:

Погляньмо на Java-код JTS для наведеного вище алгоритму:

public static double ofLine(CoordinateSequence pts) {
    int n = pts.size();
    if (n <= 1) {
        return 0.0;
    } else {
        double len = 0.0;
        Coordinate p = new Coordinate();
        pts.getCoordinate(0, p);
        double x0 = p.x;
        double y0 = p.y;

        for(int i = 1; i < n; ++i) {
            pts.getCoordinate(i, p);
            double x1 = p.x;
            double y1 = p.y;
            double dx = x1 - x0;
            double dy = y1 - y0;
            len += Math.sqrt(dx * dx + dy * dy);
            x0 = x1;
            y0 = y1;
        }

        return len;
    }
}

Як бачимо, це скалярний алгоритм, де кожна пара точок обробляється ітераційно в циклі. Цей підхід може бути покращений, якщо ми застосуємо векторизацію.

Спочатку нам потрібно буде визначити тип даних для роботи з векторними геометріями. Звичайно, ми можемо розширити існуючий клас LineString з JTS, але для чистоти експерименту краще визначити окремий тип даних. Наш тип даних буде просто зберігати два масиви: один для координат широти, а інший для довготи.

public record VectorLine(double[] latitudes, double[] longitudes) {

    public static VectorLine of(LineString geom)
    {
        int dim = geom.getNumPoints();
        double[] xCoords = new double[dim];
        double[] yCoords = new double[dim];

        var coords = geom.getCoordinateSequence();
        for (int i = 0; i < dim; i++)
        {
            var coordinate = coords.getCoordinate(i);
            xCoords[i] = coordinate.x;
            yCoords[i] = coordinate.y;
        }
        return new VectorLine(xCoords, yCoords);
    }
}

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

Вхідна ламана розглядається як вектор точок:

Далі розділимо вхідні дані на 4 вектори: 2 для широти і 2 для довготи:

Розбиття вектора на осі необхідне для виконання векторного віднімання для знаходження dx і dy. Довжина у визначенні вектора буде виглядати наступним чином:

Тепер подивімося на частину коду, який реалізує алгоритм, наведений вище.

public class VectorLengthCalc
{
    public static final VectorSpecies<Double> SPECIES = DoubleVector.SPECIES_256;

    public static double calculate(VectorLine vl) {
        double length = 0D;

        int upperBound = SPECIES.loopBound(vl.latitudes().length);
        int i = 0;
        for (; i < upperBound; i += SPECIES.length()) {
            var vX1 = DoubleVector.fromArray(SPECIES, vl.latitudes(), i);
            var vX2 = DoubleVector.fromArray(SPECIES, vl.latitudes(), i + 1);
            var vY1 = DoubleVector.fromArray(SPECIES, vl.longitudes(), i);
            var vY2 = DoubleVector.fromArray(SPECIES, vl.longitudes(), i + 1);

            var vdx = vX2.sub(vX1);
            var vdy = vY2.sub(vY1);
            var vHypotenuses = vdx.mul(vdx).add(vdy.mul(vdy)).lanewise(VectorOperators.SQRT);

            length += vHypotenuses.reduceLanes(VectorOperators.ADD);
        }

        for (; i < vl.latitudes().length - 1; i++)
        {
            var dx = vl.latitudes()[i + 1] - vl.latitudes()[i];
            var dy = vl.longitudes()[i + 1] - vl.longitudes()[i];
            length += Math.sqrt(dx * dx + dy * dy);
        }

        return length;
    }
}

Статична константа SPECIES відображає тип векторної архітектури. Вона містить інформацію про тип векторного примітиву (ціле значення або значення з плаваючою комою), кількість векторних рядків та форму вектора (довжина в бітах кожної векторної інструкції). На моїй машині стоїть процесор Intel i-7 — 9750H, який підтримує AVX2 і SSE4, тому я вибрав 256-бітний векторний вид для чисел з плаваючою комою подвійної точності. Але розробники векторного API радять використовувати методи VectorShape.preferredShape() та VectorSpecies.ofPreferred() для вибору форми з урахуванням платформи.

Після того, як ми визначилися з видами векторів та отримали вхідні дані, нам необхідно задати границі для нашого векторного циклу. Векторний підхід за визначенням повинен бути безцикловим, але векторна інструкція обмежена по рядках, тобто вона може оперувати векторами з певною кількістю елементів (якщо нам потрібно відняти 2 вектори по 16 елементів, а векторна інструкція оперує тільки з 8-елементними векторами, то ми зробимо це віднімання за 2 ітерації), тому цикл тут все ж таки потрібен.

Для ефективного використання векторних інструкцій, в залежності від форми вектора та довжини вхідних даних, необхідно встановити межі циклу. Для цього використовується метод SPECIES.loopBound(). З опису методу: «Функція управління циклом, яка повертає найбільше кратне значення VLENGTH, яке менше або дорівнює заданому значенню довжини. Тут VLENGTH є результатом функції this.length(), а довжина інтерпретується як кількість рядків. Отримане значення R задовольняє цю нерівність: R <= length < R+VLENGTH. Зокрема, цей метод обчислює довжину — floorMod(length, VLENGTH), де floorMod обчислює значення залишку шляхом округлення його частки в бік від’ємної нескінченності. Оскільки VLENGTH є ступенем двійки, то результат також дорівнює length & ~(VLENGTH — 1).» можна зробити висновок, що найкращою верхньою межею буде найближче до вхідних даних число довжини, яке є кратним до числа смуг вектора.

Проілюструємо це на прикладі. Маємо ламану з 17 точками:

Згідно з нашим алгоритмом векторизації, з цих вхідних даних ми отримаємо 4 вектори, що складаються з 16 елементів кожен. Якщо наша векторна інструкція може обробляти 4-рядкові вектори, то верхня межа буде дорівнювати 16. Якщо у нас ламана з 19 точками, то ми отримаємо 4 вектори, що складаються з 18 елементів і при тій же векторній інструкції верхня межа все одно буде 16. Отже, у нас залишається ще 2 точки для обробки. Саме тому у нашому прикладі коду є другий цикл, який займається такими «залишками», що не вкладаються у векторний цикл. Другий цикл використовує стандартний підхід скалярних обчислень.

Тепер подивимось, що відбувається всередині векторного циклу. Спочатку ми створюємо «вектори віднімання» з вхідних даних.

var vX1 = DoubleVector.fromArray(SPECIES, vl.latitudes(), i);

Цей рядок є доволі самоописовим. Він створює DoubleVector X1 і поміщає туди значення i із вхідних даних. Після ініціалізації «векторів віднімання» виконуємо обчислення вектора згідно з алгоритмом.

var vdx = vX2.sub(vX1);

У цьому рядку виконується віднімання вектора. Оскільки наша векторна інструкція оперує з 4-смуговими подвійними векторами, то це означає, що за один такт процесора від 4 значень вектора Х1 буде віднято 4 значення вектора Х2. Те ж саме робиться і для довготи (Y1 та Y2).

var vHypotenuses = 
    vdx.mul(vdx).add(vdy.mul(vdy)).lanewise(VectorOperators.SQRT);

У цьому рядку ми виконуємо відразу 4 векторні операції. Йдемо зліва направо. Спочатку приводимо значення вектора, який містить результати віднімання Х, до степеня 2. Це робиться простим множенням вектора на самого себе. Потім до отриманого вектора додаємо аналогічне векторне множення значень Y. І, нарешті, до отриманого вектора, який містить суму 2-х множень, застосовуємо операцію lanewise (операції, які будуть виконуватися над кожним елементом вектора), в нашому випадку це операція добування квадратного кореня. Таким чином, ми отримали вектор, який містить довжини 4-х відрізків (нагадаємо, що ми оперуємо з 4-смуговими векторами).

length += vHypotenuses.reduceLanes(VectorOperators.ADD);

Тут ми виконуємо операцію векторної редукції (операція над елементами вектора, яка дає в результаті скалярний результат). У нашому випадку ми підсумовуємо довжини 4-х відрізків і додаємо результат до змінної-акумулятора повної довжини.

У порівнянні з простим скалярним алгоритмом, цей векторний підхід здається складнішим і все ще має допоміжний скалярний цикл для «залишкових» значень. Тож чи дійсно ми виграємо в продуктивності від векторизації тут? Ми можемо відповісти на це питання, вимірявши ефективність між скалярним і векторним алгоритмами.

Ми будемо використовувати JMH (Java Measurement Harness) для порівняльного аналізу наших обчислень. JMH можна легко додати як залежність Maven або Gradle. Вхідними даними використаємо дані деяких геометрій берегової лінії Землі (взяті з відкритих джерел), визначену у форматі WKT (Well Known Text) в системі координат WGS 84.

Буде використано 3 приклади: геометрія з 17 точками, геометрія з 19 точками та геометрія з 250 точками. Як я вже згадував вище, для геометрії з 17 точками буде використовуватися тільки векторний цикл, для геометрії з 19 точками будуть задіяні як векторний, так і скалярний цикли. Вхідні дані будуть мати вигляд типу:

public static final String WKT_LS_19_POINTS = "LINESTRING (57.73720703125002 -20.098437500000003, 57.65654296875002 -19.98994140625001, " +
        "57.575781250000006 -19.997167968750006, 57.51503906250002 -20.055957031250003, 57.486425781250006 -20.14394531250001, " +
        "57.416015625 -20.18378906250001, 57.3857421875 -20.228613281250006, 57.36210937500002 -20.33759765625001, " +
        "57.36513671875002 -20.40644531250001, 57.31767578125002 -20.42763671875001, 57.32832031250001 -20.450000000000003, " +
        "57.38330078125 -20.503710937500003, 57.52480468750002 -20.51318359375, 57.65126953125002 -20.48486328125, " +
        "57.706640625000006 -20.434863281250003, 57.72500000000002 -20.36884765625001, 57.78066406250002 -20.326953125000003, " +
        "57.7919921875 -20.21259765625001, 57.73720703125002 -20.098437500000003) ";

Для використання JMH необхідно визначити клас бенчмарку, який буде містити налаштування вимірювання (кількість ітерацій, одиниця виміру тощо) та методи бенчмарку.

@OutputTimeUnit(TimeUnit.MILLISECONDS)
@State(Scope.Benchmark)
@BenchmarkMode(Mode.Throughput)
@Warmup(iterations = 3, time = 1)
@Measurement(iterations = 5, time = 1)
@Fork(value = 1, jvmArgsPrepend = "--add-modules=jdk.incubator.vector")
public class VectorBenchMark

Тут ми також надаємо аргумент JVM, який вказує віртуальній машині включити модулі інкубатора, інакше JVM зазнає невдачі з виключенням невідомого модуля. Далі нам потрібно буде підготувати вхідні дані: розібрати WKT та створити об’єкти VectorLine.

@Setup
public void setup() {
    geom17Points = LengthExperiment.getGeom(WktGeometries.WKT_LS_17_POINTS);
    geom19Points = LengthExperiment.getGeom(WktGeometries.WKT_LS_19_POINTS);
    geom250Points = LengthExperiment.getGeom(WktGeometries.WKT_LS_250_POINTS);

    vectorLine17Points = VectorLine.of((LineString) geom17Points);
    vectorLine19Points = VectorLine.of((LineString) geom19Points);
    vectorLine250Points = VectorLine.of((LineString) geom250Points);
}

А потім для кожної геометрії ми додамо еталонні методи для скалярного та векторного підходів.

@Benchmark
public double lengthVector17Ponts() {
    return VectorLengthCalc.calculate(vectorLine17Points);
}

@Benchmark
public double lengthScalar17Points() {
    return LengthExperiment.calcLengthScalar((LineString)geom17Points);
}

Для цього експерименту я використовував maven. Після запуску «mvn clean install» для цього проeкту, JMH створить uber-jar під назвою benhcmarks.jar. Для проведення експерименту нам потрібно запустити цей jar, вказавши параметром наш тестовий клас:

java -jar benchmarks.jar VectorBenchMark

Потім JMH запустить ітерації прогріву та вимірювання для кожного бенчмарк-методу, результат буде виглядати наступним чином:

# Run progress: 0.00% complete, ETA 00:00:48
# Fork: 1 of 1
WARNING: Using incubator modules: jdk.incubator.vector
# Warmup Iteration   1: 42437.357 ops/ms
# Warmup Iteration   2: 42365.393 ops/ms
# Warmup Iteration   3: 43891.023 ops/ms
Iteration   1: 44255.728 ops/ms
Iteration   2: 44147.808 ops/ms
Iteration   3: 44355.650 ops/ms
Iteration   4: 43774.046 ops/ms
Iteration   5: 44365.232 ops/ms

Result "com.luxoft.yz.VectorBenchMark.lengthScalar17Points":
  44179.693 ±(99.9%) 936.769 ops/ms [Average]
  (min, avg, max) = (43774.046, 44179.693, 44365.232), stdev = 243.276
  CI (99.9%): [43242.924, 45116.462] (assumes normal distribution)

І, нарешті, JMH опублікує порівняльну таблицю.

Benchmark                               Mode  Cnt      Score      Error   Units

VectorBenchMark.lengthScalar17Points   thrpt    5  44179.693 ±  936.769  ops/ms

VectorBenchMark.lengthScalar19Points   thrpt    5  39127.693 ± 1101.871  ops/ms

VectorBenchMark.lengthScalar250Points  thrpt    5   2845.169 ±   73.238  ops/ms

VectorBenchMark.lengthVector17Ponts    thrpt    5  71690.001 ± 4155.375  ops/ms

VectorBenchMark.lengthVector19Ponts    thrpt    5  53060.632 ±  969.685  ops/ms

VectorBenchMark.lengthVector250Ponts   thrpt    5   5336.413 ±  313.708  ops/ms

Ми використовували режим пропускної здатності, який вимірює ефективність у середній кількості операцій за одну мілісекунду (чим більше число — тим краща продуктивність). Адаптуємо цей показник до більш інформативного вигляду.

Отже, ми бачимо, що для кожного випадку векторизація принесла нам значний приріст продуктивності. Розглядаючи випадки з відносно невеликою кількістю точок, можна зробити висновок, що коли алгоритм використовує тільки векторний цикл (випадок 17 точок), то продуктивність зростає вдвічі, коли алгоритм використовує як скалярний, так і векторний цикли (19 точок), то приріст продуктивності менший — близько третини. Але якщо розглянути випадок з відносно великою кількістю точок, то повернення до скалярного циклу є несуттєвим для продуктивності (для цього випадку вона зросла майже вдвічі).

Отже, коли цей інкубатор вийде у світ, це може бути гарний час для того, щоб вдихнути свіже дихання в існуючі бібліотеки Java або створити нові :)

Якщо ви хочете погратися з Java Vector API або з описаним експериментом, повний код є в репозиторії GitHub. Також майте на увазі, що деякі системи вимагають явного включення інструкцій AVX або SSE через BIOS або системні прапори.

Підписуйтеся на Telegram-канал «DOU #tech», щоб не пропустити нові технічні статті