mxnet
op_attr_types.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
26 #ifndef MXNET_OP_ATTR_TYPES_H_
27 #define MXNET_OP_ATTR_TYPES_H_
28 
29 #include <mshadow/tensor.h>
30 #include <nnvm/op_attr_types.h>
31 
32 #include <vector>
33 #include <functional>
34 
35 #include "./base.h"
36 #include "./ndarray.h"
37 #include "./engine.h"
38 #include "./resource.h"
39 
40 namespace mxnet {
41 
42 using nnvm::NodeAttrs;
43 
45 enum OpReqType {
58 };
59 
66 struct OpContext {
68  bool need_grad;
70  bool is_train;
76  std::vector<Resource> requested;
82  template<typename xpu>
83  inline mshadow::Stream<xpu>* get_stream() const {
84  return run_ctx.get_stream<xpu>();
85  }
86 };
87 
89 enum class ExecType {
91  kSync,
96  kAsync,
109 };
110 
112 enum class DispatchMode {
113  kUndefined = -1,
114  // dispatch on FCompute or FStatefulCompute
115  kFCompute,
116  // dispatch on FComputeEx or FStatefulComputeEx, if available
117  kFComputeEx,
118  // dispatch on FCompute or FStatefulCompute, and performs storage fallback
120  // special dispatch mode for variables
121  kVariable,
122 };
123 
128 class OpStatePtr {
129  public:
130  /* \brief Create a OpStatePtr with state of type T.
131  * \param args Arguments passed to T's constructor.
132  */
133  template<typename T, typename... Args>
134  static OpStatePtr Create(Args&&... args) {
135  OpStatePtr ret;
136  auto state = new T(std::forward<Args>(args)...);
137  auto var = Engine::Get()->NewVariable();
138  ret.ptr_.reset(
139  new OpState(var, state),
140  [](OpState* p) {
141  Engine::Get()->DeleteVariable([](RunContext s) {}, Context::CPU(), p->var);
142  delete reinterpret_cast<T*>(p->state);
143  delete p;
144  });
145 
146  return ret;
147  }
148  /* \brief Get engine variable associated with this state */
150  return ptr_->var;
151  }
152  /* \brief Get state of type T */
153  template<typename T>
154  T& get_state() const {
155  return *reinterpret_cast<T*>(ptr_->state);
156  }
157  /* \brief clear state */
158  void reset() {
159  ptr_.reset();
160  }
161  /* \brief checks whether the managed object is managed only by the current
162  OpStatePtr instance */
163  bool unique() const {
164  return ptr_.unique();
165  }
166  /* \brief Whether state is empty */
167  explicit operator bool() const {
168  return ptr_ ? true : false;
169  }
170 
171  private:
172  /* \brief state structure */
173  struct OpState {
174  engine::VarHandle var;
175  void* state;
176 
177  OpState(engine::VarHandle var_, void* state_) : var(var_), state(state_) {}
178  OpState(const OpState& other) = delete;
179  OpState& operator=(const OpState& other) = delete;
180  };
181  /* \brief shared pointer to state */
182  std::shared_ptr<OpState> ptr_;
183 };
184 
197 using FCreateOpState = std::function<OpStatePtr (const NodeAttrs& attrs,
198  Context ctx,
199  const std::vector<TShape>& in_shape,
200  const std::vector<int>& in_type)>;
204 using FExecType = std::function<ExecType (const NodeAttrs& attrs)>;
212 using FStatefulCompute = std::function<void (const OpStatePtr& state,
213  const OpContext& ctx,
214  const std::vector<TBlob>& inputs,
215  const std::vector<OpReqType>& req,
216  const std::vector<TBlob>& outputs)>;
224 using FStatefulComputeEx = std::function<void (const OpStatePtr& state,
225  const OpContext& ctx,
226  const std::vector<NDArray>& inputs,
227  const std::vector<OpReqType>& req,
228  const std::vector<NDArray>& outputs)>;
235 using FResourceRequest = std::function<
236  std::vector<ResourceRequest> (const NodeAttrs& n)>;
243 using FResourceRequestEx = std::function<
244  std::vector<ResourceRequest> (const NodeAttrs& n,
245  const int dev_mask,
246  const DispatchMode dispatch_mode)>;
252 using FNDArrayFunction = std::function<void (const nnvm::NodeAttrs& attrs,
253  const std::vector<NDArray>& inputs,
254  std::vector<NDArray>* outputs)>;
260 using FCompute = std::function<void (const nnvm::NodeAttrs& attrs,
261  const OpContext& ctx,
262  const std::vector<TBlob>& inputs,
263  const std::vector<OpReqType>& req,
264  const std::vector<TBlob>& outputs)>;
270 using FComputeEx = std::function<void (const nnvm::NodeAttrs& attrs,
271  const OpContext& ctx,
272  const std::vector<NDArray>& inputs,
273  const std::vector<OpReqType>& req,
274  const std::vector<NDArray>& outputs)>;
275 
282 using FInferStorageType = std::function<bool (const NodeAttrs& attrs,
283  const int dev_mask,
284  DispatchMode* dispatch_mode,
285  std::vector<int>* in_attrs,
286  std::vector<int>* out_attrs)>;
287 
292 using FQuantizedOp = std::function<nnvm::NodePtr (const NodeAttrs& attrs)>;
293 
300 using FNeedRequantize = std::function<bool (const NodeAttrs& attrs)>;
301 
302 } // namespace mxnet
303 
304 #endif // MXNET_OP_ATTR_TYPES_H_
std::function< void(const OpStatePtr &state, const OpContext &ctx, const std::vector< TBlob > &inputs, const std::vector< OpReqType > &req, const std::vector< TBlob > &outputs)> FStatefulCompute
Resiger a compute function for stateful operator. OpStatePtr is a pointer type, it&#39;s content is mutab...
Definition: op_attr_types.h:216
void reset()
Definition: op_attr_types.h:158
Forward/Backward are synchronize calls.
Engine that schedules all the operations according to dependency.
no operation, do not write anything
Definition: op_attr_types.h:47
write gradient to provided space
Definition: op_attr_types.h:49
namespace of mxnet
Definition: base.h:118
std::function< std::vector< ResourceRequest >(const NodeAttrs &n)> FResourceRequest
The resource request from the operator. An operator could register ResourceRequestEx, or ResourceRequest, or neither.
Definition: op_attr_types.h:236
std::function< void(const OpStatePtr &state, const OpContext &ctx, const std::vector< NDArray > &inputs, const std::vector< OpReqType > &req, const std::vector< NDArray > &outputs)> FStatefulComputeEx
Resiger a compute function for stateful operator using NDArray interface. OpStatePtr is a pointer typ...
Definition: op_attr_types.h:228
mshadow::Stream< xpu > * get_stream() const
get mshadow stream from Context
Definition: base.h:262
std::function< bool(const NodeAttrs &attrs)> FNeedRequantize
Register a function to determine if the output of a quantized operator needs to be requantized...
Definition: op_attr_types.h:300
Asynchronous function call.
bool is_train
whether it is training phase
Definition: op_attr_types.h:70
engine::VarHandle get_var() const
Definition: op_attr_types.h:149
Cross device copy operation, this is a special operator That indicates copy across devices...
execution time context. The information needed in runtime for actual execution.
Definition: base.h:249
DispatchMode
the dispatch mode of the operator
Definition: op_attr_types.h:112
T & get_state() const
Definition: op_attr_types.h:154
engine::CallbackOnComplete async_on_complete
the callback when operation completes, used by asynchronize ops
Definition: op_attr_types.h:74
std::function< OpStatePtr(const NodeAttrs &attrs, Context ctx, const std::vector< TShape > &in_shape, const std::vector< int > &in_type)> FCreateOpState
Create a Layer style, forward/backward operator. This is easy to write code that contains state...
Definition: op_attr_types.h:200
All the possible information needed by Operator.Forward and Backward This is the superset of RunConte...
Definition: op_attr_types.h:66
NDArray interface that handles array arithematics.
std::function< ExecType(const NodeAttrs &attrs)> FExecType
Execution mode of this operator.
Definition: op_attr_types.h:204
virtual VarHandle NewVariable()=0
Allocate a new variable, the variable can then be used to schedule the operation concurrently via dep...
static OpStatePtr Create(Args &&...args)
Definition: op_attr_types.h:134
std::function< nnvm::NodePtr(const NodeAttrs &attrs)> FQuantizedOp
Register a quantized node creation function based on the attrs of the node.
Definition: op_attr_types.h:292
bool need_grad
whether there is a backward phase to compute gradients.
Definition: op_attr_types.h:68
std::function< void(const nnvm::NodeAttrs &attrs, const std::vector< NDArray > &inputs, std::vector< NDArray > *outputs)> FNDArrayFunction
Register an operator called as a NDArray function.
Definition: op_attr_types.h:254
Global resource allocation handling.
virtual void DeleteVariable(SyncFn delete_fn, Context exec_ctx, VarHandle var)=0
Schedule the deletion of a variable.
Var * VarHandle
Variable pointer type, usually hold by user used to specify dependencies.
Definition: engine.h:47
perform an inplace write, This option only happen when Target shares memory with one of input argumen...
Definition: op_attr_types.h:55
A subgraph execution should happen in the main thread, instead of in the execution engine...
bool unique() const
Definition: op_attr_types.h:163
OpReqType
operation request type to Forward and Backward
Definition: op_attr_types.h:45
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)> FComputeEx
Resiger an NDArray compute function for simple stateless forward only operator.
Definition: op_attr_types.h:274
std::vector< Resource > requested
Resources requested by the operator.
Definition: op_attr_types.h:76
RunContext run_ctx
RunContext related resources.
Definition: op_attr_types.h:72
static Context CPU(int32_t dev_id=0)
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)> FCompute
Resiger a compute function for simple stateless forward only operator.
Definition: op_attr_types.h:264
OnComplete Callback to the engine, called by AsyncFn when action completes.
Definition: engine.h:56
configuration of MXNet as well as basic data structure.
static Engine * Get()
std::function< bool(const NodeAttrs &attrs, const int dev_mask, DispatchMode *dispatch_mode, std::vector< int > *in_attrs, std::vector< int > *out_attrs)> FInferStorageType
Resiger a storage and dispatch mode inference function based on storage types of the inputs and outpu...
Definition: op_attr_types.h:286
std::function< std::vector< ResourceRequest >(const NodeAttrs &n, const int dev_mask, const DispatchMode dispatch_mode)> FResourceRequestEx
The resource request from the operator. An operator could register ResourceRequestEx, or ResourceRequest, or neither.
Definition: op_attr_types.h:246
add to the provided space
Definition: op_attr_types.h:57
ExecType
the execution type of the operator
Definition: op_attr_types.h:89
Context information about the execution environment.
Definition: base.h:133
mshadow::Stream< xpu > * get_stream() const
get mshadow stream from Context
Definition: op_attr_types.h:83
Operator state. This is a pointer type, its content is mutable even if OpStatePtr is const...
Definition: op_attr_types.h:128