Interface¶

The following API is the core interface for the execution engine:

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

This API allows you to push a function (exec_fun), along with its context information and dependencies, to the engine. exec_ctx is the context information in which the exec_fun should be executed, const_vars denotes the variables that the function reads from, and mutate_vars are the variables to be modified. The engine provides the following guarantee:

The execution of any two functions that modify a common variable is serialized in their push order.

Function¶

The function type of the engine is:

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

RunContext contains runtime information, which is determined by the engine:

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

Alternatively, you could use mxnet::engine::DAGEngine::Fn, which has the same type definition.

All of the functions are executed by the engine’s internal threads. In such a model, it’s usually not a good idea to push blocking functions to the engine (usually for dealing with I/O tasks like disk, web service, UI, etc.) because it will occupy the execution thread and reduce total throughput. In that case, we provide another asynchronous function type:

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

In the AsyncFn function, you can pass the heavy part to your own threads and safely exit the body of the function. The engine doesn’t consider the function finished until the Callback function is called.

Context¶

You can specify the Context of the function to be executed within. This usually includes whether the function should be run on a CPU or a GPU, and if you specify a GPU, which GPU to use. Context is different from RunContext. Context contains device type (GPU/CPU) and device id, while RunContext contains information that can be decided only during runtime, for example, on which stream the function should be executed.

VarHandle¶

VarHandle is used to specify the dependencies of functions. The MXNet engine is designed to be decoupled from other MXNet modules. So VarHandle is like an engine-provided token you use to represent the external resources the functions can use or modify. It’s designed to be lightweight, so creating, deleting, or copying a variable incurs little overhead. Upon pushing the functions, you need to specify the variables that will be used (immutable) in the const_vars vector, and the variables that will be modified (mutable) in the mutate_vars vector. The engine uses one rule for resolving the dependencies among functions:

The execution of any two functions when one of them modifies at least one common variable is serialized in their push order.

For example, if Fn1 and Fn2 both mutate V2 then Fn2 is guaranteed to be executed after Fn1 if Fn2 is pushed after Fn1. On the other hand, if Fn1 and Fn2 both use V2, their actual execution order could be random.

This design allows the engine to schedule state-mutating operations in a manner that minimizes calls to allocate new memory. For example, the weight update function in DNN can now use the += operator to update the weights in place, rather than generating a new weight array each time.

To create a variable, use the NewVar() API. To delete a variable, use the PushDelete API.

Push and Wait¶

All Push APIs are asynchronous. The API call returns immediately regardless of whether the pushed Fn is finished or not. This allows the engine to start computing at the same time as the user thread is pushing functions. Push APIs are not thread-safe. To be specific, only one thread should make engine API calls at a time.

If you want to wait for a specific Fn to finish, include a callback function in the closure, and call the function at the end of your Fn.

If you want to wait for all Fns that involve (use or mutate) a certain variable to finish, use the WaitForVar(var) API.

If you want to wait for all pushed Fns to finish, use the WaitForAll() API.

Save Object Creation Cost¶

In some cases, you need to push several functions to the engine for a long period of time. If the computation of these functions is light, the overhead of copying lambdas and creating use/mutate variable lists becomes relatively high. We provide an API to create an OprHandle beforehand:

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

You can keep pushing the OprHandle without repeatedly creating them:

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

To delete it, call the DeleteOperator(OprHandle op) API. Ensure that the operator has finished computing before calling this API.

API Reference¶

class Engine¶

Dependency engine that schedules operations.

Public Types

typedef engine::CallbackOnComplete CallbackOnComplete¶

callback on complete

typedef std::function(RunContext)> SyncFn¶

Synchronous operation to pass to engine.

typedef std::function(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 void Stop()¶

Stop all workers in the engine.

virtual void Start()¶

Restart all workers in the engine.

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, const char *opr_name = nullptr, bool wait = false) = 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.
• opr_name: The operator name.
• wait: Whether this is a WaitForVar operation

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, bool profiling = false) = 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.
• profiling: The variable indicate whether to profile this operator.

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, const char *opr_name = nullptr, bool wait = false) = 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.
• opr_name: The operator name.
• wait: Whether this is a WaitForVar operation

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 void Throw(VarHandle var) = 0¶

Throw if threre are associated exception with var.

virtual ~Engine()¶

virtual destructor

virtual 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, const char *opr_name = nullptr)¶

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.
• opr_name: The operator name.

Template Parameters

• SyncFn: the synchronous function to be pushed.

