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

はじめに

GraphcoreのTensorflowのコード眺めました。2020年7月15日なので、11カ月ぶり？

vengineer.hatenablog.com

このブログに書いた内容さっぱり覚えていない。。。とりあえず、記録として。

Compiler

HloRunner::CreateExecute メソッドから呼び出されるのが、下記の2つのメソッド。

backend().compiler()->Compile メソッド
backend().compiler()->Backend メソッド

Compilerは、PoplarCompilerクラスで Compile メソッドは、ここにあります。 RunHloPasses メソッドの後に、RunBackendメソッドを実行して、executable を生成しています。

StatusOr<std::vector<std::unique_ptr<Executable>>> PoplarCompiler::Compile(
    std::unique_ptr<HloModuleGroup> module_group,
    std::vector<std::vector<se::StreamExecutor*>> stream_exec,
    se::DeviceMemoryAllocator* device_allocator) {
  TENSORFLOW_TRACEPOINT();
  if (module_group->empty()) {
    return std::vector<std::unique_ptr<Executable>>();
  }
  if (module_group->size() > 1) {
    return tensorflow::errors::Unimplemented(
        "Compilation of multiple HLO modules is not supported on Poplar.");
  }
  if (stream_exec.size() != 1 || stream_exec[0].size() != 1) {
    return tensorflow::errors::Unimplemented(
        "Unexpected number of StreamExecutor's.");
  }
  auto hlo_modules = module_group->ConsumeModules();
  TF_ASSIGN_OR_RETURN(auto module,
                      RunHloPasses(std::move(hlo_modules[0]), stream_exec[0][0],
                                   device_allocator));
  TF_ASSIGN_OR_RETURN(
      auto executable,
      RunBackend(std::move(module), stream_exec[0][0], device_allocator));
  std::vector<std::unique_ptr<Executable>> ret;
  ret.push_back(std::move(executable));
  return std::move(ret);
}

RunHloPasses メソッドは何もしていません。

StatusOr<std::unique_ptr<HloModule>> PoplarCompiler::RunHloPasses(
    std::unique_ptr<HloModule> module,
    perftools::gputools::StreamExecutor* executor,
    se::DeviceMemoryAllocator* device_allocator) {
  TENSORFLOW_TRACEPOINT();
  return std::move(module);
}

代わりに、RunBackend メソッドは巨大です。744行もあります。

ここでプログラムをコンパイルして、poplar::Executableに変換しています。

      poplar::Executable exec =
          poplar::compileGraph(main_graph, progs, opt_flags, progress_logging);

プログラムは、progs です。下記のように、

visitor.GetHostToDevice()
main_program
visitor.GetDeviceToHost()

を登録しています。

    // =======================================================================
    // DO NOT CHANGE THE ORDER OF THESE WITHOUT UPDATING PoplarProgramType IN
    // poplar_executor.h
    // =======================================================================
    progs.push_back(visitor.GetHostToDevice());
    progs.push_back(main_program);
    progs.push_back(visitor.GetDeviceToHost());

    // For verified transfers fuse all 3 programs in a single one.
    if (poplar_executor->UseVerifiedTransfers()) {
      poplar::program::Sequence fused_program({}, "FusedProgram");
      for (auto& prog : progs) {
        fused_program.add(prog);
      }
      progs.clear();
      progs.push_back(fused_program);
    }

Executable

Executabe::ExecuteOnStream メソッド

StatusOr<ExecutionOutput> Executable::ExecuteOnStream(
    const ServiceExecutableRunOptions* run_options,
    std::vector<ShapeTree<xla::MaybeOwningDeviceMemory>> arguments,
    HloExecutionProfile* hlo_execution_profile) {
  StatusOr<ExecutionOutput> result = ExecuteAsyncOnStream(
      run_options, std::move(arguments), hlo_execution_profile);
  Status blocking_status = run_options->stream()->BlockHostUntilDone();
  TF_RETURN_IF_ERROR(result.status());
  TF_RETURN_IF_ERROR(blocking_status);
  return result;
}

