op_attr_types.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_OP_ATTR_TYPES_H_
#define MXNET_OP_ATTR_TYPES_H_

#include <mshadow/tensor.h>
#include <nnvm/op_attr_types.h>

#include <vector>
#include <functional>

#include "./base.h"
#include "./ndarray.h"
#include "./engine.h"
#include "./resource.h"

namespace mxnet {

using nnvm::NodeAttrs;

enum OpReqType {
  kNullOp,
  kWriteTo,
  kWriteInplace,
  kAddTo
};

struct OpContext {
  int is_train;
  RunContext run_ctx;
  engine::CallbackOnComplete async_on_complete;
  std::vector<Resource> requested;
  template<typename xpu>
  inline mshadow::Stream<xpu>* get_stream() const {
    return run_ctx.get_stream<xpu>();
  }
};

enum class ExecType {
  kSync,
  kAsync,
  kLocal,
  kCrossDeviceCopy
};

enum class DispatchMode {
  kUndefined = -1,
  // dispatch on FCompute or FStatefulCompute
  kFCompute,
  // dispatch on FComputeEx or FStatefulComputeEx, if available
  kFComputeEx,
  // dispatch on FCompute or FStatefulCompute, and performs storage fallback
  kFComputeFallback,
  // special dispatch mode for variables
  kVariable,
};

class OpStatePtr {
 public:
  /* \brief Create a OpStatePtr with state of type T.
   * \param args Arguments passed to T's constructor.
   */
  template<typename T, typename... Args>
  static OpStatePtr Create(Args&&... args) {
    OpStatePtr ret;
    ret.ptr_ = std::make_shared<OpState>();
    ret.ptr_->var_ = Engine::Get()->NewVariable();
    ret.ptr_->state_.construct<T>(std::forward<Args>(args)...);

    return ret;
  }
  /* \brief Get engine variable associated with this state */
  engine::VarHandle get_var() const {
    return ptr_->var_;
  }
  /* \brief Get state of type T */
  template<typename T>
  T& get_state() const {
    return dmlc::get<T>(ptr_->state_);
  }
  /* \brief clear state */
  void reset() {
    ptr_.reset();
  }
  /* \brief Whether state is empty */
  explicit operator bool() const {
    return ptr_ ? true : false;
  }

 private:
  /* \brief state structure */
  struct OpState {
    OpState() {}
    OpState(const OpState& other) = delete;
    OpState& operator=(const OpState& other) = delete;

    ~OpState() {
      Engine::Get()->DeleteVariable([](RunContext s) {}, Context::CPU(), var_);
    }

    engine::VarHandle var_;
    dmlc::any state_;
  };
  /* \brief shared pointer to state */
  std::shared_ptr<OpState> ptr_;
};

using FCreateOpState = std::function<OpStatePtr (const NodeAttrs& attrs,
                                                 Context ctx,
                                                 const std::vector<TShape>& in_shape,
                                                 const std::vector<int>& in_type)>;
using FExecType = std::function<ExecType (const NodeAttrs& attrs)>;
using FStatefulCompute = std::function<void (const OpStatePtr& state,
                                             const OpContext& ctx,
                                             const std::vector<TBlob>& inputs,
                                             const std::vector<OpReqType>& req,
                                             const std::vector<TBlob>& outputs)>;
using FStatefulComputeEx = std::function<void (const OpStatePtr& state,
                                               const OpContext& ctx,
                                               const std::vector<NDArray>& inputs,
                                               const std::vector<OpReqType>& req,
                                               const std::vector<NDArray>& outputs)>;
using FResourceRequest = std::function<
  std::vector<ResourceRequest> (const NodeAttrs& n)>;
using FNDArrayFunction = std::function<void (const nnvm::NodeAttrs& attrs,
                                             const std::vector<NDArray>& inputs,
                                             std::vector<NDArray>* outputs)>;
using FCompute = std::function<void (const nnvm::NodeAttrs& attrs,
                                     const OpContext& ctx,
                                     const std::vector<TBlob>& inputs,
                                     const std::vector<OpReqType>& req,
                                     const std::vector<TBlob>& outputs)>;
using FComputeEx = std::function<void (const nnvm::NodeAttrs& attrs,
                                       const OpContext& ctx,
                                       const std::vector<NDArray>& inputs,
                                       const std::vector<OpReqType>& req,
                                       const std::vector<NDArray>& outputs)>;

using FInferStorageType = std::function<bool (const NodeAttrs& attrs,
                                              const int dev_mask,
                                              DispatchMode* dispatch_mode,
                                              std::vector<int>* in_attrs,
                                              std::vector<int>* out_attrs)>;

}  // namespace mxnet

#endif  // MXNET_OP_ATTR_TYPES_H_