operator_util.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef MXNET_OPERATOR_UTIL_H_
#define MXNET_OPERATOR_UTIL_H_
 
#ifdef _MSC_VER
#pragma warning(disable : 4503)  // disable warning: decorated name length exceeded.
#endif
 
#include <dmlc/registry.h>
#include <dmlc/parameter.h>
#include <map>
#include <vector>
#include <string>
#include <utility>
#include "./base.h"
#include "./operator.h"
 
#if DMLC_USE_CXX11
#include <functional>
#endif
 
namespace mxnet {
namespace op {
struct GradFunctionArgument {
  TBlob data;
};
 
struct Input0 : GradFunctionArgument {};
struct Input1 : GradFunctionArgument {};
 
struct OutputValue : GradFunctionArgument {};
struct OutputGrad : GradFunctionArgument {};
 
struct EnvArguments {
  real_t scalar;
  std::vector<std::pair<std::string, std::string> > kwargs;
  std::vector<Resource> resource;
};
 
typedef void (*SourceFunction)(const EnvArguments& env, TBlob* ret, OpReqType req, RunContext ctx);
 
typedef mxnet::TShape (*SourceShapeFunction)(const EnvArguments& env);
 
typedef void (*UnaryFunction)(const TBlob& src,
                              const EnvArguments& env,
                              TBlob* ret,
                              OpReqType req,
                              RunContext ctx);
typedef mxnet::TShape (*UnaryShapeFunction)(const mxnet::TShape& src, const EnvArguments& env);
 
typedef void (*UnaryGradFunctionT0)(const OutputGrad& out_grad,
                                    const EnvArguments& env,
                                    TBlob* in_grad,
                                    OpReqType req,
                                    RunContext ctx);
typedef void (*UnaryGradFunctionT1)(const OutputGrad& out_grad,
                                    const OutputValue& out_value,
                                    const EnvArguments& env,
                                    TBlob* in_grad,
                                    OpReqType req,
                                    RunContext ctx);
typedef void (*UnaryGradFunctionT2)(const OutputGrad& out_grad,
                                    const Input0& in_data0,
                                    const EnvArguments& env,
                                    TBlob* in_grad,
                                    OpReqType req,
                                    RunContext ctx);
typedef void (*BinaryFunction)(const TBlob& lhs,
                               const TBlob& rhs,
                               const EnvArguments& env,
                               TBlob* ret,
                               OpReqType req,
                               RunContext ctx);
 
typedef mxnet::TShape (*BinaryShapeFunction)(const mxnet::TShape& lhs,
                                             const mxnet::TShape& rhs,
                                             const EnvArguments& env);
typedef void (*BinaryGradFunctionT0)(const OutputGrad& out_grad,
                                     const EnvArguments& env,
                                     TBlob* lhs_grad,
                                     TBlob* rhs_grad,
                                     OpReqType req_lhs_grad,
                                     OpReqType req_rhs_grad,
                                     RunContext ctx);
typedef void (*BinaryGradFunctionT1)(const OutputGrad& out_grad,
                                     const Input0& lhs,
                                     const Input1& rhs,
                                     const EnvArguments& env,
                                     TBlob* lhs_grad,
                                     TBlob* rhs_grad,
                                     OpReqType req_lhs_grad,
                                     OpReqType req_rhs_grad,
                                     RunContext ctx);
 
enum SimpleOpInplaceOption {
  kNoInplace,
  kInplaceInOut,
  kInplaceOutIn,
  kInplaceLhsOut,
  kInplaceOutLhs
};
 
enum SimpleOpScalarOption { kScalarBeforeArray, kArrayBeforeScalar };
 
enum SimpleOpRegOption { kNotRegisterSymbolic, kRegisterSymbolic };
 
class SimpleOpRegEntry {
 public:
  typedef SimpleOpRegEntry TSelf;
  std::string name;
  virtual TSelf& set_symbol_op_name(char const* symbol_name) = 0;
  virtual TSelf& set_enable_scalar(bool enable_scalar,
                                   SimpleOpScalarOption type_mask = kArrayBeforeScalar) = 0;
  virtual TSelf& set_enable_kwargs(bool enable_kwargs) = 0;
  virtual TSelf& set_resource_request(const std::vector<ResourceRequest>& reqs) = 0;
  virtual TSelf& set_resource_request(ResourceRequest req) = 0;
  virtual TSelf& set_shape_function(SourceShapeFunction fshapeinfer) = 0;
  virtual TSelf& set_shape_function(UnaryShapeFunction fshapeinfer) = 0;
  virtual TSelf& set_shape_function(BinaryShapeFunction fshapeinfer) = 0;
  virtual TSelf& set_function(int dev_mask,
                              SourceFunction fsource,
                              SimpleOpRegOption register_symbolic = kRegisterSymbolic) = 0;
  virtual TSelf& set_function(int dev_mask,
                              UnaryFunction funary,
                              SimpleOpInplaceOption inplace_in_out,
                              SimpleOpRegOption register_symbolic = kRegisterSymbolic) = 0;
  virtual TSelf& set_function(int dev_mask,
                              BinaryFunction fbinary,
                              SimpleOpInplaceOption inplace_lhs_out,
                              SimpleOpRegOption register_symbolic = kRegisterSymbolic) = 0;
  virtual TSelf& set_gradient(int dev_mask,
                              UnaryGradFunctionT0 fgrad,
                              SimpleOpInplaceOption inplace_out_in_grad) = 0;
  virtual TSelf& set_gradient(int dev_mask,
                              UnaryGradFunctionT1 fgrad,
                              SimpleOpInplaceOption inplace_out_in_grad) = 0;
  virtual TSelf& set_gradient(int dev_mask,
                              UnaryGradFunctionT2 fgrad,
                              SimpleOpInplaceOption inplace_out_in_grad) = 0;
  virtual TSelf& set_gradient(int dev_mask,
                              BinaryGradFunctionT0 fgrad,
                              SimpleOpInplaceOption inplace_out_lhs_grad) = 0;
  virtual TSelf& set_gradient(int dev_mask,
                              BinaryGradFunctionT1 fgrad,
                              SimpleOpInplaceOption inplace_out_lhs_grad) = 0;
  virtual TSelf& describe(const std::string& description) = 0;
  virtual TSelf& add_arguments(const std::vector<dmlc::ParamFieldInfo>& args) = 0;
  virtual ~SimpleOpRegEntry() {}
};
 
class SimpleOpRegistry {
 public:
  SimpleOpRegEntry& __REGISTER_OR_FIND__(char const* name);
  inline static const SimpleOpRegEntry* Find(const std::string& name) {
    return Get()->fmap_.at(name);
  }
  static SimpleOpRegistry* Get();
 
 private:
  // destructor
  ~SimpleOpRegistry();
  std::map<std::string, SimpleOpRegEntry*> fmap_;
};
 
#define ASSIGN_DISPATCH(out, req, exp) \
  {                                    \
    switch (req) {                     \
      case kNullOp:                    \
        break;                         \
      case kWriteTo:                   \
      case kWriteInplace:              \
        (out) = (exp);                 \
        break;                         \
      case kAddTo:                     \
        (out) += (exp);                \
        break;                         \
      default:                         \
        LOG(FATAL) << "not reached";   \
    }                                  \
  }
 
#define MXNET_SPECIAL_MAX_NDIM 5
 
//--------------------------------------------------------------
// The following part are API Registration of Simple Operators
//--------------------------------------------------------------
#define MXNET_REGISTER_SIMPLE_OP(Name, DEV)                                               \
  static ::mxnet::op::SimpleOpRegEntry& __make_##SimpleOpRegEntry##_##Name##__##DEV##__ = \
      ::mxnet::op::SimpleOpRegistry::Get()->__REGISTER_OR_FIND__(#Name)
 
}  // namespace op
}  // namespace mxnet
#endif  // MXNET_OPERATOR_UTIL_H_