Skip to content
W3S101

Apache MXNet - 统一算子 API

Apache MXNet - Unified Operator API (SimpleOp for Custom C++ Operators)

Section titled “Apache MXNet - Unified Operator API (SimpleOp for Custom C++ Operators)”

本章深入探讨 Apache MXNet 中的 SimpleOp API，这是一个用于创建新的低级 operator 的 C++ 接口。这是一个高级主题，主要面向需要通过定制的计算（开箱即用功能中未提供）来扩展 MXNet 核心功能的开发者。这对于性能关键的部分、实现 novel algorithms 或集成 specialized hardware accelerators 可能很有必要。SimpleOp 允许对 operator 行为进行细粒度控制，包括 shape inference、forward/backward computation logic 和 MXNet engine 中的 resource management。

SimpleOp 概述

Section titled “SimpleOp 概述”

SimpleOp 旨在简化 operator 的创建，特别是对于 unary (one input tensor) 和 binary (two input tensors) 操作，它们是深度学习模型中大部分数学计算的基础。通过在 C++ 中定义这些基本的 operator 元素，它们可以高效地集成到 MXNet 的 execution engine 中，然后由 engine 将它们纳入 computation graphs 并应用各种 optimizations。

我们将通过概念化创建一个用于 ‘Smooth L1 Loss’ 的 custom operator 来演示 SimpleOp 机制。Smooth L1 loss 是一种 robust loss function，常用于 object detection 中的 bounding box regression。对于小误差，它表现得像 L2 loss；对于大误差，它表现得像 L1 loss。对于输入 x（例如代表 prediction - target）和 scalar sigma，其定义如下：

loss(x, sigma) = 0.5 * (sigma*x)^2      if |x| < 1/(sigma^2)
                 |x| - 0.5/(sigma^2)    otherwise

此 loss 关于 x 的 derivative，我们称之为 loss'(x, sigma)，定义如下：

loss'(x, sigma) = sigma^2 * x          if |x| < 1/(sigma^2)
                  sign(x)              otherwise

虽然在神经网络中的完整应用涉及 upstream gradients 和潜在的多个输入（predictions, targets），但如果 x 被视为此 custom operator 的直接输入 tensor，而 sigma 被视为 parameter，则可以使用 SimpleOp 实现核心的 loss(x, sigma) 和 loss'(x, sigma) 计算。

定义 Shape Inference 函数

Section titled “定义 Shape Inference 函数”

MXNet 的 execution engine 要求在进行任何计算之前已知所有 tensor 的 shapes。因此，custom operator 必须提供一个 shape function。此函数根据输入 tensor 的 shapes 和任何 operator-specific parameters 来推断输出 tensor 的 shape(s)。

在 SimpleOp 上下文中，shape functions 的典型 C++ type signatures 可能如下所示：

#include <mxnet/base.h> // For mxnet::TShape
#include <vector>      // For std::vector


// Forward declaration of EnvArguments if used by shape function
// struct EnvArguments;


// Example for a unary operator shape function
typedef bool (*UnaryShapeInferFn)(const mxnet::TShape& src_shape,
                                  mxnet::TShape* out_shape,
                                  const std::vector<mxnet::resource::ResourceRequest>& reqs /* or Env */);


// Example for a binary operator shape function
typedef bool (*BinaryShapeInferFn)(const mxnet::TShape& lhs_shape,
                                   const mxnet::TShape& rhs_shape,
                                   mxnet::TShape* out_shape,
                                   const std::vector<mxnet::resource::ResourceRequest>& reqs /* or Env */);

mxnet::TShape 表示 NDArray 的 dimensions。如果没有提供 custom shape function，或者默认行为足够（例如，对于输出 shape 等于输入 shape 的 element-wise unary operator），MXNet 可能会使用默认 inference。然而，为了清晰和自定义逻辑，明确定义是 best practice。

对于我们的 Smooth L1 loss 示例，如果它是 element-wise 操作，输出 shape 将与输入 shape (x) 相同。一个简单的 shape function 将断言这一点：

// This C++ snippet would typically be in a header or an inline implementation file (-inl.h)
// to be included by both CPU (.cc) and GPU (.cu) source files.


// Placeholder for EnvArguments if it carries shape-relevant info
struct EnvArguments { /* ... real_t scalar_sigma, etc. ... */ };


inline bool SmoothL1Shape(const mxnet::TShape& src,
                          mxnet::TShape* out_shape,
                          const EnvArguments& env) {
  // For SmoothL1, output shape is the same as input source shape
  *out_shape = src;
  return true; // Indicates successful shape inference
}

