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 <string>
 
#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 {
  bool need_grad;
  bool 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>();
  }
#if MXNET_USE_CUDA
 
  inline SyncedGPUAuxStream get_gpu_aux_stream() const {
    return run_ctx.get_gpu_aux_stream();
  }
#endif
};
 
enum class ExecType {
  kSync,
  kAsync,
  kCrossDeviceCopy,
  kSubgraphExec,
};
 
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,
};
 
enum class QuantizeType {
  // This operator doesn't support quantization
  kNone = 0,
  // This operator can get huge benefit from quantization, thus must be quantized
  kMust,
  // This operator support quantization, but will be decided depending on the connection
  kSupport,
};
 
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;
    auto state = new T(std::forward<Args>(args)...);
    auto var   = Engine::Get()->NewVariable();
    ret.ptr_.reset(new OpState(var, state), [](OpState* p) {
      Engine::Get()->DeleteVariable([](RunContext s) {}, Context::CPU(), p->var);
      delete reinterpret_cast<T*>(p->state);
      delete p;
    });
 
    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 *reinterpret_cast<T*>(ptr_->state);
  }
  /* \brief clear state */
  void reset() {
    ptr_.reset();
  }
  /* \brief checks whether the managed object is managed only by the current
            OpStatePtr instance */
  bool unique() const {
    return ptr_.unique();
  }
  /* \brief Whether state is empty */
  explicit operator bool() const {
    return ptr_ ? true : false;
  }
 
 private:
  /* \brief state structure */
  struct OpState {
    engine::VarHandle var;
    void* state;
 
    OpState(engine::VarHandle var_, void* state_) : var(var_), state(state_) {}
    OpState(const OpState& other) = delete;
    OpState& operator=(const OpState& other) = delete;
  };
  /* \brief shared pointer to state */
  std::shared_ptr<OpState> ptr_;
};
 
using FCreateOpState = std::function<OpStatePtr(const NodeAttrs& attrs,
                                                Context ctx,
                                                const mxnet::ShapeVector& in_shape,
                                                const std::vector<int>& in_type)>;
 
using THasDeterministicOutput = bool;
 
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 FResourceRequestEx =
    std::function<std::vector<ResourceRequest>(const NodeAttrs& n,
                                               const int dev_mask,
                                               const DispatchMode dispatch_mode)>;
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)>;
 
using FQuantizable = std::function<QuantizeType(const NodeAttrs& attrs)>;
 
using FQuantizedOp = std::function<nnvm::ObjectPtr(const NodeAttrs& attrs)>;
 
using FNeedRequantize = std::function<bool(const NodeAttrs& attrs)>;
 
using FAvoidQuantizeInput = std::function<
    bool(const NodeAttrs& attrs, const size_t index, const std::string quantize_granularity)>;
 
using FNeedAsymQuantizeInput = std::function<bool(const NodeAttrs& attrs, const size_t index)>;
 
using FAvoidDequantizeOutput = std::function<bool(const NodeAttrs& attrs, const size_t index)>;
 
using FNeedCalibrateInput = std::function<std::vector<int>(const NodeAttrs& attrs)>;
 
using FNeedCalibrateOutput = std::function<std::vector<int>(const NodeAttrs& attrs)>;
 
#if MXNET_USE_CUDA
 
using FIsCUDAGraphsCompatible = std::function<bool(const NodeAttrs& attrs, const bool is_train)>;
 
#endif
 
}  // namespace mxnet
 
#endif  // MXNET_OP_ATTR_TYPES_H_