現在従事しているプロジェクトのプロダクトはJava製ですが、一部の実装にnative (C++)実装を含んでおり、JNIでJava-nativeコミュニケーションを実装しています。
今回は少し前にnativeに計算させるデータの渡し方を修正したので、改めてその実装のパフォーマンスを計ってみようという記事です。

array実装版

とりあえず修正前の実装と同様のサンプルコードです。
(計測に使った最終的なコードはこちら。)

nativeに渡すデータは、複数のlongのセットのコレクションです。array版ではすべてをいったんlong[][]にして一度に渡します。

Application.java

package jp.co.atware.ac2021.jni.app;

...

public class Application {

    public static void main(String... args) {
        System.out.println("JNI performance by array");
        var nativeApi = new NativeApi();

        // prepare
        var size = Integer.parseInt(args[0]);

        // perform
        var start = LocalTime.now();

        var dataStream = Data.generateDataStream(size);
        var data = new long[size][];
        var dataList = dataStream.collect(toList());
        for (var i = 0; i < size; i++) {
            data[i] = dataList.get(i).toPrimitiveArray();
        }
        var res = nativeApi.calcByArray(data);

        var end = LocalTime.now();

        // result
        var dur = Duration.between(start, end).toMillis() / 1000.0;
        System.out.printf("%sS (start: %s / end: %s)\n", dur, start, end);
        System.out.printf("calculated value: %d\n", res);
    }
}

Data.java

package jp.co.atware.ac2021.jni.lib;

...

public class Data {

    private final List<Long> values;

    private Data(List<Long> values) {
        this.values = values;
    }

    public long[] toPrimitiveArray() {
        var raw = new long[values.size()];
        IntStream.range(0, values.size()).forEach(i -> raw[i] = values.get(i));
        return raw;
    }

    public static Stream<Data> generateDataStream(int size) {
        return IntStream.range(0, size)
            .map(i -> i % 10 + 1)
            .mapToObj(i -> IntStream.range(0, i).mapToLong(v -> i).boxed().collect(toList()))
            .map(Data::new);
    }
}

NativeApi.java

package jp.co.atware.ac2021.jni.lib;

public class NativeApi {

    public NativeApi() {
        System.loadLibrary("jni-performance-sample");
    }

    public native long calcByArray(long[][] data);
}

native側は単純にすべてのvalueを足して返すだけにします。(headerは省略。)

#include "jp_co_atware_ac2021_jni_lib_NativeApi.h"

#include <jni.h>

jlong Java_jp_co_atware_ac2021_jni_lib_NativeApi_calcByArray(JNIEnv* env, jobject, jobjectArray data) {

    jlong result = 0;
    auto dataSize = env->GetArrayLength(data);

    for (jsize i = 0; i < dataSize; ++i) {
        auto d = static_cast<jlongArray>(env->GetObjectArrayElement(data, i));
        auto ds = env->GetArrayLength(d);
        auto dPtr = env->GetLongArrayElements(d, JNI_FALSE);
        for (jsize di = 0; di < ds; ++di) {
            result += dPtr[di];
        }
        env->ReleaseLongArrayElements(d, dPtr, JNI_ABORT);
    }
    return result;
}

これで計測します。5回実行の平均は 0.4064 secs でした。

Data (collection) size = 1,000,000
JVM heap = 256MB

まぁ、こんなもんでしょうか。

メモリ……が……足りない……

1M(1,000,000)のcollectionはなんとかなりましたが、実際のプロダクトではこの数倍~数十倍のデータを扱います。実際data sizeが2Mになるとlist -> arrayの部分でout of memoryで落ちます。

Application.java

for (var i = 0; i < size; i++) {
    data[i] = dataList.get(i).toPrimitiveArray();
}

十数年前のPCでもないしメモリはあるので、とりあえずheapを増やしておきます。(デプロイして使っている環境では1,024MBで運用しています。)

Data size = 2,000,000
JVM heap = 1,024MB