定义 Forward Computation 函数

Section titled “定义 Forward Computation 函数”

Forward function 执行 operator 的核心计算。它接收输入 tensors 并产生输出 tensors。

Forward functions 的概念性 C++ type signatures：

#include <mxnet/blob.h> // For mxnet::TBlob
#include <mxnet/operator_common.h> // For OpReqType, RunContext


// For unary operators
typedef void (*UnaryForwardFn)(const mxnet::TBlob& src, // Input data tensor
                               const EnvArguments& env, // Environment/scalar arguments
                               mxnet::TBlob* ret,       // Output data tensor
                               mxnet::OpReqType req,    // Request type (e.g., write, add)
                               const mxnet::RunContext& ctx);  // Runtime context (stream, etc.)


// For binary operators (similar, with lhs and rhs TBlobs)
typedef void (*BinaryForwardFn)(const mxnet::TBlob& lhs,
                                const mxnet::TBlob& rhs,
                                const EnvArguments& env,
                                mxnet::TBlob* ret,
                                mxnet::OpReqType req,
                                const mxnet::RunContext& ctx);

mxnet::TBlob (Tensor Blob) 是一个容器，用于存储 tensor 的 data pointer、shape 和 type information。mxnet::RunContext 提供运行时信息，例如 GPU 操作的 CUDA stream 或 CPU threading resources。

mxnet::OpReqType 指定如何将结果写入输出 ret tensor：

enum OpReqType {
  kNullOp,       // Do nothing (e.g., for gradients not needed or not computed)
  kWriteTo,      // Write result directly to 'ret' (overwriting existing content)
  kWriteInplace, // Write result in-place (if 'ret' shares memory with an input)
  kAddTo         // Add computed result to existing values in 'ret'
};

对于 Smooth L1 loss forward pass，该函数将使用 MShadow（MXNet 的 tensor math library）进行高效、device-agnostic 的计算：

#include <mshadow/tensor.h> // For mshadow::Tensor, mshadow::Stream, etc.


// Assume EnvArguments contains sigma, e.g., env.scalar_sigma
// Assume smooth_l1_loss_forward_map is an MShadow functor (see 'Building Tensor Operations with MShadow')


template<typename xpu> // xpu can be mshadow::cpu or mshadow::gpu
void SmoothL1Forward(const mxnet::TBlob& src_blob, const EnvArguments& env,
                     mxnet::TBlob* ret_blob, mxnet::OpReqType req,
                     const mxnet::RunContext& run_ctx) {
  using namespace mshadow;
  Stream<xpu>* s = run_ctx.get_stream<xpu>(); // Get the computation stream


  // Assuming env.scalar_sigma holds the sigma value for the loss
  mxnet::real_t sigma = env.scalar_sigma;


  // MSHADOW_TYPE_SWITCH dispatches based on tensor data type (float, double, etc.)
  MSHADOW_TYPE_SWITCH(ret_blob->type_flag_, DType, {
    // Get MShadow Tensor views of the TBlobs
    Tensor<xpu, 2, DType> out = ret_blob->get_with_shape<xpu, 2, DType>(ret_blob->shape_.get<2>(), s);
    Tensor<xpu, 2, DType> in_data = src_blob.get_with_shape<xpu, 2, DType>(src_blob.shape_.get<2>(), s);


    // Perform the computation using MShadow expressions and the custom functor
    // 'smooth_l1_loss_forward_map' would implement the loss(x, sigma) formula.
    // ASSIGN_DISPATCH handles OpReqType (kWriteTo, kAddTo, etc.)
    ASSIGN_DISPATCH(out, req, smooth_l1_loss_forward_map(in_data, ScalarExp<DType>(sigma)));
  });
}

定义 Gradient (Backward) Computation 函数

Section titled “定义 Gradient (Backward) Computation 函数”

Gradient functions 计算 operator 输出关于其输入的 derivative，并乘以来自后续层 (out_grad) 的传入 gradient。这对于 backpropagation 至关重要。

Gradient functions 的概念性 C++ type signatures。它们通常依赖于 out_grad，并可选地依赖于 forward pass 的输入 (in_data_N) 或输出 (out_value_N)。

// For a unary operator, computing gradient w.r.t. its single input
// This example assumes gradient depends on out_grad and original input_data
typedef void (*UnaryBackwardFnUseInput)(const mxnet::TBlob& out_grad_blob, // Gradient from layer above
                                        const mxnet::TBlob& in_data_blob,  // Original input to fwd pass
                                        const EnvArguments& env,
                                        mxnet::TBlob* in_grad_blob, // Gradient to compute (w.r.t. input)
                                        mxnet::OpReqType req,
                                        const mxnet::RunContext& ctx);