CallbackOnComplete CreateCallback(void (*callback)(Engine *, void *, const dmlc::Error *), void *param, )¶

factory function to create OnComplete callback.

Parameters

• callback: th static callback function.
• param: the paramter passed to callback.

virtual int bulk_size() const¶

query current limit for bulk size

virtual int set_bulk_size(int)¶

set maximum limit for bulk size

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.

Operator Interface¶

Forward is the core operator interface:

    virtual void Forward(const OpContext &ctx,
                         const std::vector<TBlob> &in_data,
                         const std::vector<OpReqType> &req,
                         const std::vector<TBlob> &out_data,
                         const std::vector<TBlob> &aux_states) = 0;

The OpContext structure is:

           struct OpContext {
             int is_train;
             RunContext run_ctx;
             std::vector<Resource> requested;
           }

It describes whether the operator is in the train or test phase, which device the operator should be run on (in run_ctx), and requested resources (covered in the following sections).

• in_data and out_data represent the input and output tensors, respectively. All of the tensor spaces have been allocated by the system.
• req denotes how the computation results are written into the out_data. In other words, req.size() == out_data.size() and req[i] correspond to the write type of out_data[i].
• The OpReqType is defined as:

           enum OpReqType {
             kNullOp,
             kWriteTo,
             kWriteInplace,
             kAddTo
           };

Normally, the types of all out_data should be kWriteTo, meaning that the provided out_data tensor is a raw memory block, so the operator should write results directly into it. In some cases, for example when calculating the gradient tensor, it would be great if we could accumulate the result, rather than directly overwrite the tensor contents so that no extra space needs to be created each time. In such a case, the corresponding req type is set as kAddTo, indicating that a += should be called.

• aux_states is intentionally designed for auxiliary tensors used to help computation. Currently, it is useless.

Aside from the Forward operator, you could optionally implement the Backward interface:

    virtual void Backward(const OpContext &ctx,
                          const std::vector<TBlob> &out_grad,
                          const std::vector<TBlob> &in_data,
                          const std::vector<TBlob> &out_data,
                          const std::vector<OpReqType> &req,
                          const std::vector<TBlob> &in_grad,
                          const std::vector<TBlob> &aux_states);

This interface follows the same design principle as the Forward interface, except that out_grad, in_data, and out_data are given, and the operator computes in_grad as the results. The naming strategy is similar to Torch’s convention, and can be summarized in following figure:

[input/output semantics figure]

Some operators might not require all of the following: out_grad, in_data and out_data. You can specify these dependencies with the DeclareBackwardDependency interface in OperatorProperty.

Operator Property¶

One convolution might have several implementations, and you might want to switch among them to achieve the best performance. Therefore, we separate the operator semantic interfaces from the implementation interface (Operator class) into the OperatorProperty class. The OperatorProperty interface consists of:

• InferShape:

           virtual bool InferShape(mxnet::ShapeVector *in_shape,
                                   mxnet::ShapeVector *out_shape,
                                   mxnet::ShapeVector *aux_shape) const = 0;

This interface has two purposes:

• Tell the system the size of each input and output tensor, so it can allocate space for them before the Forward and Backward call.
• Perform a size check to make sure that there isn’t an obvious error before running. The shape in in_shape is set by the system (from the out_shape of the previous operators). It returns false when there is not enough information to infer shapes or throws an error when the shape is inconsistent.
• Request Resources: Operations like cudnnConvolutionForward need a work space for computation. If the system can manage that, it could then perform optimizations, like reuse the space, and so on. MXNet defines two interfaces to achieve this:

           virtual std::vector<ResourceRequest> ForwardResource(
               const mxnet::ShapeVector &in_shape) const;
           virtual std::vector<ResourceRequest> BackwardResource(
               const mxnet::ShapeVector &in_shape) const;

The ResourceRequest structure (in resource.h) currently contains only a type flag:

           struct ResourceRequest {
             enum Type {
               kRandom,  // get a mshadow::Random object
               kTempSpace,  // request temporary space
             };
             Type type;
           };

If ForwardResource and BackwardResource return non-empty arrays, the system offers the corresponding resources through the ctx parameter in the Forward and Backward interface of Operator. Basically, to access those resources, simply write:

           auto tmp_space_res = ctx.requested[kTempSpace].get_space(some_shape, some_stream);
           auto rand_res = ctx.requested[kRandom].get_random(some_stream);

For an example, see src/operator/cudnn_convolution-inl.h.

