Interface¶

执行引擎的核心接口如下:

virtual void PushSync(Fn exec_fun, Context exec_ctx,
                      std::vector<VarHandle> const& const_vars,
                      std::vector<VarHandle> const& mutate_vars) = 0;

这个 API 用户将一个 函数 (exec_fun)连同它的上下文信息以及依赖关系 push 到执行引擎. exec_ctx 是 exec_fun 执行的上下文环境. const_vars 代表的是函数只有读取权限的变量, mutate_vars 表示的是函数可以修改的变量. 先不考虑具体的细节, 执行引擎保证下面的规则:

任意两个会修改同一个变量的函数,会根据它们 push 到引擎的顺序进行序列化.

Function¶

执行引擎需要的函数类型按照下面方式来定义:

using Fn = std::function<void(RunContext)>;

RunContext 包含了引擎确定的运行时信息:

struct RunContext {
    // stream pointer which could be safely cast to
    // cudaStream_t* type
    void *stream;
};

用户也可以使用 mxnet::engine::DAGEngine::Fn 定义作为第二种选择, 它们的类型是一样的.

所有的函数都会被 engine 内部的线程来执行. 在这个模型中, 我们不鼓励用户将 阻塞 (blocking) 函数 push 到引擎 ( 通常是处理 I/O 任务的函数, 比如读取硬盘, web 服务, UI, 等等). 因为阻塞函数会占用执行线程, 同时降低了这个系统的吞吐量. 这种情况下, 我们提供了另外的一种 asynchronous 函数类型:

using Callback = std::function<void()>;
using AsyncFn = std::function<void(RunContext, Callback)>;

在 AsyncFn 函数中, 用户可以将重要的计算交由自己的线程来执行, 同时不用等待函数执行结束. 除非异步函数的 Callback 被调用, 否则引擎不会考虑函数是否已经结束的事情.

Context¶

用户可以指定函数执行的需要的 Context. 这个 Context 一般包括函数是否执行在 CPU 或者 GPU 上, 如果是 GPU, 那么具体是哪个 GPU. Context 和 RunContext 是不一样的. Context 包括设备类型 (gpu/cpu) 和设备 id, 而 RunContext 包含的是只有在运行时才可以确定的信息, 比如说函数要在哪个 stream 上执行.

VarHandle¶

VarHandle 是 用来指定函数的依赖关系的. MXNet 执行引擎的设计目的是为了接口和 MXNet 的其他模块解耦合. 所以 VarHandle 类似引擎为用户提供用来代表函数需要或者会修改的外部资源的一个令牌. 它被设计成轻量级的, 所以创建, 删除或者拷贝一个变量只需要一点点开销. 对于正在推送到引擎的函数, 用户需要在 const_vars vector 里指定需要的不可变变量, 在 mutate_vars vector 里指定会被修改的变量. 执行引擎解析函数之间的依赖关系唯一的规则是:

**任意两个会修改同一个变量的函数,会根据它们 push 到引擎的顺序进行序列化.*

举个例子, 如果 Fn1和Fn2都要修改 V2, 那么如果 Fn2 后于 Fn1 push 到引擎, 那么引擎会保证 Fn2 会在 Fn1 之后执行. 另一方面, 如果 Fn1 和 Fn2 都使用但不修改 V2, 那么它们具体的执行顺序是任意的.

这种设计方式可以允许引擎可以调度 状态改变(state-mutating) 操作. 比如说, 在 DNN 中权重更新的函数可以用 += 操作来原地更新权重, 而不是每次都产生一个新的权重数组.

如果要创建一个变量, 使用 NewVar() API. 如果要删除一个变量, 使用 PushDelete API.

Push & Wait¶

所有的 Push API 是异步的. API 调用会在调用之后马上返回, 而不管 Fn 是否执行完与否. 这允许执行引擎可以在用户的线程推送函数到引擎的时候马上开始计算. 所有的 Push API 不是 thread-safe 的. 具体来说, 就是应该只有一个线程来调用 API.

如果你想要等待一个具体的 Fn 完成, 需要包含一个 callback 函数, 然后在你的 Fn 的最后调用前面的回调函数.

如果你要等待所有的对一个确定的变量的使用 (读取/修改) 的 Fn, 那么你应该调用 WaitForVar(var) API.

如果你要等待所有推送到引擎的 Fn 都结束, 那么需要调用 WaitForAll() API.

Save Object Creation Cost¶

有些情况下, 你需要推送几个函数到引擎很多次. 如果这些函数的计算是轻量级的, 那么拷贝 Lambda 表达式和创建读/写变量的列表的开销的比例就会相当高. 我们提供了一个 API 来提前创建 OprHandle:

