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_