同じように5回実行し、平均は 0.6266 secs。だんだんdata sizeを増やしていきます。

data size	avg.
3M	0.9039
5M	1.4679
10M	NaN

10MでまたOOMになりました。このくらいのdata sizeもまだ扱うレベルなので困ります。
(実際の問題はちょっと違っていて、OOMで落ちず、list -> arrayの箇所で処理が止まったままになり、native callはされず……という状態でした。)

Streamのままnativeで消費する

arrayを使うのはメモリ効率が悪いのと、また一度Listにする必要も実際にはなかったので、Streamのままnativeでaccumlateするように実装を変更しました。Streamのままと言いましたが、単純にIteratorにして回すだけですが、途中で長大なListを作らずに回すタイミングで遅延評価(= create Data instance)できるのでメモリ効率はよくなります。Stream APIができてこういうことがやりやすくなりましたね。

では実装を追加、変更します。(抜粋)

Application.java

    public static void main(String... args) {
        ...
        var dataStream = Data.generateDataStream(size);
        var res = nativeApi.calcByStream(StreamIterator.of(dataStream));
        ...
    }

StreamIterator.java

package jp.co.atware.ac2021.jni.lib;

...

public class StreamIterator {

    private final Iterator<Data> iter;

    private StreamIterator(Stream<Data> dataStream) {
        this.iter = dataStream.iterator();
    }

    public static StreamIterator of(Stream<Data> dataStream) {
        return new StreamIterator(dataStream);
    }

    public boolean hasNext() {
        return iter.hasNext();
    }

    public long[] next() {
        return iter.next().toPrimitiveArray();
    }
}

NativeApi.java

    public native long calcByStream(StreamIterator iter);

jlong Java_jp_co_atware_ac2021_jni_lib_NativeApi_calcByStream(JNIEnv* env, jobject, jobject iter) {
    auto streamIterator = env->FindClass("jp/co/atware/ac2021/jni/lib/StreamIterator");
    auto hasNext = env->GetMethodID(streamIterator, "hasNext", "()Z");
    auto next = env->GetMethodID(streamIterator, "next", "()[J");
    jlong result = 0;

    while (env->CallBooleanMethod(iter, hasNext)) {
        auto d = static_cast<jlongArray>(env->CallObjectMethod(iter, next));
        auto ds = env->GetArrayLength(d);
        auto dPtr = env->GetLongArrayElements(d, JNI_FALSE);
        for (jsize di = 0; di < ds; ++di) {
            result += dPtr[di];
        }
        env->ReleaseLongArrayElements(d, dPtr, JNI_ABORT);
        env->DeleteLocalRef(d);
    }
    return result;
}

オーバヘッドが大きそうですがこれで計測します。先の結果を併せて比較します。

JVM heap = 256MB

Data Size	avg. - array	avg. - stream
1M	0.4064	0.4923
2M	NaN	0.8762

JVM heap = 1,024MB

Data Size	avg. - array	avg. - stream
1M	0.3475	0.5843
2M	0.6266	0.9702
3M	0.9039	1.3964
5M	1.4679	2.1834
10M	NaN	4.1892

遅いですね。プロダクトでは比較計測はせず(できず)、この実装で処理が進んでちゃんと結果を得られるようになったでここまでにしました。
今回記事を書くモチベーションとして、Stream版の実装のパフォーマンスを見たかったので、その結果を得られました。やはりJava-nativeのコミュニケーションが多いほどパフォーマンスが落ちますね。まぁ順当な結果と言えるでしょう。

(数百万回native -> JVMのmethod callをするので、正直劇的に遅くなるかな？と思っていましたが、そこまでじゃないという印象。)

Java-nativeのコミュニケーションを減らす

Streamのままnativeとのコミュニケーションを減らすために、nativeで必要になる値をchunkにしつつStreamのまま回します。

