diff --git a/mindspore/ccsrc/runtime/framework/graph_compiler.cc b/mindspore/ccsrc/runtime/framework/graph_compiler.cc index 39516a32376..8670298f4a5 100644 --- a/mindspore/ccsrc/runtime/framework/graph_compiler.cc +++ b/mindspore/ccsrc/runtime/framework/graph_compiler.cc @@ -319,6 +319,8 @@ GraphId GraphCompiler::CompileGraphImpl(const KernelGraphPtr &graph, const Devic // 'KernelMod' is real executive object of kernel. device_context->CreateKernel(graph->execution_order()); + device_context->PreprocessBeforeRunGraph(graph); + if (ms_context->get_param(MS_CTX_EXECUTION_MODE) == kGraphMode) { // Create device address for all anf nodes of graph. CreateDeviceAddress(graph, device_context); @@ -376,6 +378,8 @@ GraphId GraphCompiler::CompileGraph(const session::OpRunInfo &op_run_info, const // Generate 'KernelMod' for kernel in graph. device_context->CreateKernel(graph->execution_order()); + device_context->PreprocessBeforeRunSingleOpGraph(graph); + // Create device address for all anf nodes of graph. CreateDeviceAddress(graph, device_context); diff --git a/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.cc b/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.cc index eea6ed71ed8..8af9b65addf 100644 --- a/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.cc +++ b/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.cc @@ -88,11 +88,6 @@ void CPUDeviceContext::OptimizeGraph(const KernelGraphPtr &graph) const { // Run final optimization. opt::CommonFinalOptimization(graph); - - // Remove reorder after PS feature finish adapting push/pull in auto_monad. - auto execution_order = graph->execution_order(); - AnfAlgo::ReorderPosteriorExecList(NOT_NULL(&execution_order)); - graph->set_execution_order(execution_order); } void CPUDeviceContext::OptimizeSingleOpGraph(const KernelGraphPtr &graph) const { @@ -104,9 +99,7 @@ void CPUDeviceContext::OptimizeSingleOpGraph(const KernelGraphPtr &graph) const void CPUDeviceContext::OptimizeGraphImpl(const KernelGraphPtr &graph) const { auto optimizer = std::make_shared(); auto pm = std::make_shared(); - pm->AddPass(std::make_shared("insert_cast_cpu")); pm->AddPass(std::make_shared("insert_format_transform_op_cpu")); - pm->AddPass(std::make_shared()); optimizer->AddPassManager(pm); (void)optimizer->Optimize(graph); graph->SetExecOrderByDefault(); @@ -142,6 +135,30 @@ void CPUDeviceContext::CreateKernel(const std::vector &nodes) const { } } +namespace { +void ProcessCast(const KernelGraphPtr &graph) { + auto optimizer = std::make_shared(); + auto pm = std::make_shared(); + pm->AddPass(std::make_shared("insert_cast_cpu")); + MS_LOG(INFO) << "Insert cast pass"; + pm->AddPass(std::make_shared()); + optimizer->AddPassManager(pm); + (void)optimizer->Optimize(graph); + graph->SetExecOrderByDefault(); +} +} // namespace + +void CPUDeviceContext::PreprocessBeforeRunGraph(const KernelGraphPtr &graph) const { + ProcessCast(graph); + + // Remove reorder after PS feature finish adapting push/pull in auto_monad. + auto execution_order = graph->execution_order(); + AnfAlgo::ReorderPosteriorExecList(NOT_NULL(&execution_order)); + graph->set_execution_order(execution_order); +} + +void CPUDeviceContext::PreprocessBeforeRunSingleOpGraph(const KernelGraphPtr &graph) const { ProcessCast(graph); } + bool CPUDeviceContext::LaunchKernel(const CNodePtr &kernel, const std::vector &inputs, const std::vector &workspace, const std::vector &outputs, bool) const { diff --git a/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.h b/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.h index b666c06f2b7..187effd8b83 100644 --- a/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.h +++ b/mindspore/ccsrc/runtime/hardware/cpu/cpu_device_context.h @@ -48,6 +48,9 @@ class CPUDeviceContext : public DeviceContext { void SetOperatorInfo(const std::vector &nodes) const override; void CreateKernel(const std::vector &nodes) const override; + void PreprocessBeforeRunGraph(const KernelGraphPtr &graph) const override; + void PreprocessBeforeRunSingleOpGraph(const KernelGraphPtr &graph) const override; + bool LaunchKernel(const CNodePtr &kernel, const std::vector &inputs, const std::vector &workspace, const std::vector &outputs, bool is_dynamic_shape = false) const override; diff --git a/mindspore/ccsrc/runtime/hardware/device_context.h b/mindspore/ccsrc/runtime/hardware/device_context.h index 94a8a684669..3998dd80cc1 100644 --- a/mindspore/ccsrc/runtime/hardware/device_context.h +++ b/mindspore/ccsrc/runtime/hardware/device_context.h @@ -86,6 +86,11 @@ class DeviceContext { // 'KernelMod' is real executive object of kernel. virtual void CreateKernel(const std::vector &nodes) const = 0; + // Adjust kernel graph before run graph, used in Graph Mode. + virtual void PreprocessBeforeRunGraph(const KernelGraphPtr &graph) const {} + // Adjust single op kernel graph before run graph, used in PyNative Mode. + virtual void PreprocessBeforeRunSingleOpGraph(const KernelGraphPtr &graph) const {} + // Infer kernel shape and update abstract info for dynamic shape kernel. virtual void UpdateDynamicShape(const CNodePtr &kernel) const { AnfAlgo::InferShape(kernel); }