Apache MXNet - Unified Operator API

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

This chapter delves into the SimpleOp API in Apache MXNet, a C++ interface for creating new low-level operators. This is an advanced topic primarily for developers who need to extend MXNet’s core functionality with custom computations not available out-of-the-box. This might be necessary for performance-critical sections, implementing novel algorithms, or integrating specialized hardware accelerators. SimpleOp allows fine-grained control over operator behavior, including shape inference, forward/backward computation logic, and resource management within the MXNet engine.

SimpleOp Overview

SimpleOp is designed to streamline the creation of operators, particularly for unary (one input tensor) and binary (two input tensors) operations, which form the bulk of mathematical computations in deep learning models. By defining these fundamental operator elements in C++, they can be efficiently integrated into MXNet’s execution engine, which then incorporates them into computation graphs and applies various optimizations.

We will illustrate the SimpleOp mechanism by conceptualizing the creation of a custom operator for ‘Smooth L1 Loss’. Smooth L1 loss is a robust loss function often used in object detection for bounding box regression. It behaves like L2 loss for small errors and L1 loss for large errors. For an input x (representing, for example, prediction - target) and a scalar sigma, it can be defined as:

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

The derivative of this loss with respect to x, let’s call it loss'(x, sigma), is:

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

While the full application in a neural network involves upstream gradients and potentially multiple inputs (predictions, targets), SimpleOp can be used to implement the core loss(x, sigma) and loss'(x, sigma) computations if x is treated as the direct input tensor to this custom operator, and sigma as a parameter.

Defining Shape Inference Functions

MXNet’s execution engine requires all tensor shapes to be known before any computation occurs. Therefore, a custom operator must provide a shape function. This function infers the output tensor shape(s) based on the input tensor shape(s) and any operator-specific parameters.

The typical C++ type signatures for shape functions in the context of SimpleOp might look like this:

#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 represents the dimensions of an NDArray. If no custom shape function is provided, or if the default behavior is sufficient (e.g., for element-wise unary operators where output shape equals input shape), MXNet might use a default inference. However, explicit definition is best practice for clarity and custom logic.

For our Smooth L1 loss example, if it’s an element-wise operation, the output shape will be the same as the input shape (x). A simple shape function would assert this:

// 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
}

Defining Forward Computation Functions

The forward function performs the core computation of the operator. It takes input tensors and produces output tensors.

Conceptual C++ type signatures for forward functions:

#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) is a container for a tensor’s data pointer, shape, and type information. mxnet::RunContext provides runtime information, such as the CUDA stream for GPU operations or CPU threading resources.

mxnet::OpReqType specifies how the result should be written to the output 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'
};

For the Smooth L1 loss forward pass, the function would use MShadow (MXNet’s tensor math library) for efficient, device-agnostic computation:

#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)));
  });
}

Defining Gradient (Backward) Computation Functions

Gradient functions compute the derivative of the operator’s output(s) with respect to its input(s), multiplied by the incoming gradient from the subsequent layer (out_grad). This is essential for backpropagation.

Conceptual C++ type signatures for gradient functions. They often depend on out_grad, and optionally on the forward pass’s inputs (in_data_N) or outputs (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);

For Smooth L1 loss, the backward function computes out_grad * loss'(in_data, sigma). in_data is the x from the forward pass.

// 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)));
  });
}

Registering the SimpleOp with MXNet

After defining the shape, forward, and gradient functions, the operator must be registered with MXNet’s runtime. This makes it discoverable by the symbolic graph system (NNVM/TVM based) and the imperative NDArray API.

Registration is typically done using C++ macros provided by MXNet, such as MXNET_OPERATOR_REGISTER_UNARY, MXNET_OPERATOR_REGISTER_BINARY, or more specific ones if the operator needs to declare parameters (DMLC_REGISTER_PARAMETER).

A conceptual registration structure in a .cc (CPU) or .cu (GPU) file:

// 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

The SimpleOpInplaceOption enum or similar flags control whether inputs/outputs can share memory for efficiency. This depends on the operator’s logic (e.g., if an input is modified or needed unmodified by the backward pass).

Using EnvArguments for Scalars and Resources

EnvArguments (or a similar mechanism like operator parameters/attributes) is used to pass additional information to operator functions beyond the main tensor inputs. This can include:

• Scalar values: Like sigma in our Smooth L1 loss example.
• Keyword arguments: Strings or other types controlling operator behavior (less common for SimpleOp, more for Python-level ops).
• Requested Resources: Such as pre-allocated temporary workspace memory or random number generators, managed by MXNet’s ResourceManager.

The EnvArguments struct (customized per operator needs) might be defined as:

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...
};

If an operator needs temporary memory, it declares a ResourceRequest during registration. MXNet’s ResourceManager then provides these resources via EnvArguments (or a similar context struct) during execution.

Accessing requested resources within an operator function:

// 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*/);

In our Smooth L1 loss, env.scalar_sigma would be populated from an operator attribute or a scalar input. The registration process must enable and correctly map this scalar argument.

Building Tensor Operations with MShadow

Custom C++ operators in MXNet heavily rely on MShadow, MXNet’s template-based tensor math library, for efficient and device-agnostic (CPU/GPU) computation. You may need to define custom MShadow “mappers” or “functors” when:

• The desired computation is not readily available as a standard MShadow expression or operator.
• The operation is complex, non-element-wise (e.g., softmax, reductions with specific logic), or requires specialized low-level control.

Example: MShadow mappers for the forward and backward Smooth L1 loss calculations. These are 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

These MShadow Map functions would then be invoked within the C++ SmoothL1Forward and SmoothL1Backward functions using MShadow’s expression system, e.g., mshadow::expr::F<smooth_l1_loss_forward_map>(input_tensor_expr, scalar_sigma_expr).