• Backward dependency: Let’s look at two different operator signatures (we name all of the arguments for demonstration purposes):

           void FullyConnectedForward(TBlob weight, TBlob in_data, TBlob out_data);
           void FullyConnectedBackward(TBlob weight, TBlob in_data, TBlob out_grad, TBlob in_grad);

           void PoolingForward(TBlob in_data, TBlob out_data);
           void PoolingBackward(TBlob in_data, TBlob out_data, TBlob out_grad, TBlob in_grad);

Note that out_data in FullyConnectedForward is not used by FullyConnectedBackward, while PoolingBackward requires all of the arguments of PoolingForward. Therefore, for FullyConnectedForward, the out_data tensor once consumed could be safely freed because the backward function will not need it. This provides a chance for the system to collect some tensors as garbage as soon as possible. To specify this situation, we provide an interface:

          virtual std::vector<int> DeclareBackwardDependency(
               const std::vector<int> &out_grad,
               const std::vector<int> &in_data,
               const std::vector<int> &out_data) const;

The int element of the argument vector is an ID to distinguish different arrays. Let’s see how this interface specifies different dependencies for FullyConnected and Pooling:

          std::vector<int> FullyConnectedProperty::DeclareBackwardDependency(
              const std::vector<int> &out_grad,
              const std::vector<int> &in_data,
              const std::vector<int> &out_data) const {
            return {out_grad[0], in_data[0]};  // NOTE: out_data[0] is NOT included
          }
          std::vector<int> PoolingProperty::DeclareBackwardDependency(
              const std::vector<int> &out_grad,
              const std::vector<int> &in_data,
              const std::vector<int> &out_data) const {
            return {out_grad[0], in_data[0], out_data[0]};
          }

• In place Option: To further save the cost of memory allocation, you can use in-place updates. They are appropriate for element-wise operations when the input tensor and output tensor have the same shape. You specify and in-place update with the following interface:

           virtual std::vector<std::pair<int, void*>>    ElewiseOpProperty::ForwardInplaceOption(
               const std::vector<int> &in_data,
               const std::vector<void*> &out_data) const {
             return { {in_data[0], out_data[0]} };
           }
           virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::BackwardInplaceOption(
               const std::vector<int> &out_grad,
               const std::vector<int> &in_data,
               const std::vector<int> &out_data,
               const std::vector<void*> &in_grad) const {
             return { {out_grad[0], in_grad[0]} }
           }

This tells the system that the in_data[0] and out_data[0] tensors could share the same memory spaces during Forward, and so do out_grad[0] and in_grad[0] during Backward.

Important: Even if you use the preceding specification, it’s not guaranteed that the input and output tensors will share the same space. In fact, this is only a suggestion for the system, which makes the final decision. However, in either case, the decision is completely transparent to you, so the actual Forward and Backward implementation does not need to consider that.

• Expose Operator to Python: Because of the restrictions of C++, you need user to implement following interfaces:

           // initial the property class from a list of key-value string pairs
           virtual void Init(const vector<pair<string, string>> &kwargs) = 0;
           // return the parameters in a key-value string map
           virtual map<string, string> GetParams() const = 0;
           // return the name of arguments (for generating signature in python)
           virtual vector<string> ListArguments() const;
           // return the name of output values
           virtual vector<string> ListOutputs() const;
           // return the name of auxiliary states
           virtual vector<string> ListAuxiliaryStates() const;
           // return the number of output values
           virtual int NumOutputs() const;
           // return the number of visible outputs
           virtual int NumVisibleOutputs() const;

Create an Operator from the Operator Property¶

OperatorProperty includes all semantic attributes of an operation. It’s also responsible for creating the Operator pointer for actual computation.

    virtual Operator* CreateOperator(Context ctx) const = 0;

    class ConvolutionOp {
     public:
      void Forward( ... ) { ... }
      void Backward( ... ) { ... }
    };
    class ConvolutionOpProperty : public OperatorProperty {
     public:
      Operator* CreateOperator(Context ctx) const {
        return new ConvolutionOp;
      }
    };

    #include 
    struct ConvolutionParam : public dmlc::Parameter<ConvolutionParam> {
      mxnet::TShape kernel, stride, pad;
      uint32_t num_filter, num_group, workspace;
      bool no_bias;
    };

    class ConvolutionOp {
     public:
      ConvolutionOp(ConvolutionParam p): param_(p) {}
      void Forward( ... ) { ... }
      void Backward( ... ) { ... }
     private:
      ConvolutionParam param_;
    };
    class ConvolutionOpProperty : public OperatorProperty {
     public:
      void Init(const vector<pair<string, string>& kwargs) {
        // initialize param_ using kwargs
      }
      Operator* CreateOperator(Context ctx) const {
        return new ConvolutionOp(param_);
      }
     private:
      ConvolutionParam param_;
    };

    DMLC_REGISTER_PARAMETER(ConvolutionParam);
    MXNET_REGISTER_OP_PROPERTY(Convolution, ConvolutionOpProperty);