对于 Smooth L1 loss，backward function 计算 out_grad * loss'(in_data, sigma)。in_data 是 forward pass 中的 x。

// Assume smooth_l1_loss_backward_map is an MShadow functor for loss'(x, sigma)


template<typename xpu>
void SmoothL1Backward(const mxnet::TBlob& ograd_blob, const mxnet::TBlob& data_blob,
                      const EnvArguments& env, mxnet::TBlob* igrad_blob,
                      mxnet::OpReqType req, const mxnet::RunContext& run_ctx) {
  using namespace mshadow;
  Stream<xpu>* s = run_ctx.get_stream<xpu>();
  mxnet::real_t sigma = env.scalar_sigma;


  MSHADOW_TYPE_SWITCH(igrad_blob->type_flag_, DType, {
    Tensor<xpu, 2, DType> igrad = igrad_blob->get_with_shape<xpu, 2, DType>(igrad_blob->shape_.get<2>(), s);
    Tensor<xpu, 2, DType> ograd = ograd_blob.get_with_shape<xpu, 2, DType>(ograd_blob.shape_.get<2>(), s);
    Tensor<xpu, 2, DType> data = data_blob.get_with_shape<xpu, 2, DType>(data_blob.shape_.get<2>(), s);


    # Chain rule: input_gradient = output_gradient * local_derivative
    ASSIGN_DISPATCH(igrad, req, ograd * smooth_l1_loss_backward_map(data, ScalarExp<DType>(sigma)));
  });
}

在 MXNet 中注册 SimpleOp

Section titled “在 MXNet 中注册 SimpleOp”

在定义了 shape、forward 和 gradient functions 后，必须将 operator 注册到 MXNet 的 runtime 中。这使得符号图系统（基于 NNVM/TVM）和命令式 NDArray API 都能发现它。

注册通常使用 MXNet 提供的 C++ macros 完成，例如 MXNET_OPERATOR_REGISTER_UNARY, MXNET_OPERATOR_REGISTER_BINARY，或如果 operator 需要声明参数 (DMLC_REGISTER_PARAMETER)，则使用更具体的宏。

在 .cc (CPU) 或 .cu (GPU) 文件中的概念性注册结构：

// In a .cc file (e.g., smooth_l1_op.cc) or .cu file (e.g., smooth_l1_op.cu)
#include <mxnet/operator_util.h> // For registration macros
// ... includes for shape, forward, backward functions, and EnvArguments struct ...


namespace mxnet {
namespace op {


// Define a parameter struct if the operator has configurable attributes beyond EnvArguments::scalar
// struct SmoothL1Param : public dmlc::Parameter<SmoothL1Param> {
//   float sigma_param_attr; // Example attribute if not using EnvArguments::scalar
//   DMLC_DECLARE_PARAMETER(SmoothL1Param) { /* ... field declarations ... */ }
// };


// Actual registration call (macros and structure may vary with MXNet version)
// This is a simplified, illustrative example for a unary operator
// The real macros handle CPU/GPU dispatch, NDArray/Symbol API, etc.


// This registration sets up the operator property, linking functions.
// It needs to be done for each device (CPU, GPU) the operator supports.
// bool SmoothL1OpProp::InferShape(...) { /* calls SmoothL1Shape */ }
// void SmoothL1OpProp::Forward(...) { /* calls SmoothL1Forward */ }
// void SmoothL1OpProp::Backward(...) { /* calls SmoothL1Backward */ }


// Then register this property with a name:
// MXNET_REGISTER_OP_PROPERTY(SmoothL1, SmoothL1OpProp)
// .describe("Computes Smooth L1 loss element-wise.")
// .add_argument("data", "NDArray-or-Symbol", "Input data tensor.")
// .add_argument("scalar_sigma", "float", "Sigma parameter for Smooth L1 loss (passed via EnvArguments).");


// The registration process is complex. Refer to existing operator implementations
// in the MXNet source tree (e.g., src/operator/) for current best practices.
// Key aspects include specifying input/output arguments, attributes, and linking
// the C++ implementation functions for different execution contexts and data types.


} // namespace op
} // namespace mxnet

SimpleOpInplaceOption enum 或类似的 flags 控制输入/输出是否可以共享内存以提高效率。这取决于 operator 的逻辑（例如，输入是否在 backward pass 中被修改或需要未修改）。

