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 
25 #ifndef MXNET_OP_ATTR_TYPES_H_
26 #define MXNET_OP_ATTR_TYPES_H_
27 
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  int is_train;
74  std::vector<Resource> requested;
80  template<typename xpu>
81  inline mshadow::Stream<xpu>* get_stream() const {
82  return run_ctx.get_stream<xpu>();
83  }
84 };
85 
87 enum class ExecType {
89  kSync,
94  kAsync,
96  kLocal,
104 };
105 
110 class OpStatePtr {
111  public:
112  /* \brief Create a OpStatePtr with state of type T.
113  * \param args Arguments passed to T's constructor.
114  */
115  template<typename T, typename... Args>
116  static OpStatePtr Create(Args&&... args) {
117  OpStatePtr ret;
118  ret.ptr_ = std::make_shared<OpState>();
119  ret.ptr_->var_ = Engine::Get()->NewVariable();
120  ret.ptr_->state_.construct<T>(std::forward<Args>(args)...);
121 
122  return ret;
123  }
124  /* \brief Get engine variable associated with this state */
126  return ptr_->var_;
127  }
128  /* \brief Get state of type T */
129  template<typename T>
130  T& get_state() const {
131  return dmlc::get<T>(ptr_->state_);
132  }
133  /* \brief clear state */
134  void reset() {
135  ptr_.reset();
136  }
137  /* \brief Whether state is empty */
138  explicit operator bool() const {
139  return ptr_ ? true : false;
140  }
141 
142  private:
143  /* \brief state structure */
144  struct OpState {
145  OpState() {}
146  OpState(const OpState& other) = delete;
147  OpState& operator=(const OpState& other) = delete;
148 
149  ~OpState() {
150  Engine::Get()->DeleteVariable([](RunContext s) {}, Context::CPU(), var_);
151  }
152 
153  engine::VarHandle var_;
154  dmlc::any state_;
155  };
156  /* \brief shared pointer to state */
157  std::shared_ptr<OpState> ptr_;
158 };
159 
172 using FCreateOpState = std::function<OpStatePtr (const NodeAttrs& attrs,
173  Context ctx,
174  const std::vector<TShape>& in_shape,
175  const std::vector<int>& in_type)>;
179 using FExecType = std::function<ExecType (const NodeAttrs& attrs)>;
187 using FStatefulCompute = std::function<void (const OpStatePtr& state,
188  const OpContext& ctx,
189  const std::vector<TBlob>& inputs,
190  const std::vector<OpReqType>& req,
191  const std::vector<TBlob>& outputs)>;
199 using FStatefulComputeEx = std::function<void (const OpStatePtr& state,
200  const OpContext& ctx,
201  const std::vector<NDArray>& inputs,
202  const std::vector<OpReqType>& req,
203  const std::vector<NDArray>& outputs)>;
209 using FResourceRequest = std::function<
210  std::vector<ResourceRequest> (const NodeAttrs& n)>;
216 using FNDArrayFunction = std::function<void (const nnvm::NodeAttrs& attrs,
217  const std::vector<NDArray>& inputs,
218  std::vector<NDArray>* outputs)>;
224 using FCompute = std::function<void (const nnvm::NodeAttrs& attrs,
225  const OpContext& ctx,
226  const std::vector<TBlob>& inputs,
227  const std::vector<OpReqType>& req,
228  const std::vector<TBlob>& outputs)>;
229 } // namespace mxnet
230 
231 #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:191
void reset()
Definition: op_attr_types.h:134
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:126
std::function< std::vector< ResourceRequest >(const NodeAttrs &n)> FResourceRequest
The resource request from the operator.
Definition: op_attr_types.h:210
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:203
mshadow::Stream< xpu > * get_stream() const
get mshadow stream from Context
Definition: base.h:251
Asynchronous function call.
engine::VarHandle get_var() const
Definition: op_attr_types.h:125
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:238
Run this operator on the scheduling thread without pushing to engine.
T & get_state() const
Definition: op_attr_types.h:130
engine::CallbackOnComplete async_on_complete
the callback when operation completes, used by asynchronize ops
Definition: op_attr_types.h:72
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:175
All the possible information needed by Operator.Forward and Backward This is the superset of RunConte...
Definition: op_attr_types.h:66
int is_train
whether it is training phase
Definition: op_attr_types.h:68
std::function< ExecType(const NodeAttrs &attrs)> FExecType
Execution mode of this operator.
Definition: op_attr_types.h:179
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:116
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:218
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:46
perform an inplace write, Target shares memory with one of input arguments. This option only happen w...
Definition: op_attr_types.h:55
OpReqType
operation request type to Forward and Backward
Definition: op_attr_types.h:45
std::vector< Resource > requested
Resources requested by the operator.
Definition: op_attr_types.h:74
RunContext run_ctx
RunContext related resources.
Definition: op_attr_types.h:70
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:228
OnComplete Callback to the engine, called by AsyncFn when action completes.
Definition: engine.h:55
static Engine * Get()
add to the provided space
Definition: op_attr_types.h:57
ExecType
the execution type of the operator
Definition: op_attr_types.h:87
Context information about the execution environment.
Definition: base.h:141
mshadow::Stream< xpu > * get_stream() const
get mshadow stream from Context
Definition: op_attr_types.h:81
Operator state. This is a pointer type, its content is mutable even if OpStatePtr is const...
Definition: op_attr_types.h:110