From c0e47202b956cedbe39eca35c7cd6c73189d811a Mon Sep 17 00:00:00 2001
From: caifubi <caifubi1@huawei.com>
Date: Sat, 8 Jan 2022 14:18:04 +0800
Subject: [PATCH] fix bug of dynamic shape

---
 .../kernel_compiler/tbe/tbe_kernel_compile.cc |  3 +++
 .../backend/kernel_compiler/tbe/tbe_utils.cc  | 21 ++++++++++++-------
 mindspore/ccsrc/utils/utils.h                 |  1 +
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_kernel_compile.cc b/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_kernel_compile.cc
index b74b5cfc184..c0abf4a0817 100644
--- a/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_kernel_compile.cc
+++ b/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_kernel_compile.cc
@@ -695,6 +695,9 @@ void TbeKernelCompileManager::DistributeCompileTask(const std::vector<CNodePtr>
     auto json_name = json_creator->GetJsonName();
     auto full_name = node->fullname_with_scope();
     full_name_to_json_name_[full_name] = json_name;
+    if (AnfAlgo::IsDynamicShape(node)) {
+      AnfAlgo::SetNodeAttr(kAttrJsonFileName, MakeValue(json_name), node);
+    }
     // save all task io size info for gen kernel mod
     SaveIOSizeInfo(kernel_json, json_name);
     if (tbe::TbeUtils::SearchCache(json_name, false) != nullptr && !is_need_rebuild_) {
diff --git a/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_utils.cc b/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_utils.cc
index d29dbbaae2f..7bd3bcd6b9f 100644
--- a/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_utils.cc
+++ b/mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_utils.cc
@@ -376,15 +376,20 @@ bool KernelMeta::ReadIndex(const std::string &bin_dir) {
 void TbeUtils::GetCompileInfo(const AnfNodePtr &node, std::string *compile_info, bool *get_flag) {
   MS_EXCEPTION_IF_NULL(node);
   MS_LOG(INFO) << "Get compile info from json file start. [" << node->fullname_with_scope() << "]";
-  auto json_creator = std::make_shared<kernel::BuildTbeJsonCreator>();
-  MS_EXCEPTION_IF_NULL(json_creator);
-  nlohmann::json kernel_json;
-  if (!json_creator->GenJson(node, &kernel_json)) {
-    MS_LOG(WARNING) << "Gen kernel json failed [" << node->fullname_with_scope() << "]";
-    *get_flag = false;
-    return;
+  std::string json_name;
+  if (AnfAlgo::HasNodeAttr(kAttrJsonFileName, node->cast<CNodePtr>())) {
+    json_name = AnfAlgo::GetNodeAttr<std::string>(node, kAttrJsonFileName);
+  } else {
+    auto json_creator = std::make_shared<kernel::BuildTbeJsonCreator>();
+    MS_EXCEPTION_IF_NULL(json_creator);
+    nlohmann::json kernel_json;
+    if (!json_creator->GenJson(node, &kernel_json)) {
+      MS_LOG(WARNING) << "Gen kernel json failed [" << node->fullname_with_scope() << "]";
+      *get_flag = false;
+      return;
+    }
+    json_name = json_creator->GetJsonName();
   }
-  auto json_name = json_creator->GetJsonName();
   auto config_path = TbeUtils::GetOpDebugPath();
   std::string path = config_path + kCceKernelMeta + json_name + kJsonSuffix;
   if (path.size() > PATH_MAX) {
diff --git a/mindspore/ccsrc/utils/utils.h b/mindspore/ccsrc/utils/utils.h
index 2be4b569d45..d33fec33d21 100644
--- a/mindspore/ccsrc/utils/utils.h
+++ b/mindspore/ccsrc/utils/utils.h
@@ -513,6 +513,7 @@ constexpr auto kAttrCustAicpu = "cust_aicpu";
 constexpr auto kAttrIsInternalOutputNopNode = "is_internal_output_nop_node";
 constexpr auto kAttrIsUBFusionOp = "is_ub_fusion_op";
 constexpr auto kAttrMicro = "micro";
+constexpr auto kAttrJsonFileName = "json_file_name";
 
 // custom operator func type
 constexpr auto kCustomTypeAOT = "aot";