OneFlow源码解析：算子指令在虚拟机中的执行

撰文｜郑建华、赵露阳

1 

Op在虚拟机里的执行

1.1 PhysicalRun和InstructionsBuilder

上一篇文章《OneFlow源码解析：Op、Kernel与解释器》中提到：

PhysicalRun接受一个lambda函数作为参数，这里即InstructionsBuilder->Call方法，该方法接受kernel、input/output的eager blob object、kernel执行的上下文作为参数。Call方法实际会完成OpCall指令的构建，并最终将其派发至vm指令列表中，等待VM实际调度执行。

这个PhysicalRun函数里包裹着一个lambda函数：

JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
    return builder->Call(xxx);
}));

其中，lambda函数接受一个InstructionsBuilder指针(builder)，并调用builder->Call方法，用于实际完成Op指令在VM中的构建。而PhysicalRun（https://github.com/Oneflow-Inc/oneflow/blob/88f147d50e75d1644e552ed445dd58f9b5121ea5/oneflow/core/framework/instructions_builder.h#L160）在 oneflow/core/framework/instructions_builder.h中定义，其接受lambda函数作为模版参数(CallbackT)：

// Make VM instructions with instruction builder and run instructions with physical/local view.
template<typename CallbackT>
Maybe<void> PhysicalRun(const CallbackT& Build) {
  vm::InstructionList instruction_list;
  InstructionsBuilder instructions_builder(&instruction_list);
  JUST(Build(&instructions_builder));
  JUST(vm::Run(instructions_builder.mut_instruction_list()));
  return Maybe<void>::Ok();
}

可见，PhysicalRun函数中，首先初始化一个InstructionsBuilder，然后将InstructionsBuilder指针作为参数传给lambda函数，完成实际指令的构建；最后通过vm::Run()方法将该指令发送至VM，等候VM实际调度和执行。Run方法如下：

Maybe<void> Run(vm::InstructionList* instruction_list) {
  auto* virtual_machine = JUST(SingletonMaybe());
  JUST(virtual_machine->Receive(instruction_list));
  return Maybe<void>::Ok();
}

可以看见，Run()方法获取了全局单例的VM对象指针，然后通过vm的Receive()方法，将该条指令发送给虚拟机（所以这里Run其实有点歧义，更贴切的意思，其实是指令发送或传送）。

这个VirtualMachine->Receive方法很重要，会在后面的第2.章节中详细展开。

1.2 InstructionsBuilder

上面PhysicalRun函数中的InstructionsBuilder，类似一个指令构建的helper，InstructionsBuilder的系列方法配合指令策略(InstructionPolicy)，可以帮助构建不同类型的vm指令。

从InstructionsBuilder

（https://github.com/Oneflow-Inc/oneflow/blob/88f147d50e75d1644e552ed445dd58f9b5121ea5/oneflow/core/framework/instructions_builder.h#L47）的定义中，我们可以看到指令的构建方法，其中常用方法如下：

// 用于lazy mode(nn.Graph)
// Build VM execution instructions with NNGraph's inputs/outputs/parameters for NNGraph execution.
Maybe<void> LaunchLazyJob(const vm::EagerBlobObjectListPtr& inputs,
                          const vm::EagerBlobObjectListPtr& outputs,
                          const vm::EagerBlobObjectListPtr& parameters,
                          const std::shared_ptr& nn_graph);
 
 
 
 
// 用于全局同步，同步等待所有指令调用完成
Maybe<void> GlobalSync();
 
 
// 用于Tensor内存释放(归还allocator)
Maybe<void> ReleaseTensor(const std::shared_ptr& eager_blob_object);
 
 
// 操作Tensor实际内存（blob）
template<typename T>
Maybe<void> AccessBlobByCallback(
    const T tensor,
    const std::function<void(ep::Stream*, const std::shared_ptr&)>& callback,
    const std::string& modifier);
 
 
// 最常用的指令构建方法，用于构造op执行所需的OpCall指令
Maybe<void> Call(const std::shared_ptr& opkernel,
                   vm::EagerBlobObjectList&& input_eager_blob_objects,
                   vm::EagerBlobObjectList&& output_eager_blob_objects,
                   const one::OpExprInterpContext& ctx, Symbol stream);

1.3 InstructionPolicy

InstructionPolicy

（https://github.com/Oneflow-Inc/oneflow/blob/88f147d50e75d1644e552ed445dd58f9b5121ea5/oneflow/core/vm/instruction_policy.h#L34）——指令策略，通常用于配合InstructionsBuilder实际构建出不同的vm指令。InstructionPolicy的子类实现如下：

这些子类的InstructionPolicy可近似认为是指令类型。如，用于Op执行的OpCallInstructionPolicy、用于Tensor内存释放的ReleaseTensorInstructionPolicy、用于屏障阻塞的BarrierInstructionPolicy等。

以Op执行为例：

JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
    return builder->Call(xxx);
}));

实际上是通过InstructionsBuilder的Call方法

（https://github.com/Oneflow-Inc/oneflow/blob/88f147d50e75d1644e552ed445dd58f9b512