Interface Summary¶

We’ve almost covered the entire interface required to define a new operator. Let’s do a recap:

• Use the Operator interface to write your computation logic (Forward and Backward).
• Use the OperatorProperty interface to:
  • Pass the parameter to the operator class (you can use the Init interface).
  • Create an operator using the CreateOperator interface.
  • Correctly implement the operator description interface, such as the names of arguments, etc.
  • Correctly implement the InferShape interface to set the output tensor shape.
  • [Optional] If additional resources are needed, check ForwardResource and BackwardResource.
  • [Optional] If Backward doesn’t need all of the input and output of Forward, check DeclareBackwardDependency.
  • [Optional] If in-place update is supported, check ForwardInplaceOption and BackwardInplaceOption.
• Register the OperatorProperty class and the parameter class.

Define Shapes¶

The mshadow library requires explicit memory allocation. As a consequence, all data shapes must be provided before any calculation occurs. Before we proceed with defining functions and gradient, let’s check input data shape consistency and provide output shape.

    typedef mxnet::TShape (*UnaryShapeFunction)(const mxnet::TShape& src,
                                         const EnvArguments& env);
    typedef mxnet::TShape (*BinaryShapeFunction)(const mxnet::TShape& lhs,
                                          const mxnet::TShape& rhs,
                                          const EnvArguments& env);

You can use mshadow::TShape to check input data shape and designate output data shape. If you don’t define this function, the default output shape is the same as the input shape. In the case of a binary operator, the shape of lhs and rhs is checked as the same by default.

You can also use shape functions to check if any additional arguments and resources are present. Refer to the additional usages of EnvArguments to accomplish this.

Before we start on our smooth l1 loss example, we define a XPU to cpu or gpu in the header smooth_l1_unary-inl.h implementation so that we reuse the same code in smooth_l1_unary.cc and smooth_l1_unary.cu.

    #include 
    #if defined(__CUDACC__)
    #define XPU gpu
    #else
    #define XPU cpu
    #endif

In our smooth l1 loss example, it’s okay to use the default behavior whereby the output has the same shape as the source. Written explicitly, it is:

    inline mxnet::TShape SmoothL1Shape_(const mxnet::TShape& src,
                                 const EnvArguments& env) {
      return mxnet::TShape(src);
    }

Define Functions¶

Create a unary or binary function with one output: mshadow::TBlob.

    typedef void (*UnaryFunction)(const TBlob& src,
                                  const EnvArguments& env,
                                  TBlob* ret,
                                  OpReqType req,
                                  RunContext ctx);
    typedef void (*BinaryFunction)(const TBlob& lhs,
                                   const TBlob& rhs,
                                   const EnvArguments& env,
                                   TBlob* ret,
                                   OpReqType req,
                                   RunContext ctx);

• Functions are differentiated by the types of input arguments.
• RunContext ctx contains information needed during runtime for execution.

        struct RunContext {
          void *stream;  // the stream of the device, can be NULL or Stream* in GPU mode
          template<typename xpu> inline mshadow::Stream<xpu>* get_stream() // get mshadow stream from Context
        }  // namespace mxnet

mshadow::stream *s = ctx.get_stream(); is an example of obtaining a stream from ctx.

• OpReqType req denotes how computation results are written into ret.

        enum OpReqType {
          kNullOp,  // no operation, do not write anything
          kWriteTo,  // write gradient to provided space
          kWriteInplace,  // perform an in-place write
          kAddTo  // add to the provided space
        };

A macro is defined in operator_util.h for a simplified use of OpReqType. ASSIGN_DISPATCH(out, req, exp) checks req and performs an assignment.