使用 EnvArguments 传递 Scalars 和 Resources

Section titled “使用 EnvArguments 传递 Scalars 和 Resources”

EnvArguments（或类似的机制，如 operator parameters/attributes）用于向 operator functions 传递主 tensor 输入之外的附加信息。这可以包括：

  • Scalar values： 例如我们 Smooth L1 loss 示例中的 sigma。
  • Keyword arguments： 控制 operator 行为的字符串或其他类型（对于 SimpleOp 不太常见，对于 Python-level ops 更常见）。
  • Requested Resources： 例如预分配的临时 workspace memory 或 random number generators，由 MXNet 的 ResourceManager 管理。

EnvArguments struct（根据 operator 需要定制）可以定义为：

struct EnvArguments {
  mxnet::real_t scalar_sigma; // For the sigma in Smooth L1 loss
  // std::vector<mxnet::resource::Resource> resource_handles; // Handles to requested resources
  // Other members as needed...
};

如果 operator 需要临时内存，它会在注册期间声明一个 ResourceRequest。然后，MXNet 的 ResourceManager 会在执行期间通过 EnvArguments（或类似的 context struct）提供这些 resources。

在 operator function 中访问请求的 resources：

// Example pseudocode for accessing a resource within an operator function:
// mxnet::resource::Resource temp_space_res = env.resource_handles[/*index of temp space request*/];
// mshadow::TensorContainer<xpu, DType> temp_workspace =
//   temp_space_res.get_space_typed<xpu, DType>(/*shape_needed*/, s /*stream*/);

在我们的 Smooth L1 loss 中，env.scalar_sigma 将从 operator attribute 或 scalar 输入填充。注册过程必须启用并正确映射此 scalar argument。

使用 MShadow 构建 Tensor 操作

Section titled “使用 MShadow 构建 Tensor 操作”

MXNet 中的 custom C++ operators 大量依赖于 MShadow，它是 MXNet 基于 template 的 tensor math library，用于高效且 device-agnostic (CPU/GPU) 的计算。当出现以下情况时，您可能需要定义 custom MShadow “mappers” 或 “functors”：

  • 所需的计算不是 MShadow 标准 expression 或 operator 中现成的。
  • 操作复杂、非 element-wise（例如，带有特定逻辑的 softmax、reductions）或需要 specialized low-level control。

示例：用于 forward 和 backward Smooth L1 loss 计算的 MShadow mappers。这些是 element-wise operations。

#include <mshadow/tensor_base.h> // For MSHADOW_XINLINE, mshadow::real_t, etc.
#include <cmath> // For std::abs


namespace mshadow_op {


// Functor for Smooth L1 Loss forward pass: loss(x, sigma)
struct smooth_l1_loss_forward_map {
  // x_val is element from input tensor, sigma_val is the scalar sigma parameter
  MSHADOW_XINLINE static mshadow::real_t Map(mshadow::real_t x_val, mshadow::real_t sigma_val) {
    if (sigma_val == 0) return 0; // Or handle as error, depending on desired behavior
    mshadow::real_t sigma2_inv = 1.0f / (sigma_val * sigma_val);
    mshadow::real_t abs_x = std::abs(x_val);
    if (abs_x < sigma2_inv) {
      return 0.5f * x_val * x_val * (sigma_val * sigma_val);
    } else {
      return abs_x - 0.5f * sigma2_inv;
    }
  }
};


// Functor for Smooth L1 Loss backward pass: loss'(x, sigma)
struct smooth_l1_loss_backward_map {
  // x_val is element from original input tensor, sigma_val is the scalar sigma parameter
  MSHADOW_XINLINE static mshadow::real_t Map(mshadow::real_t x_val, mshadow::real_t sigma_val) {
    if (sigma_val == 0) return 0;
    mshadow::real_t sigma2_inv = 1.0f / (sigma_val * sigma_val);
    mshadow::real_t abs_x = std::abs(x_val);
    if (abs_x < sigma2_inv) {
      return x_val * (sigma_val * sigma_val);
    } else {
      return (x_val > 0.0f) ? 1.0f : ((x_val < 0.0f) ? -1.0f : 0.0f); // sign(x_val)
    }
  }
};


} // namespace mshadow_op

这些 MShadow Map functions 随后将通过 MShadow 的 expression system 在 C++ 的 SmoothL1Forward 和 SmoothL1Backward 函数中被调用，例如 mshadow::expr::F<smooth_l1_loss_forward_map>(input_tensor_expr, scalar_sigma_expr)。