You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by ha...@apache.org on 2018/04/29 05:49:27 UTC
[incubator-mxnet] branch master updated: add FResourceRequestEx
interface (#10649)
This is an automated email from the ASF dual-hosted git repository.
haibin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new c090f95 add FResourceRequestEx interface (#10649)
c090f95 is described below
commit c090f956a850a3301b89cb856b16351752d7b70a
Author: Ziyue Huang <zy...@gmail.com>
AuthorDate: Sun Apr 29 13:49:18 2018 +0800
add FResourceRequestEx interface (#10649)
* FResourceRequestEx
* check for FResourceRequestEx in advance of FResourceRequest
* address comments
---
include/mxnet/op_attr_types.h | 13 ++++++++++++-
src/executor/attach_op_resource_pass.cc | 14 ++++++++++++--
src/imperative/imperative_utils.h | 10 +++++++---
src/operator/optimizer_op.cc | 10 +++++++---
4 files changed, 38 insertions(+), 9 deletions(-)
diff --git a/include/mxnet/op_attr_types.h b/include/mxnet/op_attr_types.h
index 800872b..3969d84 100644
--- a/include/mxnet/op_attr_types.h
+++ b/include/mxnet/op_attr_types.h
@@ -213,13 +213,24 @@ using FStatefulComputeEx = std::function<void (const OpStatePtr& state,
const std::vector<OpReqType>& req,
const std::vector<NDArray>& outputs)>;
/*!
- * \brief The resource request from the operator
+ * \brief The resource request from the operator.
+ * An operator could register ResourceRequestEx, or ResourceRequest, or neither.
*
* \note Register under "FResourceRequest"
*/
using FResourceRequest = std::function<
std::vector<ResourceRequest> (const NodeAttrs& n)>;
/*!
+ * \brief The resource request from the operator.
+ * An operator could register ResourceRequestEx, or ResourceRequest, or neither.
+ *
+ * \note Register under "FResourceRequestEx"
+ */
+using FResourceRequestEx = std::function<
+ std::vector<ResourceRequest> (const NodeAttrs& n,
+ const int dev_mask,
+ const DispatchMode dispatch_mode)>;
+/*!
* \brief Register an operator called as a NDArray function
*
* \note Register under "FNDArrayFunction"
diff --git a/src/executor/attach_op_resource_pass.cc b/src/executor/attach_op_resource_pass.cc
index 9a7ed09..6818662 100644
--- a/src/executor/attach_op_resource_pass.cc
+++ b/src/executor/attach_op_resource_pass.cc
@@ -33,9 +33,12 @@ namespace exec {
Graph AttachOpResources(Graph g) {
static auto& fresource =
nnvm::Op::GetAttr<FResourceRequest>("FResourceRequest");
+ static auto& fresource_ex =
+ nnvm::Op::GetAttr<FResourceRequestEx>("FResourceRequestEx");
auto& op_execs = nnvm::get<OpExecVector>(*g.attrs.at("op_execs"));
const auto& vctx = g.GetAttr<ContextVector>("context");
const auto& vdispatch = g.GetAttr<DispatchModeVector>("dispatch_mode");
+ const auto& dev_masks = g.GetAttr<DevMaskVector>("dev_mask");
const auto& idx = g.indexed_graph();
// Use global resource pool for each executor for now.
std::map<Context, Resource> cached_temp;
@@ -47,8 +50,15 @@ Graph AttachOpResources(Graph g) {
auto& requested = op_execs[nid]->op_ctx.requested;
requested.clear();
const auto op = inode.source->op();
- if (fresource.count(op) != 0) {
- auto reqs = fresource[op](inode.source->attrs);
+ const bool rsc_req = (fresource.count(op) != 0);
+ const bool rsc_ex_req = (fresource_ex.count(op) != 0);
+ CHECK(!(rsc_req && rsc_ex_req))
+ << "An operator could not register both ResourceRequestEx and ResourceRequest";
+ if (rsc_req || rsc_ex_req) {
+ auto reqs = rsc_ex_req ? fresource_ex[op](inode.source->attrs,
+ dev_masks[nid],
+ vdispatch[nid])
+ : fresource[op](inode.source->attrs);
// Get the resource of temporal space.
for (const ResourceRequest& req : reqs) {
if (req.type == ResourceRequest::kTempSpace) {
diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h
index 86683f9..10a011e 100644
--- a/src/imperative/imperative_utils.h
+++ b/src/imperative/imperative_utils.h
@@ -193,6 +193,7 @@ inline void SetDependency(const nnvm::NodeAttrs& attrs,
const DispatchMode dispatch_mode) {
static auto& fmutate = nnvm::Op::GetAttr<nnvm::FMutateInputs>("FMutateInputs");
static auto& ftmp_resource = nnvm::Op::GetAttr<FResourceRequest>("FResourceRequest");
+ static auto& ftmp_resource_ex = nnvm::Op::GetAttr<FResourceRequestEx>("FResourceRequestEx");
std::vector<engine::VarHandle>& read_vars = *p_read_vars;
std::vector<engine::VarHandle>& write_vars = *p_write_vars;
@@ -202,10 +203,13 @@ inline void SetDependency(const nnvm::NodeAttrs& attrs,
if (fmutate.count(attrs.op)) {
mutate_idx = fmutate[attrs.op](attrs);
}
-
- if (ftmp_resource.count(attrs.op)) {
+ const bool rsc_req = (ftmp_resource.count(attrs.op) != 0);
+ const bool rsc_ex_req = (ftmp_resource_ex.count(attrs.op) != 0);
+ if (rsc_req || rsc_ex_req) {
int ntmp = 0;
- auto resource_reqs = ftmp_resource[attrs.op](attrs);
+ auto resource_reqs = rsc_ex_req ? ftmp_resource_ex[attrs.op](attrs,
+ static_cast<int>(ctx.dev_mask()), dispatch_mode)
+ : ftmp_resource[attrs.op](attrs);
for (const auto& req : resource_reqs) {
switch (req.type) {
case ResourceRequest::kTempSpace:
diff --git a/src/operator/optimizer_op.cc b/src/operator/optimizer_op.cc
index fe0be9d..0e73c1d 100644
--- a/src/operator/optimizer_op.cc
+++ b/src/operator/optimizer_op.cc
@@ -344,9 +344,13 @@ only the row slices whose indices appear in grad.indices are updated (for both w
[](const nnvm::NodeAttrs& attrs) {
return std::vector<uint32_t>{2};
})
-.set_attr<FResourceRequest>("FResourceRequest",
- [](const NodeAttrs& attrs) {
- return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+.set_attr<FResourceRequestEx>("FResourceRequestEx",
+ [](const NodeAttrs& attrs, const int dev_mask, const DispatchMode dispatch_mode) {
+ std::vector<ResourceRequest> request;
+ if (dispatch_mode == DispatchMode::kFComputeEx) {
+ request.push_back(ResourceRequest::kTempSpace);
+ }
+ return request;
})
.set_attr<FCompute>("FCompute<cpu>", SGDMomUpdate<cpu>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SGDMomUpdateEx<cpu>)
--
To stop receiving notification emails like this one, please contact
haibin@apache.org.