Application.java

    public static void main(String... args) {
        ...
        // prepare
        var size = Integer.parseInt(args[0]);
        var chunkSize = Integer.parseInt(args[1]);

        // perform
        var start = LocalTime.now();

        var dataStream = Data.generateDataStream(size);
        var res = nativeApi.calcByStreamChunk(StreamIterator.of(dataStream, chunkSize));
        ...
    }

StreamIterator.java

    private final Iterator<Data> iter;
    private final int chunkSize;

    private StreamIterator(Stream<Data> dataStream, int chunkSize) {
        this.iter = dataStream.iterator();
        this.chunkSize = chunkSize;
    }
    public static StreamIterator of(Stream<Data> dataStream, int chunkSize) {
        return new StreamIterator(dataStream, chunkSize);
    }

    public long[][] nextChunk() {
        var chunk = new ArrayList<long[]>(chunkSize);
        while (iter.hasNext() && chunk.size() < chunkSize) {
            chunk.add(iter.next().toPrimitiveArray());
        }
        return chunk.stream().toArray(long[][]::new);
    }

jlong Java_jp_co_atware_ac2021_jni_lib_NativeApi_calcByStreamChunk(JNIEnv* env, jobject, jobject iter) {
    auto streamIterator = env->FindClass("jp/co/atware/ac2021/jni/lib/StreamIterator");
    auto hasNext = env->GetMethodID(streamIterator, "hasNext", "()Z");
    auto nextChunk = env->GetMethodID(streamIterator, "nextChunk", "()[[J");
    jlong result = 0;

    while (env->CallBooleanMethod(iter, hasNext)) {
        auto dc = static_cast<jobjectArray>(env->CallObjectMethod(iter, nextChunk));
        auto dcs = env->GetArrayLength(dc);
        for (jsize dci = 0; dci < dcs; ++dci) {
            auto d = static_cast<jlongArray>(env->GetObjectArrayElement(dc, dci));
            auto ds = env->GetArrayLength(d);
            auto dPtr = env->GetLongArrayElements(d, JNI_FALSE);
            for (jsize di = 0; di < ds; ++di) {
                result += dPtr[di];
            }
            env->ReleaseLongArrayElements(d, dPtr, JNI_ABORT);
            env->DeleteLocalRef(d);
        }
        env->DeleteLocalRef(dc);
    }
    return result;
}

chunk sizeを変えながら計測した結果はこちら。これも同様に5回実行の平均。

JVM heap = 1,024MB

data size	avg. - chunk 1K	avg. - chunk 10K	avg. - chunk 50K	chunk 100K
1M	0.3942	0.4145	0.4402	0.4402
2M	0.6294	0.6898	0.7276	0.7352
3M	0.8426	0.9282	1.0214	1.0116
5M	1.2692	1.4232	1.5586	1.5588
10M	2.3632	2.7088	2.9113	2.9528

chunkが小さい方がパフォーマンスがいい感じですね。ちょっと意外。
これまでの結果と併せて3つの実装の比較をしてみます。

Data Size	avg. - array	avg. - stream	avg. - chunk 1K
1M	0.3475	0.5843	0.3942
2M	0.6266	0.9702	0.6294
3M	0.9039	1.3964	0.8426
5M	1.4679	2.1834	1.2692
10M	NaN	4.1892	2.3632

data size 1Mではarray ver.がいいですが、2Mではほぼ同等、3M~ではchunk ver.に軍配が上がりました。

おわりに

今回はnative側の実装をStreamを回す部分だけを実装し、native -> JVMのAPI callの部分ではexception handlingはしていません。(env->ExceptionCheck()、env->ExceptionOccurred()など。) この辺りをちゃんと実装するとまた少し結果は変わるでしょう。
この結果から今のプロダクトでの実装がまだパフォーマンスを改善できそうなことがわかったので、あとで修正チケットをあげようと思います。

Note

environment
- Ubuntu 20.04
- Intel Core i9-9900K 3.6GHz
- memory 64GB
- OpenJDK 11.0.2
source https://github.com/kikuchi-m/ac2021-jni-sample
Java Native Interface Specification Contents (Java 11)