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) {
|
||||
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];
|
||||
}
|
||||
}
|
||||
|
@ -171,18 +171,43 @@ static void EraseNodeFromExecOrder(const AnfNodePtr &node, const NotNull<std::ve
|
|||
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) {
|
||||
std::set<KernelGraphPtr> memo;
|
||||
std::vector<std::pair<AnfNodePtr, AnfNodePtr>> link_list;
|
||||
// Insert Assign
|
||||
ChildGraphDataAssign(kg, NOT_NULL(&link_list), NOT_NULL(&memo));
|
||||
memo.clear();
|
||||
// Reuse Parameter
|
||||
ReuseParameter(kg, link_list);
|
||||
// replace call by label goto / label switch
|
||||
memo.clear();
|
||||
(void)ProcessKernelGraph(kg, nullptr, nullptr, NOT_NULL(&memo));
|
||||
memo.clear();
|
||||
// assign label resource
|
||||
device::ascend::AscendLabelAssign::GetInstance().AssignLabel(kg);
|
||||
AttachChildGraphToReturnNode(kg, NOT_NULL(&memo));
|
||||
}
|
||||
|
||||
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 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 void AttachChildGraphToReturnNode(NotNull<KernelGraphPtr> graph,
|
||||
const NotNull<std::set<KernelGraphPtr> *> memo);
|
||||
// root graph order
|
||||
static bool CheckLabelIndex(uint32_t order_index, uint32_t label_index, const CNodePtr &cnode,
|
||||
NotNull<KernelGraphPtr> graph);
|
||||
|
|
Loading…
Reference in New Issue