Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!694 Refactor vm module for multigraph sink. 

Merge pull request !694 from rick_sanchez/master
This commit is contained in:
mindspore-ci-bot 2020-04-28 10:42:26 +08:00 committed by Gitee
parent 8bf35b2b0c 6ae8345cad
commit c1813671db
11 changed files with 561 additions and 108 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

80
mindspore/ccsrc/session/ascend_session.cc
View File

@ -800,45 +800,77 @@ void AscendSession::UpdateGraphOrder(GraphId to_graph_id) {
}
}
size_t AscendSession::SetChildGraphInput(const KernelGraphPtr &graph, const AnfNodePtr &node, size_t input_index) {
auto output_num = AnfAlgo::GetOutputTensorNum(node);
if (output_num > 1 && !AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem)) {
return input_index + output_num;
}
auto &graph_inputs = graph->inputs();
auto &valid_inputs = graph->ValidInputs();
if (valid_inputs[input_index]) {
SetChildGraphParameter(node, graph_inputs[input_index]);
} else {
MS_LOG(DEBUG) << "Invalid input arg: " << node->DebugString();
}
return ++input_index;
}
size_t AscendSession::SetChildGraphInput(const KernelGraphPtr &graph, const ValuePtr &value, size_t input_index) {
MS_EXCEPTION_IF_NULL(value);
if (!value->isa<Tensor>()) {
MS_LOG(EXCEPTION) << "Value Node should be a tensor, unexpected value: " << value->ToString();
}
auto &graph_inputs = graph->inputs();
SetChildGraphParameter(value->cast<TensorPtr>(), graph_inputs[input_index]);
return ++input_index;
}
size_t AscendSession::SetChildGraphInput(const KernelGraphPtr &graph, const VectorRef &vec_args, size_t input_index) {
auto index = input_index;
for (auto &arg : vec_args) {
if (utils::isa<AnfNodePtr>(arg)) {
// arg is a anf node
auto node = utils::cast<AnfNodePtr>(arg);
index = SetChildGraphInput(graph, node, input_index);
} else if (utils::isa<ValuePtr>(arg)) {
// arg is a tensor
auto value = utils::cast<ValuePtr>(arg);
index = SetChildGraphInput(graph, value, input_index);
} else {
MS_LOG(EXCEPTION) << "Unexpected arg type " << arg.ToString();
}
}
return index;
}
void AscendSession::SetChildGraphInput(GraphId g, const VectorRef &args) {
MS_LOG(INFO) << "Set input of graph " << g;
auto to_graph = GetGraph(g);
MS_EXCEPTION_IF_NULL(to_graph);
DumpGraphInputArgs(args);
UpdateGraphOrder(g);
std::vector<AnfNodePtr> graph_inputs = to_graph->inputs();
auto valid_inputs = to_graph->ValidInputs();
auto &graph_inputs = to_graph->inputs();
auto real_args = GetRealArgs(to_graph, args);
size_t input_index = 0;
for (size_t i = 0; i < real_args.size(); i++) {
if (input_index >= graph_inputs.size()) {
MS_LOG(EXCEPTION) << "input_index " << input_index << " out of range size " << graph_inputs.size();
}
if (utils::isa<AnfNodePtr>(real_args[i])) {
auto &real_arg = real_args[i];
if (utils::isa<AnfNodePtr>(real_arg)) {
// arg is a anf node
auto real_arg = utils::cast<AnfNodePtr>(real_args[i]);
auto real_arg_output_num = AnfAlgo::GetOutputTensorNum(real_arg);
if (!AnfAlgo::CheckPrimitiveType(real_arg, prim::kPrimTupleGetItem) && real_arg_output_num > 1) {
input_index += real_arg_output_num;
continue;
}
if (valid_inputs[input_index]) {
SetChildGraphParameter(real_arg, graph_inputs[input_index]);
} else {
MS_LOG(DEBUG) << "Invalid input arg" << real_arg->DebugString();
}
input_index++;
} else if (utils::isa<ValuePtr>(args[i])) {
auto value = utils::cast<ValuePtr>(args[i]);
MS_EXCEPTION_IF_NULL(value);
auto node = utils::cast<AnfNodePtr>(real_arg);
input_index = SetChildGraphInput(to_graph, node, input_index);
} else if (utils::isa<ValuePtr>(real_arg)) {
// arg is a tensor
if (!value->isa<Tensor>()) {
MS_LOG(EXCEPTION) << "Value Node should be a tensor, unexpected value: " << value->ToString();
}
SetChildGraphParameter(value->cast<TensorPtr>(), graph_inputs[input_index]);
input_index++;
auto value = utils::cast<ValuePtr>(real_arg);
input_index = SetChildGraphInput(to_graph, value, input_index);
} else if (utils::isa<VectorRef>(real_arg)) {
// arg is a VectorRef
auto vec_args = utils::cast<VectorRef>(real_arg);
input_index = SetChildGraphInput(to_graph, vec_args, input_index);
} else {
MS_LOG(EXCEPTION) << "Unexpected arg type " << args[i].ToString();
MS_LOG(EXCEPTION) << "Unexpected arg type " << real_arg.ToString();
}
}
MS_LOG(INFO) << "Finish!";

4
mindspore/ccsrc/session/ascend_session.h
View File

@ -79,6 +79,10 @@ class AscendSession : public SessionBasic {
void RunOpHardwareOptimize(const std::shared_ptr<session::KernelGraph> &kernel_graph) const;
void RunOpExecTask(const std::shared_ptr<KernelGraph> &kernel_graph) const;
size_t SetChildGraphInput(const KernelGraphPtr &graph, const AnfNodePtr &node, size_t input_index);
size_t SetChildGraphInput(const KernelGraphPtr &graph, const ValuePtr &value, size_t input_index);
size_t SetChildGraphInput(const KernelGraphPtr &graph, const VectorRef &vec_args, size_t input_index);
// merge execution order list of child graphs
void MergeGraphExecOrder();
// insert assion op to sync data bettween different graphs

2
mindspore/ccsrc/session/kernel_graph.h
View File

@ -88,7 +88,7 @@ class KernelGraph : public FuncGraph {
void set_executable(bool executable) { executable_ = executable; }
// set invalid inputs for control sink
std::vector<bool> *MutableValidInputs() { return &valid_inputs_; }
std::vector<bool> ValidInputs() { return valid_inputs_; }
const std::vector<bool> &ValidInputs() const { return valid_inputs_; }
private:
// remove value node form graph

9
mindspore/ccsrc/utils/base_ref.h
View File

@ -228,6 +228,8 @@ T cast(const BaseRef &handle) {
class VectorRef : public BaseRef {
public:
using value_type = BaseRef;
VectorRef() {}
explicit VectorRef(const std::vector<BaseRef> &elements) : elements_(elements) {}
VectorRef(const const_iterator &begin, const const_iterator &end) : elements_(begin, end) {}
@ -251,6 +253,13 @@ class VectorRef : public BaseRef {
return elements_[dim];
}
BaseRef &operator[](const std::size_t &dim) {
if (dim >= size()) {
MS_LOG(EXCEPTION) << "Out of the size of the tuple.";
}
return elements_[dim];
}
uint32_t type() const override { return tid(); }
std::string ToString() const override;
std::vector<BaseRef> &elements() { return elements_; }

106
mindspore/ccsrc/vm/backend.cc
View File

@ -143,6 +143,66 @@ void MsBackend::SetSwitchGraph() {
}
}
// convert node from formal parameter to actual parameter,
// and actual parameter is graph user's formal parameter.
// get top while graph's parameter in recall while.
AnfNodePtr MsBackend::ConvertGraphInput(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
std::unordered_map<AnfNodePtr, size_t> params_index;
auto result = node;
auto graph = result->func_graph();
while (func_graph != graph) {
auto iter = graph_user_inputs_.find(graph);
if (iter == graph_user_inputs_.end()) {
break;
}
params_index.clear();
auto &params = graph->parameters();
for (size_t i = 0; i < params.size(); ++i) {
params_index[params[i]] = i;
}
graph = iter->second.first;
auto &inputs = iter->second.second;
result = inputs[params_index[result]];
}
return result;
}
void MsBackend::SetGraphUserInputs(const FuncGraphPtr &func_graph, const FuncGraphPtr &user,
const AnfNodePtrList &inputs) {
if (graph_user_inputs_.find(func_graph) != graph_user_inputs_.end()) {
return;
}
graph_user_inputs_[func_graph] = {user, inputs};
}
void MsBackend::RecallGraphInput(const FuncGraphPtr &func_graph, const VectorRef &args, const BaseRef &c) {
std::unordered_map<AnfNodePtr, size_t> params_index;
auto &params = func_graph->parameters();
for (size_t i = 0; i < params.size(); ++i) {
params_index[params[i]] = i;
}
// recall all child graphs in this while
auto &graph_inputs = graph_inputs_[c];
for (auto &iter : graph_inputs) {
auto &graph = iter.first;
auto &old_args = iter.second;
auto &result = graph_id_map_[graph];
auto &inputs = result.inputs;
for (size_t i = 0; i < inputs.size(); ++i) {
auto input = ConvertGraphInput(func_graph, inputs[i]);
auto it = params_index.find(input);
if (it != params_index.end()) {
old_args[i] = args[it->second];
}
}
sess_->SetChildGraphInput(graph, old_args);
}
graph_inputs_.erase(c);
}
// compile set input output
VectorRef MsBackend::MsSimuRunGraph(const GraphId &g, const VectorRef &args) {
MS_LOG(DEBUG) << "set graph input:" << g;
@ -150,13 +210,20 @@ VectorRef MsBackend::MsSimuRunGraph(const GraphId &g, const VectorRef &args) {
sess_->SetChildGraphInput(g, args);
if (is_switch_call_) {
bool curr_cond = simu_cond_map_[curr_switch_].curr_cond;
MS_LOG(DEBUG) << "switch call MsSimuRunGraph:" << curr_cond;
if (0 == simu_cond_map_[curr_switch_].cond_graph_map.count(curr_cond)) {
MS_LOG(DEBUG) << "switch call MsSimuRunGraph:" << curr_cond << ", " << g;
simu_cond_map_[curr_switch_].cond_graph_map[curr_cond] = g;
SetSwitchGraph();
if (!curr_switch_.is_null()) {
// push this {g, args} to all user while graph_inputs for nest while,
// when current condition recall over delete this cond in graph_inputs.
for (auto &iter : graph_inputs_) {
iter.second.push_back({g, args});
}
if (graph_inputs_.find(curr_switch_) == graph_inputs_.end()) {
graph_inputs_[curr_switch_].push_back({g, args});
}
}
bool curr_cond = simu_cond_map_[curr_switch_].curr_cond;
MS_LOG(DEBUG) << "switch call MsSimuRunGraph:" << curr_cond << ", " << g;
simu_cond_map_[curr_switch_].cond_graph_map[curr_cond] = g;
SetSwitchGraph();
}
std::vector<BaseRef> outputs;
@ -205,42 +272,17 @@ VectorRef MsBackend::MsRunGraph(const GraphId &g, const VectorRef &args) {
return outputs;
}
void MsBackend::SetSimuCondFlag(const BaseRef &c, int flag) {
MS_LOG(DEBUG) << "while set cond :" << c.ToString() << ", " << simu_cond_map_.size();
if (simu_cond_map_.find(c) == simu_cond_map_.end()) {
MS_LOG(EXCEPTION) << "error c not find";
}
simu_cond_map_[c].flag = flag;
}
int MsBackend::GetSimuCondFlag(const BaseRef &c) {
BaseRef cond = c;
if (cond.is_null()) {
MS_LOG(DEBUG) << "get curr_switch";
cond = curr_switch_;
}
if (simu_cond_map_.find(cond) == simu_cond_map_.end()) {
MS_LOG(ERROR) << "error c not find";
return -1;
}
return simu_cond_map_[cond].flag;
}
SwitchCondStatus MsBackend::SetSimuCond(const BaseRef &c, bool value) {
MS_LOG(DEBUG) << "set cond :" << c.ToString() << ", " << simu_cond_map_.size();
CondGraph cond_graph;
cond_graph.curr_cond = value;
if (simu_cond_map_.find(c) == simu_cond_map_.end()) {
cond_graph.flag = 0;
simu_cond_map_[c] = cond_graph;
}
if (simu_cond_map_[c].cond_graph_map.count(value)) {
if (value == true) {
return kCondAlreadyRun;
}
return kCondAlreadyRun;
}
simu_cond_map_[c].curr_cond = value;
MS_LOG(DEBUG) << "end set cond ";

16
mindspore/ccsrc/vm/backend.h
View File

@ -16,9 +16,11 @@
#ifndef MINDSPORE_CCSRC_VM_BACKEND_H_
#define MINDSPORE_CCSRC_VM_BACKEND_H_
#include <string>
#include <list>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include "ir/anf.h"
#include "vm/segment_runner.h"
@ -45,6 +47,8 @@ class Backend {
virtual bool GetCond(const BaseRef &c, bool *value);
virtual void SetSwitchGraph() {}
virtual void SetSwitchActive(const BaseRef &, bool) {}
virtual void RecallGraphInput(const FuncGraphPtr &, const VectorRef &, const BaseRef &) {}
virtual void SetGraphUserInputs(const FuncGraphPtr &, const FuncGraphPtr &, const AnfNodePtrList &) {}
void set_curr_switch(const BaseRef &value) {
curr_switch_ = value;
@ -54,8 +58,6 @@ class Backend {
BaseRef curr_switch() { return curr_switch_; }
virtual void Link(GraphId) {}
virtual LinConvertResult GetMultiGraphRun(const FuncGraphPtr &) { return LinConvertResult(); }
virtual void SetSimuCondFlag(const BaseRef &, int) {}
virtual int GetSimuCondFlag(const BaseRef &) { return 0; }
LinConvertResult multi_result() { return multi_result_; }
void set_multi_result(const LinConvertResult &value) { multi_result_ = value; }
@ -75,11 +77,11 @@ class Backend {
bool simu_flag_;
LinConvertResult multi_result_;
AnfNodePtr final_output_;
std::unordered_map<FuncGraphPtr, std::pair<FuncGraphPtr, AnfNodePtrList>> graph_user_inputs_;
};
struct CondGraph {
bool curr_cond;
int flag;
std::unordered_map<bool, GraphId> cond_graph_map;
};
@ -97,15 +99,17 @@ class MsBackend : public Backend {
void SetSwitchGraph() override;
void SetSwitchActive(const BaseRef &c, bool cond) override;
void RecallGraphInput(const FuncGraphPtr &, const VectorRef &, const BaseRef &) override;
void SetGraphUserInputs(const FuncGraphPtr &, const FuncGraphPtr &, const AnfNodePtrList &) override;
void Link(GraphId) override;
AnfNodePtr ConvertGraphInput(const FuncGraphPtr &, const AnfNodePtr &);
LinConvertResult GetMultiGraphRun(const FuncGraphPtr &g) override;
void SetSimuCondFlag(const BaseRef &c, int flag) override;
int GetSimuCondFlag(const BaseRef &c) override;
private:
session::SessionPtr sess_;
std::unordered_map<BaseRef, CondGraph, BaseRefHash> simu_cond_map_;
std::unordered_map<GraphId, LinConvertResult> graph_id_map_;
std::unordered_map<BaseRef, std::list<std::pair<GraphId, VectorRef>>, BaseRefHash> graph_inputs_;
};
} // namespace compile
} // namespace mindspore

23
mindspore/ccsrc/vm/transform.cc
View File

@ -390,6 +390,16 @@ void CompileGraph::AddTailCall(const AnfNodePtr &fn, size_t size) {
void CompileGraph::AddPartial(const CNodePtr &node) {
auto inputs = node->inputs();
VectorRef args;
auto fn = inputs[1];
if (!IsValueNode<FuncGraph>(fn)) {
MS_LOG(EXCEPTION) << "The type of 1st input of node must be FuncGraph";
}
if (backend_->is_multi_graph_sink()) {
auto func_graph = GetValueNode<FuncGraphPtr>(fn);
args.emplace_back(func_graph);
AnfNodePtrList outs(inputs.begin() + 2, inputs.end());
backend_->SetGraphUserInputs(func_graph, node->func_graph(), outs);
}
for (size_t i = 1; i < inputs.size(); i++) {
args.emplace_back(Ref(inputs[i]));
}
@ -442,12 +452,17 @@ void CompileGraph::AddPrimitive(const CNodePtr &node, const PrimitivePtr &prim)
}
int CompileGraph::AddCall(const FuncGraphPtr &graph, const CNodePtr &node) {
auto node_inputs = node->inputs();
AnfNodePtr fn = node_inputs[0];
auto inputs = node->inputs();
AnfNodePtr fn = inputs[0];
if (backend_->is_multi_graph_sink() && IsValueNode<FuncGraph>(fn)) {
auto func_graph = GetValueNode<FuncGraphPtr>(fn);
AnfNodePtrList outs(inputs.begin() + 1, inputs.end());
backend_->SetGraphUserInputs(func_graph, node->func_graph(), outs);
}
(void)Ref(fn);
size_t size = node_inputs.size();
size_t size = inputs.size();
for (size_t i = size - 1; i > 0; i--) {
AddInput(node_inputs[i]);
AddInput(inputs[i]);
}
if (node == graph->output()) {
AddTailCall(fn, size);

116
mindspore/ccsrc/vm/vm.cc
View File

@ -32,7 +32,8 @@ namespace compile {
// Arguments:
// fn_: Callable function.
// args_: Sequence of function args.
StructPartial::StructPartial(int fn, const VectorRef &args) : fn_(fn), args_(args) {}
// fg_: Graph of function.
StructPartial::StructPartial(int fn, const VectorRef &args, const FuncGraphPtr &fg) : fn_(fn), args_(args), fg_(fg) {}
std::ostream &operator<<(std::ostream &os, const StructPartial &other) {
os << "partial(" << other.fn_ << ", " << other.args_.ToString() << ")";
@ -40,7 +41,7 @@ std::ostream &operator<<(std::ostream &os, const StructPartial &other) {
}
bool operator==(const StructPartial &lhs, const StructPartial &rhs) {
return (lhs.fn_ == rhs.fn_ && lhs.args_ == rhs.args_);
return (lhs.fn_ == rhs.fn_ && lhs.args_ == rhs.args_ && lhs.fg_ == rhs.fg_);
}
StructSimuSwitch::StructSimuSwitch(const BaseRef &fn, const BaseRef &value) : fn_(fn), value_(value) {}
@ -242,16 +243,6 @@ void FinalVM::InstTailCall(const VectorRef &args) {
int nargs = utils::cast<int>(args[2]);
auto new_jmp = Ref(jmp);
if (backend_->simu_flag()) {
if (backend_->GetSimuCondFlag(BaseRef()) == 2) {
MS_LOG(DEBUG) << "invoke while call tail first";
Pop(height);
Push(1);
Popp();
return;
}
}
MoveStack(nargs, height);
MS_LOG(DEBUG) << "TailCall pushp:" << pc_ << ", jmp:" << jmp;
DoJmp(new_jmp);
@ -291,8 +282,30 @@ void FinalVM::InstReturn(const VectorRef &args) {
MS_LOG(DEBUG) << "End";
}
void FinalVM::InstPartial(const VectorRef &args) {
MS_LOG(DEBUG) << "Start";
void FinalVM::InstSimuPartial(const VectorRef &args) {
const size_t args_size = 2;
if (args.size() < args_size) {
MS_LOG(ERROR) << __FUNCTION__ << " requires " << args_size << " or more parameters, while the input size is "
<< args.size() << ".";
return;
}
auto &node = args[0];
if (!utils::isa<FuncGraphPtr>(node)) {
MS_LOG(ERROR) << "The type of 1st input of node must be FuncGraph";
return;
}
auto fg = utils::cast<FuncGraphPtr>(node);
int fn_ = utils::cast<int>(args[1]);
auto fn = utils::cast<int>(Ref(fn_));
MS_LOG(DEBUG) << "Partial argssize:" << args.size();
std::vector<BaseRef> outs(args.size() - 2);
(void)std::transform(args.begin() + 2, args.end(), outs.begin(),
[&, this](const BaseRef &a) { return Ref(utils::cast<int>(a)); });
Push(std::make_shared<StructPartial>(fn, VectorRef(outs), fg));
}
void FinalVM::InstRealPartial(const VectorRef &args) {
const size_t args_size = 1;
if (args.size() < args_size) {
MS_LOG(ERROR) << __FUNCTION__ << " requires " << args_size << " or more parameters, while the input size is "
@ -304,10 +317,18 @@ void FinalVM::InstPartial(const VectorRef &args) {
auto fn = utils::cast<int>(Ref(fn_));
MS_LOG(DEBUG) << "Partial argssize:" << args.size();
std::vector<BaseRef> outs(args.size() - 1);
(void)std::transform(args.begin() + 1, args.end(), outs.begin(),
[&, this](const BaseRef &a) { return Ref(utils::cast<int>(a)); });
Push(std::make_shared<StructPartial>(fn, VectorRef(outs)));
}
void FinalVM::InstPartial(const VectorRef &args) {
MS_LOG(DEBUG) << "Start";
if (backend_->is_multi_graph_sink()) {
InstSimuPartial(args);
} else {
InstRealPartial(args);
}
MS_LOG(DEBUG) << "End";
}
@ -328,43 +349,57 @@ void FinalVM::InstSimuSwitch(const VectorRef &args) {
bool bool_value = cond;
SwitchCondStatus cond_stat = backend_->SetSimuCond(c, bool_value);
int cond_flag = backend_->GetSimuCondFlag(c);
MS_LOG(DEBUG) << "Simu switch cond:" << cond << ", " << cond_flag << ", " << c.cast<AnfNodePtr>()->DebugString();
if (cond_flag == 2) {
Popp();
Popp();
backend_->SetSimuCondFlag(c, 0);
return;
}
if (cond_stat == kCondAlreadyRun) {
MS_LOG(DEBUG) << "switch alreay run bool while true jmp";
if (cond_flag == 0) {
MS_LOG(DEBUG) << "switch second run bool while true jmp";
backend_->SetSwitchActive(c, true);
Push(std::make_shared<StructSimuSwitch>(Ref(vtrue), c));
Pushsp();
backend_->SetSimuCondFlag(c, 1);
return;
} else if (cond_flag == 1) {
MS_LOG(DEBUG) << "switch first run bool while if jmp";
Push(std::make_shared<StructSimuSwitch>(Ref(vfalse), c));
(void)backend_->SetSimuCond(c, false);
backend_->SetSimuCondFlag(c, 2);
return;
} else {
MS_LOG(EXCEPTION) << "error cond not find";
return;
BaseRef jmp = Ref(vtrue);
if (utils::isa<StructPartial>(jmp)) {
auto new_jmp = utils::cast<std::shared_ptr<StructPartial>>(jmp);
backend_->RecallGraphInput(new_jmp->fg_, new_jmp->args_, c);
}
cond_jmp_[c] = Ref(vfalse);
Push(static_cast<int>(cond_stat));
Popp();
backend_->SetSwitchActive(c, bool_value);
return;
}
if (bool_value) {
Push(std::make_shared<StructSimuSwitch>(Ref(vtrue), c));
Pushsp();
} else {
MergeJmpArgs(Ref(vfalse), c);
Push(std::make_shared<StructSimuSwitch>(Ref(vfalse), c));
}
}
void FinalVM::MergeJmpArgs(const BaseRef &jmp, const BaseRef &c) {
auto iter = cond_jmp_.find(c);
if (iter == cond_jmp_.end()) {
return;
}
auto old_jmp = utils::cast<std::shared_ptr<StructPartial>>(iter->second);
auto new_jmp = utils::cast<std::shared_ptr<StructPartial>>(jmp);
auto &old_args = old_jmp->args_;
auto &new_args = new_jmp->args_;
for (size_t i = 0; i < new_args.size(); ++i) {
auto &old_arg = old_args[i];
auto &new_arg = new_args[i];
if (utils::isa<VectorRef>(old_arg)) {
auto old_vec_ref = utils::cast<VectorRef>(old_arg);
if (utils::isa<VectorRef>(new_arg)) {
auto new_vec_ref = utils::cast<VectorRef>(new_arg);
std::copy(new_vec_ref.begin(), new_vec_ref.end(), std::back_inserter(old_vec_ref));
}
new_arg = old_vec_ref;
} else if (utils::isa<VectorRef>(new_arg)) {
auto new_vec_ref = utils::cast<VectorRef>(new_arg);
new_vec_ref.push_back(old_arg);
new_arg = new_vec_ref;
} else {
new_arg = VectorRef({new_arg, old_arg});
}
}
}
void FinalVM::InstRealSwitch(const VectorRef &args) {
const size_t args_size = 3;
if (args.size() != args_size) {
@ -399,6 +434,7 @@ void FinalVM::InstSwitch(const VectorRef &args) {
} else {
InstRealSwitch(args);
}
MS_LOG(DEBUG) << "End";
}
void FinalVM::InstTuple(const VectorRef &args) {

10
mindspore/ccsrc/vm/vm.h
View File

@ -27,6 +27,9 @@
#include <utility>
#include <vector>
#include <deque>
#include <unordered_map>
#include "ir/anf.h"
#include "utils/base_ref.h"
namespace mindspore {
@ -60,13 +63,14 @@ const std::vector<std::string> inst_str{"call", "tail_call", "return", "partial
class StructPartial : public Base {
public:
// Initialize StructPartial.
StructPartial(int fn, const VectorRef &args);
StructPartial(int fn, const VectorRef &args, const FuncGraphPtr &fg = nullptr);
virtual ~StructPartial() = default;
MS_DECLARE_PARENT(StructPartial, Base)
int fn_;
VectorRef args_;
FuncGraphPtr fg_;
};
std::ostream &operator<<(std::ostream &os, const StructPartial &other);
@ -98,6 +102,8 @@ class FinalVM {
void InstTailCall(const VectorRef &args);
void InstReturn(const VectorRef &args);
void InstPartial(const VectorRef &args);
void InstSimuPartial(const VectorRef &args);
void InstRealPartial(const VectorRef &args);
void InstSwitch(const VectorRef &args);
void InstSimuSwitch(const VectorRef &args);
void InstRealSwitch(const VectorRef &args);
@ -120,6 +126,7 @@ class FinalVM {
void Pushsp();
void Popsp();
void DoJmp(const BaseRef &jmp);
void MergeJmpArgs(const BaseRef &jmp, const BaseRef &c);
private:
InstSet insts_;
@ -128,6 +135,7 @@ class FinalVM {
std::stack<int> retsp_;
int pc_;
int sp_;
std::unordered_map<BaseRef, BaseRef, BaseRefHash> cond_jmp_;
BackendPtr backend_;
const InstFunctionMap inst_function_map = {
{Instruction::kCall, [this](const VectorRef &args) { InstCall(args); }},

184
tests/st/control/test_multigraph_sink.py Normal file
View File

@ -0,0 +1,184 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed 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.
# ============================================================================
""" test_multigraph_sink """
import pytest
import numpy as np
import mindspore.nn as nn
import mindspore.context as context
from mindspore.common.tensor import Tensor
from mindspore.common import dtype as mstype
from mindspore.common import ms_function
from mindspore.ops import operations as P
def setup_module(module):
context.set_context(mode = context.PYNATIVE_MODE, save_graphs = True, device_target = "Ascend")
context.set_context(enable_task_sink = True, device_id = 0)
c1 = Tensor([2], mstype.int32)
c2 = Tensor([14], mstype.int32)
c3 = Tensor([1], mstype.int32)
c4 = Tensor([0], mstype.int32)
c5 = Tensor([14], mstype.int32)
@ms_function
def simple_if(x, y, z):
if x < y:
x = x + 1
else:
x = x + 2
x = x + 3
return x
@ms_function
def if_by_if(x, y, z):
if x < y:
x = x + 1
if y > x:
x = x + 2
x = x + 3
return x
@ms_function
def if_in_if(x, y, z):
out = c4
if x < y:
z = c4 + c4
if z < y:
z = z + 2
out = out + z
x = x + 3
out = out + x
return out
@ms_function
def simple_while(x, y, z):
y = y + 4
while x < y:
x = x + 1
x = x + 3
return x
@ms_function
def while_by_while(x, y, z):
while x < y:
x = x + 1
while z < c5:
z = z + 1
x = x + 1
x = x + 1
return x
@ms_function
def while_in_while(x, y, z):
out = c4
while x < y:
z = c4 + c4
while z < y:
z = z + 1
out = out + z
x = x + 1
out = out + x
return out
@ms_function
def while_by_while_in_while(x, y, z):
out = c4
while x < c2:
y = c4 + c4
while y < c2:
y = y + 1
out = out + y
z = c4 + c4
while z < c2:
z = z + 1
out = out + z
x = x + 1
out = out + x
return out
@ms_function
def while_in_while_in_while(x, y, z):
out = c4
while x < c2:
y = c4 + c4
while y < c2:
y = y + 1
z = c4 + c4
while z < c2:
z = z + 1
out = out + z
out = out + y
x = x + 1
out = out + x
return out
def test_simple_if():
output = simple_if(c1, c2, c3)
expect = Tensor([6], mstype.int32)
assert output == expect
def test_if_by_if():
output = if_by_if(c1, c2, c3)
expect = Tensor([8], mstype.int32)
assert output == expect
def test_if_in_if():
output = if_in_if(c1, c2, c3)
expect = Tensor([7], mstype.int32)
assert output == expect
def test_simple_while():
output = simple_while(c1, c2, c3)
expect = Tensor([21], mstype.int32)
assert output == expect
def test_while_by_while():
output = while_by_while(c1, c2, c3)
expect = Tensor([28], mstype.int32)
assert output == expect
def test_while_in_while():
output = while_in_while(c1, c2, c3)
expect = Tensor([1274], mstype.int32)
assert output == expect
def test_while_by_while_in_while():
output = while_by_while_in_while(c1, c2, c3)
expect = Tensor([350], mstype.int32)
assert output == expect
def test_while_in_while_in_while():
output = while_in_while_in_while(c1, c2, c3)
expect = Tensor([2534], mstype.int32)
assert output == expect

119
tests/ut/python/pynative_mode/test_multigraph_sink.py Normal file
View File

@ -0,0 +1,119 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed 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.
# ============================================================================
""" test_multigraph_sink """
import pytest
import numpy as np
import mindspore.nn as nn
import mindspore.context as context
from mindspore.common.tensor import Tensor
from mindspore.common import dtype as mstype
from mindspore.common import ms_function
from mindspore.ops import operations as P
def setup_module(module):
context.set_context(mode = context.PYNATIVE_MODE, save_graphs = True, device_target = "Ascend")
context.set_context(enable_task_sink = True, device_id = 0)
c1 = Tensor([2], mstype.int32)
c2 = Tensor([14], mstype.int32)
c3 = Tensor([1], mstype.int32)
c4 = Tensor([0], mstype.int32)
c5 = Tensor([14], mstype.int32)
@ms_function
def simple_if(x, y, z):
if x < y:
x = x + 1
else:
x = x + 2
x = x + 3
return x
@ms_function
def if_by_if(x, y, z):
if x < y:
x = x + 1
if y > x:
x = x + 2
x = x + 3
return x
@ms_function
def if_in_if(x, y, z):
out = c4
if x < y:
z = c4 + c4
if z < y:
z = z + 2
out = out + z
x = x + 3
out = out + x
return out
@ms_function
def simple_while(x, y, z):
y = y + 4
while x < y:
x = x + 1
x = x + 3
return x
@ms_function
def while_by_while(x, y, z):
while x < y:
x = x + 1
while z < c5:
z = z + 1
x = x + 1
x = x + 1
return x
def test_simple_if():
output = simple_if(c1, c2, c3)
expect = Tensor([6], mstype.int32)
assert output == expect
def test_if_by_if():
output = if_by_if(c1, c2, c3)
expect = Tensor([8], mstype.int32)
assert output == expect
def test_if_in_if():
output = if_in_if(c1, c2, c3)
expect = Tensor([7], mstype.int32)
assert output == expect
def test_simple_while():
output = simple_while(c1, c2, c3)
expect = Tensor([21], mstype.int32)
assert output == expect
def test_while_by_while():
output = while_by_while(c1, c2, c3)
expect = Tensor([28], mstype.int32)
assert output == expect