!28256 [GraphKernel] decouple update_state_formatter from bachend.

Merge pull request !28256 from chenlei_autodiff/gk_lite

mindspore/ccsrc/backend/optimizer/graph_kernel/add_stitch_atomic_clean_gpu.cc

@@ -203,7 +203,7 @@ bool StitchAtomicCleanInsertter::Run(const FuncGraphPtr &func_graph) {
   }
 
   if (changed) {
-    UpdateMng(mng, func_graph);
+    GkUtils::UpdateFuncGraphManager(mng, func_graph);
   }
 
   return changed;

mindspore/ccsrc/backend/optimizer/graph_kernel/core/graph_kernel_utils.cc

@@ -203,4 +203,19 @@ FuncGraphPtr GkUtils::LiteGraph2AnfGraph(const inner::LiteGraphPtr &lite_graph)
   }
   return func_graph;
 }
+
+FuncGraphManagerPtr GkUtils::GetFuncGraphManager(const FuncGraphPtr &func_graph) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  FuncGraphManagerPtr manager = func_graph->manager();
+  if (manager == nullptr) {
+    manager = Manage(func_graph, true);
+    func_graph->set_manager(manager);
+  }
+  return manager;
+}
+
+void GkUtils::UpdateFuncGraphManager(const FuncGraphManagerPtr &mng, const FuncGraphPtr &func_graph) {
+  mng->RemoveRoots();
+  mng->KeepRoots({func_graph});
+}
 }  // namespace mindspore::graphkernel

mindspore/ccsrc/backend/optimizer/graph_kernel/core/graph_kernel_utils.h

@@ -85,6 +85,16 @@ class GkUtils {
    * @brief Change lite graph to anf graph.
    */
   static FuncGraphPtr LiteGraph2AnfGraph(const inner::LiteGraphPtr &lite_graph);
+
+  /**
+   * @brief Get manager of func graph. If there is no manager, a new one will be created.
+   */
+  static FuncGraphManagerPtr GetFuncGraphManager(const FuncGraphPtr &func_graph);
+
+  /**
+   * @brief Update func graph manager. Do nothing if the manager is nullptr.
+   */
+  static void UpdateFuncGraphManager(const FuncGraphManagerPtr &mng, const FuncGraphPtr &func_graph);
 };
 }  // namespace mindspore::graphkernel
 #endif  // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_CORE_GRAPH_KERNEL_UTILS_H_

mindspore/ccsrc/backend/optimizer/graph_kernel/update_state_formatter.cc → mindspore/ccsrc/backend/optimizer/graph_kernel/core/update_state_formatter.cc

@@ -13,17 +13,16 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "backend/optimizer/graph_kernel/update_state_formatter.h"
+#include "backend/optimizer/graph_kernel/core/update_state_formatter.h"
 
 #include <vector>
 #include <set>
 #include <memory>
 #include <utility>
 #include <algorithm>
-#include "backend/session/anf_runtime_algorithm.h"
-#include "backend/kernel_compiler/common_utils.h"
-#include "backend/optimizer/graph_kernel/graph_kernel_helper.h"
+#include "ir/anf.h"
 #include "backend/optimizer/graph_kernel/core/graph_kernel_utils.h"
+#include "backend/optimizer/graph_kernel/core/graph_kernel_callback.h"
 #include "backend/optimizer/graph_kernel/core/eliminate_redundant_output.h"
 
 namespace mindspore::graphkernel {
@@ -52,7 +51,7 @@ AnfNodePtrList SpreadUpdateState::ExtendInputsOfUpdateState(const AnfNodePtrList
         auto tuple_getitem = func_graph->NewCNode({NewValueNode(prim::kPrimTupleGetItem), node, idx});
         MS_EXCEPTION_IF_NULL(tuple_getitem);
         tuple_getitem->set_abstract(node_abstract[i]);
-        tuple_getitem->set_kernel_info(std::make_shared<device::KernelInfo>());
+        Callback::Instance()->SetEmptyKernelInfo(tuple_getitem);
         result.push_back(tuple_getitem);
       }
     } else {
@@ -103,12 +102,12 @@ bool ShrinkUpdateState::Run(const FuncGraphPtr &func_graph) {
     mt_inputs.insert(mt_inputs.begin(), NewValueNode(prim::kPrimMakeTuple));
     auto mt_node = func_graph->NewCNode(mt_inputs);
     mt_node->set_abstract(std::make_shared<abstract::AbstractTuple>(abs_list));
-    mt_node->set_kernel_info(std::make_shared<device::KernelInfo>());
+    Callback::Instance()->SetEmptyKernelInfo(mt_node);
 
     AnfNodePtrList inputs = {cnode->input(0), cnode->input(1), mt_node};
     auto new_node = func_graph->NewCNode(inputs);
     new_node->set_abstract(node->abstract());
-    new_node->set_kernel_info(std::make_shared<device::KernelInfo>());
+    Callback::Instance()->SetEmptyKernelInfo(new_node);
     (void)mng->Replace(node, new_node);
     changed = true;
   }
@@ -125,7 +124,7 @@ bool ExtendOutputForUpdateState::Run(const FuncGraphPtr &func_graph) {
     if (getitems_.empty()) continue;
     FindIndexesToUpdateState(mng);
     if (indexes_.empty()) continue;
-    auto sub_func_graph = AnfAlgo::GetCNodeFuncGraphPtr(node);
+    auto sub_func_graph = GetCNodeFuncGraph(node);
     FilterIndexes(sub_func_graph);
     if (indexes_.empty()) continue;
     for (auto idx : indexes_) {
@@ -133,7 +132,7 @@ bool ExtendOutputForUpdateState::Run(const FuncGraphPtr &func_graph) {
     }
   }
   if (changed) {
-    UpdateMng(mng, func_graph);
+    GkUtils::UpdateFuncGraphManager(mng, func_graph);
     std::make_shared<SpreadUpdateState>()->Run(func_graph);
     std::make_shared<EliminateHangingOutput>()->Run(func_graph);
   }

mindspore/ccsrc/backend/optimizer/graph_kernel/update_state_formatter.h → mindspore/ccsrc/backend/optimizer/graph_kernel/core/update_state_formatter.h

@@ -14,8 +14,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#ifndef MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_UPDATE_STATE_FORMATTER_H_
-#define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_UPDATE_STATE_FORMATTER_H_
+#ifndef MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_CORE_UPDATE_STATE_FORMATTER_H_
+#define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_CORE_UPDATE_STATE_FORMATTER_H_
 
 #include <vector>
 #include "backend/optimizer/common/pass.h"
@@ -145,4 +145,4 @@ class MergeOutputForUpdateState : public opt::Pass {
   bool Run(const FuncGraphPtr &func_graph) override;
 };
 }  // namespace mindspore::graphkernel
-#endif  // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_UPDATE_STATE_FORMATTER_H_
+#endif  // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_CORE_UPDATE_STATE_FORMATTER_H_

mindspore/ccsrc/backend/optimizer/graph_kernel/decrease_compute_precision.cc

@@ -24,6 +24,7 @@
 #include "backend/optimizer/ascend/ascend_helper.h"
 #include "backend/session/anf_runtime_algorithm.h"
 #include "backend/optimizer/graph_kernel/graph_kernel_helper.h"
+#include "backend/optimizer/graph_kernel/core/graph_kernel_utils.h"
 #include "runtime/device/kernel_info.h"
 #include "backend/kernel_compiler/kernel_build_info.h"
 #include "backend/kernel_compiler/common_utils.h"
@@ -259,7 +260,7 @@ bool DecreaseComputePrecision::Run(const FuncGraphPtr &func_graph) {
     }
   }
   if (changed) {
-    UpdateMng(mng, func_graph);
+    GkUtils::UpdateFuncGraphManager(mng, func_graph);
   }
   return changed;
 }

mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.cc

@@ -495,19 +495,4 @@ inner::LiteGraphPtr AnfGraph2LiteGraph(const FuncGraphPtr &func_graph) {
   }
   return gb.Get();
 }
-
-FuncGraphManagerPtr GetFuncGraphManager(const FuncGraphPtr &func_graph) {
-  MS_EXCEPTION_IF_NULL(func_graph);
-  FuncGraphManagerPtr manager = func_graph->manager();
-  if (manager == nullptr) {
-    manager = Manage(func_graph, true);
-    func_graph->set_manager(manager);
-  }
-  return manager;
-}
-
-void UpdateMng(const FuncGraphManagerPtr &mng, const FuncGraphPtr &func_graph) {
-  mng->RemoveRoots();
-  mng->KeepRoots({func_graph});
-}
 }  // namespace mindspore::graphkernel

mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_helper.h

@@ -117,10 +117,5 @@ AbstractBasePtr GetOutputAbstract(const AnfNodePtr &node, size_t output_idx);
 
 // functions to graphkernel model
 inner::LiteGraphPtr AnfGraph2LiteGraph(const FuncGraphPtr &func_graph);
-
-// return a func_graph's manager
-FuncGraphManagerPtr GetFuncGraphManager(const FuncGraphPtr &func_graph);
-
-void UpdateMng(const FuncGraphManagerPtr &mng, const FuncGraphPtr &func_graph);
 }  // namespace mindspore::graphkernel
 #endif  // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_GRAPH_KERNEL_HELPER_H_

mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_optimization.cc

@@ -39,7 +39,7 @@
 #include "backend/optimizer/graph_kernel/optimize_assign.h"
 #include "backend/optimizer/graph_kernel/split_umonad.h"
 #include "backend/optimizer/graph_kernel/reorder_ops.h"
-#include "backend/optimizer/graph_kernel/update_state_formatter.h"
+#include "backend/optimizer/graph_kernel/core/update_state_formatter.h"
 #include "backend/optimizer/graph_kernel/axis_normalizer.h"
 #include "backend/optimizer/graph_kernel/decrease_compute_precision.h"
 #include "backend/optimizer/graph_kernel/decrease_transfer_precision.h"

mindspore/ccsrc/backend/optimizer/graph_kernel/parallel_fusion.cc

@@ -27,7 +27,7 @@
 #include "backend/kernel_compiler/common_utils.h"
 #include "frontend/operator/ops.h"
 #include "ir/func_graph_cloner.h"
-#include "backend/optimizer/graph_kernel/update_state_formatter.h"
+#include "backend/optimizer/graph_kernel/core/update_state_formatter.h"
 #include "backend/optimizer/graph_kernel/core/graph_builder.h"
 
 namespace mindspore::graphkernel {

mindspore/ccsrc/backend/optimizer/graph_kernel/raise_reduction_precision.cc

@@ -19,6 +19,7 @@
 #include "base/core_ops.h"
 #include "utils/utils.h"
 #include "backend/optimizer/common/helper.h"
+#include "backend/optimizer/graph_kernel/core/graph_kernel_utils.h"
 #include "backend/optimizer/graph_kernel/graph_kernel_helper.h"
 #include "backend/session/anf_runtime_algorithm.h"
 #include "ir/tensor.h"
@@ -120,7 +121,7 @@ bool RaiseReductionPrecision::Run(const FuncGraphPtr &func_graph) {
     }
   }
   if (changed) {
-    UpdateMng(mng, func_graph);
+    GkUtils::UpdateFuncGraphManager(mng, func_graph);
   }
   return changed;
 }

mindspore/ccsrc/backend/optimizer/graph_kernel/reorder_ops.cc

@@ -24,7 +24,7 @@
 #include "utils/log_adapter.h"
 #include "backend/session/anf_runtime_algorithm.h"
 #include "debug/anf_ir_dump.h"
-#include "backend/optimizer/graph_kernel/graph_kernel_helper.h"
+#include "backend/optimizer/graph_kernel/core/graph_kernel_utils.h"
 
 namespace mindspore::graphkernel {
 namespace {
@@ -326,7 +326,7 @@ bool ReorderOps::ReorderCastTypeInsensitive(const FuncGraphPtr &func_graph) {
   // Limitation: Assuming the type insensitive node will not change the type of input nodes, otherwise it can be seen
   //   as another cast node in some sense, such as LessEqual operator, which performs on two inputs and output a
   //   a boolean result.
-  auto mng = GetFuncGraphManager(func_graph);
+  auto mng = GkUtils::GetFuncGraphManager(func_graph);
   bool changed = false;
   auto todos = TopoSort(func_graph->get_return());
   for (const auto &anf_node : todos) {

mindspore/ccsrc/backend/optimizer/graph_kernel/tsa_atomic_add_to_first_tensor.cc

@@ -286,7 +286,7 @@ bool TsaAtomicAddToFirstTensor::Run(const FuncGraphPtr &func_graph) {
   }
 
   if (changed) {
-    UpdateMng(mng, func_graph);
+    GkUtils::UpdateFuncGraphManager(mng, func_graph);
   }
 
   return changed;