Isa@Diary

ソフトウェア開発やってます。プログラミングとか、US生活とかについて書きます。

MKLをJNI経由で呼ぶ

MKLを使ったC++ライブラリをJNI経由でJavaから呼ぼうとしたら初めうまくいかなかったものの解決したので書く。 C++ ライブラリはこんな感じ。ビルド時のオプションは https://software.intel.com/en-us/articles/intel-mkl-link-line-advisor で作った。

>ldd ./build/lib/libhoge.so
        linux-vdso.so.1 =>  (0x00007fff2c329000)
        libpthread.so.0 => /lib64/libpthread.so.0 (0x00007fc890ac9000)
        libdl.so.2 => /lib64/libdl.so.2 (0x00007fc8908c5000)
        libmkl_core.so => /build/lib/libmkl_core.so (0x00007fc88e732000)
        libmkl_intel_lp64.so => /build/lib/libmkl_intel_lp64.so (0x00007fc88dc44000)
        libmkl_intel_ilp64.so => /build/lib/libmkl_intel_ilp64.so (0x00007fc88d215000)
        libmkl_gnu_thread.so => /build/lib/libmkl_gnu_thread.so (0x00007fc88bb57000)
        libgomp.so.1 => /usr/lib64/libgomp.so.1 (0x00007fc88b94a000)
        libm.so.6 => /lib64/libm.so.6 (0x00007fc88b6c6000)
        libstdc++.so.6 => /usr/lib64/libstdc++.so.6 (0x00007fc88b341000)
        libgcc_s.so.1 => /lib64/libgcc_s.so.1 (0x00007fc88b12b000)
        libc.so.6 => /lib64/libc.so.6 (0x00007fc88ad97000)
        /lib64/ld-linux-x86-64.so.2 (0x00007fc890ee8000)
        librt.so.1 => /lib64/librt.so.1 (0x00007fc88ab8f000)

JNIライブラリは

> ldd ./build/lib/libhogejni.so
        linux-vdso.so.1 =>  (0x00007fffe6eda000)
        libpthread.so.0 => /lib64/libpthread.so.0 (0x00007f3e5a12e000)
        libdl.so.2 => /lib64/libdl.so.2 (0x00007f3e59f2a000)
        libhoge.so => not found
        libiomp5.so => not found
        libmkl_core.so => not found
        libmkl_intel_ilp64.so => not found
        libmkl_gnu_thread.so => not found
        libm.so.6 => /lib64/libm.so.6 (0x00007f3e59ca6000)
        libstdc++.so.6 => /usr/lib64/libstdc++.so.6 (0x00007f3e59921000)
        libgcc_s.so.1 => /lib64/libgcc_s.so.1 (0x00007f3e5970b000)
        libc.so.6 => /lib64/libc.so.6 (0x00007f3e59377000)
        /lib64/ld-linux-x86-64.so.2 (0x00007f3e5a34b000)

こっちは当初はMKL関連のライブラリはリンクしてなかったのだが、SOでリンクすればいいよと言われたのでリンクした。結果としてはリンクしてもダメだった。

JNIライブラリへのSystem.loadLibrary自体は成功するのだが、実際の関数を呼ぼうとすると

Intel MKL FATAL ERROR: Cannot load libmkl_avx2.so or libmkl_def.so

とあまり役に立たないエラーだけでる。このエラーで検索するとPythonでMKLのインストールがうまくいかない、みたいなのばっかりでる。

解決しないのでlibmkl_avx2.so自体をloadLibraryしようとすると

<path_to_lib>/libmkl_avx2.so: <path_to_lib>/libmkl_avx2.so:
undefined symbol: mkl_sparse_optimize_bsr_trsm_i8

というエラーが出る。このシンボルはlibmkl_gnu_thread.so内にある。(他にlibmkl_core.so内にあるものもあった)

>nm <path_to_lib>/libmkl_gnu_thread.so | grep mkl_sparse_optimize_bsr_trsm_i8
00000000004fe240 T mkl_sparse_optimize_bsr_trsm_i8

そこで先にlibmkl_gnu_thread.soをロードしておけばいいんじゃないかと思って試してみるとダメ。

SOで聞いてみるとJNIライブラリはロードの際にdlopenに渡すflagがRTLD_LOCALという後のライブラリロード時のシンボル解決には使わないという状態になっている(defaultの挙動)らしい。

これを調べるために下記のコードをgcc -shared -fPIC dlopen_trace.c -o dlopen_trace.so -ldlとやってできたライブラリをLD_PRELOADに指定することでdlopenにhookした。初めはRTLD_GLOBALを加えずに何が渡されているか見ていたのだが、大体のライブラリはRTLD_LAZYだけでロードされているようだった。C++で作ったexecutableの場合はリンクされているライブラリに対してはdlopen自体呼ばれていないので何か他の方法でロードしてるっぽい。

> less dlopen_trace.c
#define _GNU_SOURCE
#include <dlfcn.h>
#include <stdio.h>
#include <string.h>

typedef void *(*orig_dlopen_type)(const char *file, int mode);

char mkl_gnu_thread_so_name[1000] = "/build/lib/libmkl_gnu_thread.so\0";
char mkl_core_so_name[1000] = "/build/lib/libmkl_core.so\0";

void *dlopen(const char *file, int mode)
{
  fprintf(stderr, "dlopen called (mode: %d) on %s\n", mode, file);


  if ( mode & RTLD_LAZY ) {
    fprintf(stderr, "RTLD_LAZY\n");
  }

  if (mode & RTLD_NOW) {
    fprintf(stderr, "RTLD_NOW\n");
  }

  if (mode & RTLD_GLOBAL) {
    fprintf(stderr, "RTLD_GLOBAL\n");
  }

  if (mode & RTLD_LOCAL) {
    fprintf(stderr, "RTLD_LOCAL\n");
  }

  if (mode & RTLD_NODELETE) {
    fprintf(stderr, "RTLD_NODELETE\n");
  }

  if (mode & RTLD_NOLOAD) {
    fprintf(stderr, "RTLD_NOLOAD\n");
  }

  if (mode & RTLD_DEEPBIND) {
    fprintf(stderr, "RTLD_DEEPBIND\n");
  }

  if (file != NULL) {
    if (strcmp(mkl_gnu_thread_so_name, file) == 0 || strcmp(mkl_core_so_name, file) == 0) {
      fprintf(stderr, "Adding RTLD_GLOBAL\n");
      mode |= RTLD_GLOBAL;
    }
  }
  orig_dlopen_type orig_dlopen;
  orig_dlopen = (orig_dlopen_type)dlsym(RTLD_NEXT, "dlopen");
  return orig_dlopen(file, mode);
}

結果としてRTLD_GLOBALを足してやるとlibmkl_avx2.soもちゃんとJNIを使ってロードできた。実際にはJNIのライブラリにfunctionを1個生やしてdlopenを呼んでやることにした。dlopenにはファイル名を渡してやれば通常のロード時と同じ順番でライブラリを探してくれるので便利。

参考にしたものたち: