operator.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_H_
#define MXNET_OPERATOR_H_
 
#include <dmlc/base.h>
#include <dmlc/json.h>
#include <dmlc/logging.h>
#include <dmlc/registry.h>
#include <nnvm/node.h>
#include <vector>
#include <map>
#include <string>
#include <utility>
#include "./base.h"
#include "./resource.h"
#include "./op_attr_types.h"
 
namespace mxnet {
class Operator {
 public:
  virtual ~Operator() {}
  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;
  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) {
    LOG(FATAL) << "Backward is not implemented";
  }
  virtual ExecType exec_type()  // NOLINT(*) exec_type has been moved to OperatorProperty
      const final {             // NOLINT(*) exec_type has been moved to OperatorProperty
    return ExecType::kSync;
  }
};
 
#if DMLC_USE_CXX11
// OperatorProperty allows C++11, while Operator do not rely on it.
class OperatorProperty {
 public:
  virtual ~OperatorProperty() {}
  virtual void Init(const std::vector<std::pair<std::string, std::string> >& kwargs) = 0;
  virtual std::map<std::string, std::string> GetParams() const = 0;
  virtual std::vector<std::string> ListArguments() const {
    return {"data"};
  }
  virtual std::vector<std::string> ListOutputs() const {
    return {"output"};
  }
  virtual std::vector<std::string> ListAuxiliaryStates() const {
    return {};
  }
  virtual int NumOutputs() const {
    return this->ListOutputs().size();
  }
  virtual int NumVisibleOutputs() const {
    return NumOutputs();
  }
  virtual bool InferShape(mxnet::ShapeVector* in_shape,
                          mxnet::ShapeVector* out_shape,
                          mxnet::ShapeVector* aux_shape) const = 0;
  virtual bool InferType(std::vector<int>* in_type,
                         std::vector<int>* out_type,
                         std::vector<int>* aux_type) const {
    CHECK_LE(in_type->size(), this->ListArguments().size());
    int n_in = this->ListArguments().size();
    for (unsigned i = 0; i < in_type->size(); ++i) {
      CHECK(in_type->at(i) == mshadow::default_type_flag || in_type->at(i) == -1)
          << "Unsupported data type " << in_type->at(i);
    }
    in_type->clear();
    for (int i = 0; i < n_in; ++i)
      in_type->push_back(mshadow::default_type_flag);
 
    int n_out = this->ListOutputs().size();
    out_type->clear();
    for (int i = 0; i < n_out; ++i)
      out_type->push_back(mshadow::default_type_flag);
 
    int n_aux = this->ListAuxiliaryStates().size();
    aux_type->clear();
    for (int i = 0; i < n_aux; ++i)
      aux_type->push_back(mshadow::default_type_flag);
    return true;
  }
  virtual OperatorProperty* Copy() const = 0;
  virtual Operator* CreateOperator(Context ctx) const = 0;
  virtual Operator* CreateOperatorEx(Context ctx,
                                     mxnet::ShapeVector* in_shape,
                                     std::vector<int>* in_type) const {
    std::vector<int> out_type, aux_type;
    mxnet::ShapeVector out_shape, aux_shape;
    out_type.resize(this->ListOutputs().size());
    out_shape.resize(this->ListOutputs().size());
    aux_type.resize(this->ListAuxiliaryStates().size());
    aux_shape.resize(this->ListAuxiliaryStates().size());
    CHECK(InferType(in_type, &out_type, &aux_type));
    CHECK(InferShape(in_shape, &out_shape, &aux_shape));
    return CreateOperator(ctx);
  }
  virtual std::string TypeString() const = 0;
  //--------------------------------------------------------
  // All the below functions are optional to override.
  //--------------------------------------------------------
  virtual std::vector<ResourceRequest> ForwardResource(const mxnet::ShapeVector& in_shape) const {
    return std::vector<ResourceRequest>();
  }
  virtual std::vector<ResourceRequest> BackwardResource(const mxnet::ShapeVector& in_shape) const {
    return std::vector<ResourceRequest>();
  }
  virtual std::vector<int> DeclareBackwardDependency(const std::vector<int>& out_grad,
                                                     const std::vector<int>& in_data,
                                                     const std::vector<int>& out_data) const {
    // By default requires to see all the things.
    // remember to override this function to get a better performance.
    std::vector<int> ret = out_grad;
    ret.insert(ret.end(), in_data.begin(), in_data.end());
    ret.insert(ret.end(), out_data.begin(), out_data.end());
    return ret;
  }
  virtual std::vector<std::pair<int, void*> > ForwardInplaceOption(
      const std::vector<int>& in_data,
      const std::vector<void*>& out_data) const {
    return std::vector<std::pair<int, void*> >();
  }
  virtual std::vector<std::pair<int, void*> > 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 std::vector<std::pair<int, void*> >();
  }
  template <typename T>
  inline std::vector<T> BackwardInputs(const std::vector<T>& out_grad,
                                       const std::vector<T>& in_data,
                                       const std::vector<T>& out_data) const {
    int counter = 0;
    std::vector<int> out_grad_index(out_grad.size());
    std::vector<int> in_data_index(in_data.size());
    std::vector<int> out_data_index(out_data.size());
    for (size_t i = 0; i < out_grad_index.size(); ++i) {
      out_grad_index[i] = counter++;
    }
    for (size_t i = 0; i < in_data_index.size(); ++i) {
      in_data_index[i] = counter++;
    }
    for (size_t i = 0; i < out_data_index.size(); ++i) {
      out_data_index[i] = counter++;
    }
    std::vector<T> all_data;
    all_data.insert(all_data.end(), out_grad.begin(), out_grad.end());
    all_data.insert(all_data.end(), in_data.begin(), in_data.end());
    all_data.insert(all_data.end(), out_data.begin(), out_data.end());
 
    std::vector<int> ret_index =
        this->DeclareBackwardDependency(out_grad_index, in_data_index, out_data_index);
 
    std::vector<T> ret(ret_index.size());
    for (size_t i = 0; i < ret_index.size(); ++i) {
      ret[i] = all_data[ret_index[i]];
    }
    return ret;
  }
  static OperatorProperty* Create(const char* type_name);
  virtual ExecType exec_type() const {
    return ExecType::kSync;
  }
};
 
