Vengineerの妄想

人生を妄想しています。

nGraphのMLIR

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

MLIRがLLVMに統合されて、覗いておいた方がいいかな?と思って、MLIRを使っているプロジェクトを探っています。

今日は、Intel nGraphです。

github.com

CMakeList.txt の中で、

の2つが MLIR 関連の変数。NGRAPH_MLIR_ENABLE が TRUEの時有効になる。

if (NGRAPH_MLIR_ENABLE)
  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNGRAPH_MLIR_ENABLE")
  set(NGRAPH_MLIR_SOURCE_DIR ${CMAKE_SOURCE_DIR}/src/contrib/mlir)
endif()

 で、src/contrib/mlir の下を取り込むということに。

if (NGRAPH_MLIR_ENABLE)
  include(cmake/external_mlir.cmake)
endif()

 にて、CMakeのありかも。こちらで、MLIRをソースコードからビルド。NGRAPH_USE_PREBUILT_MLIR が TRUE なら、ビルドはしない。

 

ソースコード全体の CMakeList.txt では、

# This must be added before any backend that uses MLIR
if (NGRAPH_MLIR_ENABLE)
  add_subdirectory(${NGRAPH_MLIR_SOURCE_DIR})
endif()

src/ngraph/runtime/cpu の CMakeLists.txt にも、

if (NGRAPH_MLIR_ENABLE)
   set(SRC

         ${SRC} 

         builder/mlir_cpu_compiled_kernel.cpp
   )
endif()

 

実際に使われている部分は、cpu_runtime_context.hpp 

#ifdef NGRAPH_MLIR_ENABLE
    /// Maps CompiledKernel nodes to their MLIR compiler
    /// The MLIR compiler caches the compiled code on the first invocation,
    /// and may in the future support re-compilation
    std::unordered_map<ngraph::op::CompiledKernel*,
    ngraph::runtime::ngmlir::MLIRCPURuntime>
    mlir_runtimes;
#endif

 で MLIRCPURuntime を使っています。この mlir_runtimes は、上に出てきた mlir_cpu_compiled_kernel.cpp の中で使われています。

MLIRCPURuntime& mlir_runtime =
        ctx->mlir_runtimes.find(compiled_kernel)->second;
// Grab the context and initialize a core compiler
mlir::MLIRContext& context = mlir_runtime.get_context();
MLIRCompiler mlir_compiler(compiled_kernel, context);
// Compile to NG dialect
mlir_compiler.compile();
// Grab a context and initialize a CPU backend using same context
MLIRCPUBackend mlir_backend(mlir_compiler.get_module(), context);
// Codegen to LLVM dialect
mlir_backend.codegen();
// Store module into runtime, and invoke.
mlir_runtime.set_module(mlir_backend.get_module());
mlir_runtime.run(&ptr_args);

 MLIRCompiler にて、カーネルコンパイルして、MLIRCPUBackedn にて、LLVM dialectにてコード生成し、実行。

nGraph => MLIR への変換は、NgDialectConversionPass の runOnModule でやっています。

 

とういうことで、

  • MLIRCompilerにて、nGraph => MLIR に変換
  • MLIRCPUBackend にて、MLIR => LLVM に変換
  • MLIRCPURuntime にて、CPU で実行

でした。