Tensorflow to WebAssembly

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

Tensorflow to WebAssembly

TensorFlow XLA の tfcompile を利用して、WebAssembly に変換するというもの。

wasm.patch.diffを見てみると、

基本的には、tensorflow/compiler/xla/service/cpu/cpu_compiler.ccに、
LLVM内のWebAssemblyを出力するコードを追加するだけのお仕事。

こんな感じに、

引用
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
index d039132..7047d30 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
@@ -158,6 +158,11 @@ CpuCompiler::CpuCompiler() {
   LLVMInitializeAArch64TargetMC();
   LLVMInitializeAArch64AsmPrinter();
   LLVMInitializeAArch64Disassembler();
+  LLVMInitializeWebAssemblyTarget();
+  LLVMInitializeWebAssemblyTargetInfo();
+  LLVMInitializeWebAssemblyTargetMC();
+  LLVMInitializeWebAssemblyAsmPrinter();
+  LLVMInitializeWebAssemblyDisassembler();
 }

CpuCompilerのコンストラクタにて、WebAssembly関連のコードを初期化。

それ以外には、BUILDファイルに、LLVMのTargetとに、WebAssemblyを追加するだけよ。

引用
diff --git a/third_party/llvm/llvm.BUILD b/third_party/llvm/llvm.BUILD
index e1c22c8..78f38f9 100644
--- a/third_party/llvm/llvm.BUILD
+++ b/third_party/llvm/llvm.BUILD
@@ -30,6 +30,7 @@ llvm_targets = [
     "ARM",
     "NVPTX",
     "PowerPC",
+    "WebAssembly",
     "X86",
 ]

いろいろ、続く。

そんでもって、tfcompile コマンドで以下のようなパラメータを指定して実行すると、

引用
$ bazel-bin/tensorflow/compiler/aot/tfcompile \
    --target_triple="wasm32-unknown-unknown-wasm" \
    --target_cpu="generic" \
    --xla_cpu_multi_thread_eigen=false \
    --graph=/path/to/my_graph.pb \
    --config=/path/to/my_config.pbtxt \
    --out_function_object=out_model.o \
    --out_header=out_header.h \
    --out_metadata_object=out_helper.o \
    --cpp_class=MyClass

out_model.o と out_helper.o というオブジェクトが生成されるって、
ポイントは、--xla_cpu_multi_thread_eigen=false　にすること、何故なら、Single Thread でしか動かないので。

上手くやっているよね。