typedef std::function<OperatorProperty*()> OperatorPropertyFactory;
struct OperatorPropertyReg
    : public dmlc::FunctionRegEntryBase<OperatorPropertyReg, OperatorPropertyFactory> {
  inline OperatorPropertyReg& set_key_var_num_args(const std::string& key) {  // NOLINT(*)
    this->key_var_num_args = key;
    return *this;
  }
  inline OperatorPropertyReg& check_name() {
    OperatorProperty* p = this->body();
    std::string type    = p->TypeString();
    delete p;
    CHECK_EQ(this->name, type) << "Register Name and TypeString mismatch, name=\"" << this->name
                               << "\","
                               << " but TypeString=\"" << type << "\"";
    return *this;
  }
 
  std::string key_var_num_args;
};
 
//---------------------------------------------------------------------------------
// The following part are API Registration of Operators
// See also MXNET_REGISTER_SIMPLE_OP in operator_util.h for registering simple ops.
//---------------------------------------------------------------------------------
#define MXNET_REGISTER_OP_PROPERTY(name, OperatorPropertyType)                    \
  DMLC_REGISTRY_REGISTER(::mxnet::OperatorPropertyReg, OperatorPropertyReg, name) \
      .set_body([]() { return new OperatorPropertyType(); })                      \
      .set_return_type("NDArray-or-Symbol")                                       \
      .check_name()
 
#endif  // DMLC_USE_CXX11
}  // namespace mxnet
#endif  // MXNET_OPERATOR_H_