In our smooth l1 loss example, we use UnaryFunction to define the function of this operator.

    template<typename xpu>
    void SmoothL1Forward_(const TBlob& src,
                          const EnvArguments& env,
                          TBlob *ret,
                          OpReqType req,
                          RunContext ctx) {
      using namespace mshadow;
      using namespace mshadow::expr;
      mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
      real_t sigma2 = env.scalar * env.scalar;
      MSHADOW_TYPE_SWITCH(ret->type_flag_, DType, {
        mshadow::Tensor<xpu, 2, DType> out = ret->get<xpu, 2, DType>(s);
        mshadow::Tensor<xpu, 2, DType> in = src.get<xpu, 2, DType>(s);
        ASSIGN_DISPATCH(out, req,
                        F<mshadow_op::smooth_l1_loss>(in, ScalarExp<DType>(sigma2)));
      });
    }

After obtaining mshadow::Stream from RunContext, we get mshadow::Tensor from mshadow::TBlob. mshadow::F is a shortcut to initiate a mshadow expression. The macro MSHADOW_TYPE_SWITCH(type, DType, ...) handles details on different types, and the macro ASSIGN_DISPATCH(out, req, exp) checks OpReqType and performs actions accordingly. sigma2 is a special parameter in this loss, which we will cover later.

Define Gradients (Optional)¶

Create a gradient function with various types of inputs.

    // depending only on out_grad
    typedef void (*UnaryGradFunctionT0)(const OutputGrad& out_grad,
                                        const EnvArguments& env,
                                        TBlob* in_grad,
                                        OpReqType req,
                                        RunContext ctx);
    // depending only on out_value
    typedef void (*UnaryGradFunctionT1)(const OutputGrad& out_grad,
                                        const OutputValue& out_value,
                                        const EnvArguments& env,
                                        TBlob* in_grad,
                                        OpReqType req,
                                         RunContext ctx);
    // depending only on in_data
    typedef void (*UnaryGradFunctionT2)(const OutputGrad& out_grad,
                                        const Input0& in_data0,
                                        const EnvArguments& env,
                                        TBlob* in_grad,
                                        OpReqType req,
                                        RunContext ctx);

Gradient functions of binary operators have similar structures, except that Input, TBlob, and OpReqType are doubled.

GradFunctionArgument

Input0, Input, OutputValue, and OutputGrad all share the structure of GradFunctionArgument, which is defined as:

      struct GradFunctionArgument {
          TBlob data;
      }

In our smooth l1 loss example, note that it’s an f'(x), which utilizes input for the gradient calculation, so the UnaryGradFunctionT2 is suitable. To enable the chain rule of the gradient, we also need to multiply out_grad from the top to the result of in_grad.

    template<typename xpu>
    void SmoothL1BackwardUseIn_(const OutputGrad& out_grad,
                                const Input0& in_data0,
                                const EnvArguments& env,
                                TBlob *in_grad,
                                OpReqType req,
                                RunContext ctx) {
      using namespace mshadow;
      using namespace mshadow::expr;
      mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
      real_t sigma2 = env.scalar * env.scalar;
      MSHADOW_TYPE_SWITCH(in_grad->type_flag_, DType, {
        mshadow::Tensor<xpu, 2, DType> src = in_data0.data.get<xpu, 2, DType>(s);
        mshadow::Tensor<xpu, 2, DType> ograd = out_grad.data.get<xpu, 2, DType>(s);
        mshadow::Tensor<xpu, 2, DType> igrad = in_grad->get<xpu, 2, DType>(s);
         ASSIGN_DISPATCH(igrad, req,
                        ograd * F<mshadow_op::smooth_l1_gradient>(src, ScalarExp<DType>(sigma2)));
      });
    }

Register SimpleOp to MXNet¶

After creating the shape, function, and gradient, restore them into both an NDArray operator and a symbolic operator. To simplify this process, use the registration macro defined in operator_util.h.

    MXNET_REGISTER_SIMPLE_OP(Name, DEV)
    .set_shape_function(Shape)
    .set_function(DEV::kDevMask, Function<XPU>, SimpleOpInplaceOption)
    .set_gradient(DEV::kDevMask, Gradient<XPU>, SimpleOpInplaceOption)
    .describe("description");

SimpleOpInplaceOption is defined as:

    enum SimpleOpInplaceOption {
      kNoInplace,  // do not allow inplace in arguments
      kInplaceInOut,  // allow inplace in with out (unary)
      kInplaceOutIn,  // allow inplace out_grad with in_grad (unary)
      kInplaceLhsOut,  // allow inplace left operand with out (binary)
      kInplaceOutLhs  // allow inplace out_grad with lhs_grad (binary)
    };

