forked from mindspore-Ecosystem/mindspore
!3169 dumpir can dump subgraphs of ascend kernel graph
Merge pull request !3169 from zhoufeng/xiu-ba-ge
This commit is contained in:
commit
07f1d35ede
|
@ -107,7 +107,7 @@ static void ReuseParameter(NotNull<KernelGraphPtr> root_kg,
|
||||||
|
|
||||||
static CNodePtr GetNextRealKernel(const std::vector<CNodePtr> &list, size_t start) {
|
static CNodePtr GetNextRealKernel(const std::vector<CNodePtr> &list, size_t start) {
|
||||||
for (size_t i = start; i < list.size() - 1; ++i) {
|
for (size_t i = start; i < list.size() - 1; ++i) {
|
||||||
if (!IsPrimitiveCNode(list[i], prim::kPrimPartial) && AnfAlgo::IsRealKernel(list[i])) {
|
if (AnfAlgo::IsRealKernel(list[i])) {
|
||||||
return list[i];
|
return list[i];
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -171,18 +171,43 @@ static void EraseNodeFromExecOrder(const AnfNodePtr &node, const NotNull<std::ve
|
||||||
exec_order->erase(exec_iter);
|
exec_order->erase(exec_iter);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
void AscendControlParser::AttachChildGraphToReturnNode(NotNull<KernelGraphPtr> graph,
|
||||||
|
const NotNull<std::set<KernelGraphPtr> *> memo) {
|
||||||
|
if (memo->find(graph) != memo->end()) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
memo->insert(graph.get());
|
||||||
|
const std::vector<std::shared_ptr<KernelGraph>> &child_graph_order = graph->child_graph_order();
|
||||||
|
if (child_graph_order.empty()) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<AnfNodePtr> depend_inputs = {NewValueNode(std::make_shared<Primitive>(prim::kPrimPartial->name()))};
|
||||||
|
for (auto &cg : child_graph_order) {
|
||||||
|
MS_EXCEPTION_IF_NULL(cg);
|
||||||
|
auto fg = cg->cast<FuncGraphPtr>();
|
||||||
|
MS_EXCEPTION_IF_NULL(fg);
|
||||||
|
depend_inputs.emplace_back(NewValueNode(fg));
|
||||||
|
AttachChildGraphToReturnNode(NOT_NULL(cg), memo);
|
||||||
|
}
|
||||||
|
auto child_graphs = graph->NewCNode(depend_inputs);
|
||||||
|
InsertDependToGraph(graph, NOT_NULL(child_graphs));
|
||||||
|
}
|
||||||
|
|
||||||
void AscendControlParser::LinkGraph(NotNull<KernelGraphPtr> kg) {
|
void AscendControlParser::LinkGraph(NotNull<KernelGraphPtr> kg) {
|
||||||
std::set<KernelGraphPtr> memo;
|
std::set<KernelGraphPtr> memo;
|
||||||
std::vector<std::pair<AnfNodePtr, AnfNodePtr>> link_list;
|
std::vector<std::pair<AnfNodePtr, AnfNodePtr>> link_list;
|
||||||
// Insert Assign
|
// Insert Assign
|
||||||
ChildGraphDataAssign(kg, NOT_NULL(&link_list), NOT_NULL(&memo));
|
ChildGraphDataAssign(kg, NOT_NULL(&link_list), NOT_NULL(&memo));
|
||||||
|
memo.clear();
|
||||||
// Reuse Parameter
|
// Reuse Parameter
|
||||||
ReuseParameter(kg, link_list);
|
ReuseParameter(kg, link_list);
|
||||||
// replace call by label goto / label switch
|
// replace call by label goto / label switch
|
||||||
memo.clear();
|
|
||||||
(void)ProcessKernelGraph(kg, nullptr, nullptr, NOT_NULL(&memo));
|
(void)ProcessKernelGraph(kg, nullptr, nullptr, NOT_NULL(&memo));
|
||||||
|
memo.clear();
|
||||||
// assign label resource
|
// assign label resource
|
||||||
device::ascend::AscendLabelAssign::GetInstance().AssignLabel(kg);
|
device::ascend::AscendLabelAssign::GetInstance().AssignLabel(kg);
|
||||||
|
AttachChildGraphToReturnNode(kg, NOT_NULL(&memo));
|
||||||
}
|
}
|
||||||
|
|
||||||
void AscendControlParser::EraseParameter(NotNull<KernelGraphPtr> root_graph,
|
void AscendControlParser::EraseParameter(NotNull<KernelGraphPtr> root_graph,
|
||||||
|
|
|
@ -66,7 +66,8 @@ class AscendControlParser {
|
||||||
static AnfNodePtr InsertAssignToGraph(NotNull<KernelGraphPtr> kg, NotNull<AnfNodePtr> from, NotNull<AnfNodePtr> to);
|
static AnfNodePtr InsertAssignToGraph(NotNull<KernelGraphPtr> kg, NotNull<AnfNodePtr> from, NotNull<AnfNodePtr> to);
|
||||||
static std::vector<std::pair<KernelGraphPtr, std::vector<AnfNodePtr>>> ParseCallNode(NotNull<CNodePtr> call_node);
|
static std::vector<std::pair<KernelGraphPtr, std::vector<AnfNodePtr>>> ParseCallNode(NotNull<CNodePtr> call_node);
|
||||||
static std::tuple<KernelGraphPtr, std::vector<AnfNodePtr>> ParsePartial(NotNull<AnfNodePtr> node);
|
static std::tuple<KernelGraphPtr, std::vector<AnfNodePtr>> ParsePartial(NotNull<AnfNodePtr> node);
|
||||||
|
static void AttachChildGraphToReturnNode(NotNull<KernelGraphPtr> graph,
|
||||||
|
const NotNull<std::set<KernelGraphPtr> *> memo);
|
||||||
// root graph order
|
// root graph order
|
||||||
static bool CheckLabelIndex(uint32_t order_index, uint32_t label_index, const CNodePtr &cnode,
|
static bool CheckLabelIndex(uint32_t order_index, uint32_t label_index, const CNodePtr &cnode,
|
||||||
NotNull<KernelGraphPtr> graph);
|
NotNull<KernelGraphPtr> graph);
|
||||||
|
|
Loading…
Reference in New Issue