environ bug fix in control flow

mindspore/ccsrc/backend/kernel_compiler/cpu/environ/environ_cpu_get.cc

@@ -33,6 +33,7 @@ void EnvironGetCPUKernel::InitKernel(const CNodePtr &node) {
   }
 
   value_type_attr_ = TypeId(AnfAlgo::GetNodeAttr<int>(node, kEnvValueTypeAttr));
+  MS_LOG(INFO) << "The EnvironGet kernel " << node->fullname_with_scope() << " value type: " << value_type_attr_;
   handle_size_ = sizeof(int64_t);
   key_size_ = sizeof(int64_t);
 
@@ -67,6 +68,9 @@ bool EnvironGetCPUKernel::Launch(const std::vector<AddressPtr> &inputs, const st
 
   // Get env and value by handle and key.
   const auto &env = EnvironMgr::GetInstance().Get(host_handle);
+  if (env == nullptr) {
+    MS_LOG(EXCEPTION) << "Get the env failed, handle: " << host_handle << ", key: " << host_key;
+  }
   MS_EXCEPTION_IF_NULL(env);
   const auto &env_value = env->Get(host_key);
   // Default value.

mindspore/ccsrc/backend/kernel_compiler/cpu/environ/environ_cpu_set.cc

@@ -41,6 +41,7 @@ void EnvironSetCPUKernel::InitKernel(const CNodePtr &node) {
   }
 
   value_type_attr_ = TypeId(AnfAlgo::GetNodeAttr<int>(node, kEnvValueTypeAttr));
+  MS_LOG(INFO) << "The EnvironSet kernel " << node->fullname_with_scope() << " value type: " << value_type_attr_;
   handle_size_ = sizeof(int64_t);
   key_size_ = sizeof(int64_t);
 
@@ -63,9 +64,6 @@ bool EnvironSetCPUKernel::Launch(const std::vector<AddressPtr> &inputs, const st
   auto input_key = GetDeviceAddress<int64_t>(inputs, 1);
   auto input_value = GetDeviceAddress<void>(inputs, 2);
   auto output_handle = GetDeviceAddress<int64_t>(outputs, 0);
-  if (input_handle != output_handle) {
-    MS_LOG(EXCEPTION) << "The EnvSet is ref kernel and the output handle is not equal of input handle.";
-  }
 
   // Get host handle and host key.
   int64_t host_handle = input_handle[0];
@@ -81,10 +79,16 @@ bool EnvironSetCPUKernel::Launch(const std::vector<AddressPtr> &inputs, const st
 
   // Set env member.
   const auto &env = EnvironMgr::GetInstance().Get(host_handle);
-  MS_EXCEPTION_IF_NULL(env);
+  if (env == nullptr) {
+    MS_LOG(EXCEPTION) << "Get the env failed, handle: " << host_handle << ", key: " << host_key;
+  }
   auto env_value = std::make_shared<EnvironValue>(value_ptr, value_size_, value_type_attr_, kGPUDevice);
   env->Set(host_key, env_value);
 
+  // Set output handle.
+  output_handle[0] = input_handle[0];
+  MS_LOG(DEBUG) << "Get output handle: " << output_handle[0];
+
   return true;
 }
 }  // namespace kernel

mindspore/ccsrc/backend/kernel_compiler/environ_manager.cc

@@ -39,21 +39,26 @@ int64_t EnvironMgr::Create() {
 
 EnvironPtr EnvironMgr::Get(int64_t handle) {
   mutex.lock_shared();
-  if (envs_.count(handle) > 0) {
-    return envs_[handle];
-  } else {
-    return nullptr;
+  const auto &envIter = envs_.find(handle);
+  if (envIter != envs_.end()) {
+    auto &result = envIter->second;
+    mutex.unlock_shared();
+    return result;
   }
-  mutex.unlock();
+
+  mutex.unlock_shared();
+  return nullptr;
 }
 
 void EnvironMgr::Clear() {
+  mutex.lock();
   for (auto &env : envs_) {
     MS_EXCEPTION_IF_NULL(env.second);
     env.second->Clear();
   }
 
   envs_.clear();
+  mutex.unlock();
 }
 
 bool EnvironMgr::CheckEnvInput(const CNodePtr &kernel_node) {
@@ -99,17 +104,15 @@ bool EnvironMgr::IsScalarTensor(TypeId type, std::vector<size_t> shape) {
     return false;
   }
 
-  if (shape.size() != kScalarTensorShapeDim) {
-    MS_LOG(ERROR) << "The shape size is invalid: " << shape.size();
-    return false;
+  if (shape.empty()) {
+    return true;
   }
 
-  if (shape[0] != kScalarTensorShapeSize) {
-    MS_LOG(ERROR) << "The shape is invalid: " << shape[0];
-    return false;
+  if ((shape.size() == kScalarTensorShapeDim) && (shape[0] == kScalarTensorShapeSize)) {
+    return true;
   }
 
-  return true;
+  return false;
 }
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/gpu/environ/environ_gpu_get.cc

@@ -34,6 +34,7 @@ bool EnvironGetGpuKernel::Init(const CNodePtr &kernel_node) {
   }
 
   value_type_attr_ = TypeId(AnfAlgo::GetNodeAttr<int>(kernel_node, kEnvValueTypeAttr));
+  MS_LOG(INFO) << "The EnvironGet kernel " << kernel_node->fullname_with_scope() << " value type: " << value_type_attr_;
   handle_size_ = sizeof(int64_t);
   key_size_ = sizeof(int64_t);
 
@@ -84,7 +85,9 @@ bool EnvironGetGpuKernel::Launch(const std::vector<AddressPtr> &inputs, const st
 
   // Get env and value by handle and key.
   const auto &env = EnvironMgr::GetInstance().Get(host_handle);
-  MS_EXCEPTION_IF_NULL(env);
+  if (env == nullptr) {
+    MS_LOG(EXCEPTION) << "Get the env failed, handle: " << host_handle << ", key: " << host_key;
+  }
   const auto &env_value = env->Get(host_key);
   // Default value.
   auto value = input_default_value;

mindspore/ccsrc/backend/kernel_compiler/gpu/environ/environ_gpu_set.cc

@@ -43,6 +43,7 @@ bool EnvironSetGpuKernel::Init(const CNodePtr &kernel_node) {
   }
 
   value_type_attr_ = TypeId(AnfAlgo::GetNodeAttr<int>(kernel_node, kEnvValueTypeAttr));
+  MS_LOG(INFO) << "The EnvironSet kernel " << kernel_node->fullname_with_scope() << " value type: " << value_type_attr_;
   handle_size_ = sizeof(int64_t);
   key_size_ = sizeof(int64_t);
 
@@ -70,9 +71,6 @@ bool EnvironSetGpuKernel::Launch(const std::vector<AddressPtr> &inputs, const st
   auto input_key = GetDeviceAddress<int64_t>(inputs, 1);
   auto input_value = GetDeviceAddress<void>(inputs, 2);
   auto output_handle = GetDeviceAddress<int64_t>(outputs, 0);
-  if (input_handle != output_handle) {
-    MS_LOG(EXCEPTION) << "The EnvSet is ref kernel and the output handle is not equal of input handle.";
-  }
 
   // Get host handle and host key.
   int64_t host_handle = 0;
@@ -98,10 +96,18 @@ bool EnvironSetGpuKernel::Launch(const std::vector<AddressPtr> &inputs, const st
 
   // Set env member.
   const auto &env = EnvironMgr::GetInstance().Get(host_handle);
-  MS_EXCEPTION_IF_NULL(env);
+  if (env == nullptr) {
+    MS_LOG(EXCEPTION) << "Get the env failed, handle: " << host_handle << ", key: " << host_key;
+  }
   auto env_value = std::make_shared<EnvironValue>(value_ptr, value_size_, value_type_attr_, kGPUDevice);
   env->Set(host_key, env_value);
 
+  // Copy output handle.
+  CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_,
+                             cudaMemcpyAsync(output_handle, input_handle, handle_size_, cudaMemcpyDeviceToDevice,
+                                             reinterpret_cast<cudaStream_t>(stream_ptr)),
+                             "Copy output handle failed.");
+
   return true;
 }
 }  // namespace kernel

mindspore/ccsrc/backend/session/kernel_graph.h

@@ -157,6 +157,7 @@ class KernelGraph : public FuncGraph {
   AnfNodePtr GetBackendAnfByFrontAnf(const AnfNodePtr &front_anf);
   // get front anf by backend anf
   AnfNodePtr GetFrontAnfByBackendAnf(const AnfNodePtr &backend_anf) const;
+  const mindspore::HashMap<AnfNodePtr, AnfNodePtr> &backend_front_anf_map() const { return backend_front_anf_map_; }
   // check backend node whether exist in map
   bool BackendNodeExistInFrontBackendMap(const AnfNodePtr &backend_anf);
   // get value node by tensor

mindspore/ccsrc/runtime/framework/graph_scheduler.cc

@@ -1835,7 +1835,19 @@ void GraphScheduler::DumpActor(const ActorSet *actor_set, const GraphCompilerInf
     return;
   }
 
-  std::string filename = GetSaveGraphsPathName("actor_set_" + actor_set->name_ + ".ir");
+  // Get the saved actor set name.
+  auto &kernel_graphs = graph_compiler_info.graphs_;
+  MS_EXCEPTION_IF_NULL(kernel_graphs.front());
+  auto first_graph_id = kernel_graphs.front()->graph_id();
+  MS_EXCEPTION_IF_NULL(kernel_graphs.back());
+  auto last_graph_id = kernel_graphs.back()->graph_id();
+  std::string strategy = (graph_compiler_info.strategy_ == GraphExecutionStrategy::kPipeline) ? "pipeline" : "step";
+  std::string save_name = "actor_set_" + strategy + "_kernel_graph_" + std::to_string(first_graph_id);
+  if (last_graph_id != first_graph_id) {
+    save_name = save_name + "-" + std::to_string(last_graph_id);
+  }
+
+  std::string filename = GetSaveGraphsPathName(save_name + ".ir");
   std::ofstream ofs(filename);
   if (!ofs.is_open()) {
     MS_LOG(ERROR) << "Open file [" << filename << "] failed!";
@@ -1912,6 +1924,14 @@ void GraphScheduler::DumpDeviceTensorStore(const GraphCompilerInfo &graph_compil
       }
     }
     ofs << "\n";
+
+    for (auto &backend_front_map : graph->backend_front_anf_map()) {
+      MS_EXCEPTION_IF_NULL(backend_front_map.first);
+      MS_EXCEPTION_IF_NULL(backend_front_map.second);
+      MS_LOG(DEBUG) << "Graph: " << graph->graph_id()
+                    << ", backend node: " << backend_front_map.first->fullname_with_scope()
+                    << ", front node: " << backend_front_map.second->DebugString();
+    }
   }
 }
 }  // namespace runtime

mindspore/ccsrc/runtime/hardware/device_context.h

@@ -27,10 +27,6 @@
 #include "backend/session/kernel_graph.h"
 #include "backend/session/anf_runtime_algorithm.h"
 #include "backend/optimizer/common/common_backend_optimization.h"
-#ifdef ENABLE_DUMP_IR
-#include "debug/anf_ir_dump.h"
-#include "debug/dump_proto.h"
-#endif
 
 namespace mindspore {
 namespace device {