In our example, we have a gradient function that relies on input data, so the function can’t be written in place. The output gradient has no purpose after gradient computation, so the gradient can be written in place.

    MXNET_REGISTER_SIMPLE_OP(smooth_l1, XPU)
    .set_function(XPU::kDevMask, SmoothL1Forward_<XPU>, kNoInplace)
    .set_gradient(XPU::kDevMask, SmoothL1BackwardUseIn_<XPU>, kInplaceOutIn)
    .set_enable_scalar(true)
    .describe("Calculate Smooth L1 Loss(lhs, scalar)");

Remember from the discussion of shape functions that a default behavior without set_shape_function forces the inputs (if they’re binary) to be the same shape and yield the same shape for output. We’ll discuss set_enable_scalar later.

NDArray Operator Summary¶

• Create a shape function for determining the output shape.
• Create a function as the forward routine by choosing a suitable function type.
• Create a gradient as the backward routine by choosing a suitable gradient type.
• Register the operator using the registration process.

Using SimpleOp on EnvArguments¶

Some operations might need a scalar as input, such as a gradient scale, a set of keyword arguments controlling behavior, or a temporary space to speed up calculations.EnvArguments provides additional arguments and resources to make calculations more scalable and efficient.

    struct EnvArguments {
      real_t scalar;  // scalar argument, if enabled
      std::vector<std::pair<std::string, std::string> > kwargs;  // keyword arguments
      std::vector<Resource> resource;  // pointer to the resources requested
    };

More registration parameters are required to enable these additional features. To prevent confusion on parameters, scalar and kwargs can’t be present at the same time. To enable scalar, use set_enable_scalar(bool enable_scalar) in registration. Then, in forward functions and gradients, the scalar can be accessed from env.scalar as in the function parameter EnvArguments env.

To enable kwargs, use set_enable_kwargs(bool enable_kwargs) in registration. Then, in forward functions and gradients, additional arguments are contained in env.kwarg, which is defined as std::vector std::string> >. Use the DMLC parameter structure to simplify parsing keyword arguments. For more details, see the guide on parameter structure.

Additional resources like mshadow::Random and temporary memory space can also be requested and accessed from EnvArguments.resource. The registration routine is set_resource_request(ResourceRequest req) or set_resource_request(const std::vector), where mxnet::ResourceRequest is defined as:

    struct ResourceRequest {
      enum Type {  // Resource type, indicating what the pointer type is
        kRandom,  // mshadow::Random object
        kTempSpace  // A dynamic temp space that can be arbitrary size
      };
      Type type;  // type of resources
    };

Registration will request the declared resource requests from mxnet::ResourceManager, and place resources in std::vector resource in EnvArguments. To access resources, use the following:

    auto tmp_space_res = env.resources[0].get_space(some_shape, some_stream);
    auto rand_res = env.resources[0].get_random(some_stream);

For an example, see src/operator/loss_binary_op-inl.h.

In our smooth l1 loss example, a scalar input is needed to mark the turning point of a loss function. Therefore, in the registration process, we use set_enable_scalar(true), and use env.scalar in function and gradient declarations.

Crafting a Tensor Operation¶

Because computation utilizes the mshadow library and we sometimes don’t have functions readily available, we can craft tensor operations in operator implementations. If you define such functions as element-wise, you can implement them as a mxnet::op::mshadow_op. src/operator/mshadow_op.h that contains a lot of mshadow_op, for example. mshadow_op are expression mappers. They deal with the scalar case of desired functions. For details, see mshadow expression API guide.

If an operation can’t be done in an element-wise way, like the softmax loss and gradient, then you need to create a new tensor operation. You need to create as mshadow function and as mshadow::cuda function directly. For details, see the mshadow library. For an example, see src/operator/roi_pooling.cc.

In our smooth l1 loss example, we create two mappers, namely the scalar cases of smooth l1 loss and gradient.

    namespace mshadow_op {
    struct smooth_l1_loss {
      // a is x, b is sigma2
      MSHADOW_XINLINE static real_t Map(real_t a, real_t b) {
        if (a > 1.0f / b) {
          return a - 0.5f / b;
        } else if (a < -1.0f / b) {
          return -a - 0.5f / b;
        } else {
          return 0.5f * a * a * b;
        }
      }
    };
    }

The gradient, which can be found in src/operator/smooth_l1_unary-inl.h, is similar.

Beyond Two Operands¶

The new unified API is designed to fulfill the fundamentals of an operation. For operators with more than two inputs, more than one output, or that need more features, see the original Operator API.