PoplarExecutable::ExecuteAsyncOnStream メソッドのAsyncExecuteTask構造体の operator() メソッドにて、ExecuteComputeFunction メソッドが呼ばれています。

  struct AsyncExecuteTask {
    PoplarExecutable* executable;
    ServiceExecutableRunOptions run_options;
    se::DeviceMemoryBase output_buffer;
    HloExecutionProfile* hlo_execution_profile;
    std::vector<se::DeviceMemoryBase> argument_buffers;
    std::vector<Shape> argument_shapes;
    PoplarExecutor::ArgsHandleMap args_map;
    uint64 start_time_us;

    void operator()() {
      TF_CHECK_OK(executable->ExecuteComputeFunction(
          &run_options.run_options(), &output_buffer, hlo_execution_profile,
          argument_buffers, argument_shapes, args_map, start_time_us));
    }
  };

PoplarExecutable::ExecuteComputeFunction メソッドで PoplarExecuteクラスの ExecuteEngine メソッドを呼んでいます。

Status PoplarExecutable::ExecuteComputeFunction(
    const ExecutableRunOptions* run_options,
    se::DeviceMemoryBase* result_buffer,
    HloExecutionProfile* hlo_execution_profile,
    const std::vector<se::DeviceMemoryBase>& argument_buffers,
    const std::vector<Shape>& argument_shapes,
    const PoplarExecutor::ArgsHandleMap& args_map, uint64 start_time_us) {
  TENSORFLOW_TRACEPOINT();
  VLOG(2) << "Begin asynchronous engine execution " << module().name();
  se::Stream* stream = run_options->stream();
  se::StreamExecutor* executor = stream->parent();
  PoplarExecutor* poplar_executor =
      static_cast<PoplarExecutor*>(executor->implementation());
  se::DeviceMemoryAllocator* memory_allocator = run_options->allocator();

  poplar_executor->ExecuteEngine(result_buffer, executor, *this, args_map,
                                 memory_allocator, argument_buffers);

Executor

Executor は、PoplarExxecutor で、ExecuteEngine メソッドは下記のようにExecuteEngineImpメソッドを呼んでいます。

void PoplarExecutor::ExecuteEngine(se::DeviceMemoryBase* result_buffer,
                                   se::StreamExecutor* executor,
                                   PoplarExecutable& executable,
                                   const ArgsHandleMap& args_map,
                                   se::DeviceMemoryAllocator* allocator,
                                   const Args& args) {
  TENSORFLOW_TRACEPOINT();
  std::lock_guard<std::recursive_mutex> g(ipu_.Mutex());
  if (!current_status_.ok()) {
    LOG(FATAL) << current_status_.ToString();
  }
  current_status_ = ExecuteEngineImpl(result_buffer, executor, executable,
                                      args_map, allocator, args);
}

ExecuteEngineImpl メソッドも長いです。実際に実行しているところが下記の部分です。current_engine_ は、poplar::Engineのポインタです。poplar::Engineクラスの run メソッドを実行して、IPU上でプログラムを実行しています。

      // Before executing the main program, prepare the random seeds for each
      // replica.
      seed_generator_.PrepareSeedsForReplicas(current_replication_factor_);

      // Run the main engine
      current_engine_->enableExecutionProfiling();
      current_engine_->run(PoplarProgramType::MAIN_SEQUENCE);

      StopIOThreads();

run メソッドは、

    void run(unsigned prog = 0, const std::string &debugName = "")；

のようになっています。上記では、PoplarProgramType::MAIN_SEQUENCE を渡しています。この PoplarProgramType::MAIN_SEQUENCE は、poplar_executor.h にて次のように定義されています。MAIN_SEQUENCEはenum ですが、数値だと1になります。

enum PoplarProgramType {
  HOST_TO_DEVICE,
  MAIN_SEQUENCE,
  DEVICE_TO_HOST,
};

これはprogsの1番目、main_program を実行することを意味しています。

    progs.push_back(visitor.GetHostToDevice());
    progs.push_back(main_program);
    progs.push_back(visitor.GetDeviceToHost());

おわりに

Graphcore の Tensorflowのコードを眺め直しました。

TensorFlow XLA から Poplar API に置き換えてる感じです。Poplar APIを直接たたくのに比べて、かなりオーバーヘッドがありますね。

Vengineerの戯言

人生は短いけど、長いです。人生を楽しみましょう！

GraphcoreのTensorFlow XLAを覗いてみた。

はじめに

Compiler

Executable

Executor

おわりに