virtual OprHandle NewOperator(AsyncFn fn,
                              std::vector<VarHandle> const& const_vars,
                              std::vector<VarHandle> const& mutate_vars) = 0;

这样你就可以一直推送 OprHandle, 而不用每次都要创建:

virtual void Push(OprHandle op, Context exec_ctx) = 0;

可以通过调用 DeleteOperator(OprHandle op) 来删除它. 不过一定要确保它们已经执行完了.

API Reference¶

class mxnet::Engine¶

Dependency engine that schedules operations.

Public Types

typedef engine::CallbackOnComplete CallbackOnComplete¶

callback on complete

typedef std::function<void(RunContext)> SyncFn¶

Synchronous operation to pass to engine.

typedef std::function<void(RunContext, CallbackOnComplete)> AsyncFn¶

Asynchronous operation to pass to engine.

typedef engine::VarHandle VarHandle¶

Variable pointer.

typedef engine::OprHandle OprHandle¶

Operator pointer.

Public Functions

virtual void NotifyShutdown() = 0¶

Notify the engine about a shutdown, This can help engine to print less messages into display.

User do not have to call this function.

Return

0 when success, -1 when failure happens.

virtual VarHandle NewVariable() = 0¶

Allocate a new variable, the variable can then be used to schedule the operation concurrently via dependency patterns.

Return

The new variable allocated.

virtual OprHandle NewOperator(AsyncFn fn, std::vector<VarHandle> const &const_vars, std::vector<VarHandle> const &mutable_vars, FnProperty prop = FnProperty::kNormal) = 0¶

Create a new operator. The returned operator could be saved externally so that it could be resued for scheduling.

Return

The new operator allocated.

Parameters

• fn -

  The execution function.

• const_vars -

  The variables that current operation will use but not mutate.

• mutable_vars -

  The variables that current operation will mutate.

• prop -

  Property of the function.

virtual void DeleteOperator(OprHandle op) = 0¶

Delete the given operator.

The delete will not happen immediately, but will wait until all the operations using this operator are completed.

Parameters

• op -

  The operator to delete.

virtual void Push(OprHandle op, Context exec_ctx, int priority = 0) = 0¶

Push an operator to the engine.

Parameters

• op -

  The operator to push.

• exec_ctx -

  Execution context.

• priority -

  Priority of the action, as hint to the engine.

virtual void PushAsync(AsyncFn exec_fun, Context exec_ctx, std::vector<VarHandle> const &const_vars, std::vector<VarHandle> const &mutable_vars, FnProperty prop = FnProperty::kNormal, int priority = 0) = 0¶

Push an asynchronous operation to the engine.

Parameters

• exec_fun -

  Execution function, this function takes a parameter on_complete that must be called when the execution completes.

• exec_ctx -

  Execution context.

• const_vars -

  The variables that current operation will use but not mutate.

• mutable_vars -

  The variables that current operation will mutate.

• prop -

  Property of the function.

• priority -

  Priority of the action, as hint to the engine.

virtual void DeleteVariable(SyncFn delete_fn, Context exec_ctx, VarHandle var) = 0¶

Schedule the deletion of a variable.

The delete will not happen immediately, but will wait until all the operations depending on var are completed.

Parameters

• delete_fn -

  A function that will be called after the variable is deleted.

• exec_ctx -

  Execution context.

• var -

  The variable to be deleted.

virtual void WaitForVar(VarHandle var) = 0¶

Wait for a variable.

Parameters

• var -

  The variable we should wait for. This function returns when the variable is ready.

virtual void WaitForAll() = 0¶

Wait until all the activity of engine finishes.

virtual ~Engine()¶

virtual destructor

template <typename SyncFn>

void PushSync(SyncFn exec_fn, Context exec_ctx, std::vector<VarHandle> const &const_vars, std::vector<VarHandle> const &mutable_vars, FnProperty prop = FnProperty::kNormal, int priority = 0)¶

Push an synchronous operation to the engine.

Parameters

• exec_fn -

  Execution function that executes the operation.

• exec_ctx -

  Execution context.

• const_vars -

  The variables that current operation will use but not mutate.

• mutable_vars -

  The variables that current operation will mutate.

• prop -

  Property of the function.

• priority -

  Priority of the action, as hint to the engine.

Template Parameters

• SyncFn -

  the synchronous function to be pushed.

Public Static Functions

static Engine *Get()¶

Return

Engine singleton.

static std::shared_ptr<Engine> _GetSharedRef()¶

Get shared pointer reference to engine singleton. Most user should not call this function. This function is called by another singleton X who requires engine to be destructed after X.

Return

A shared pointer to Engine singleton.