support remove kernel_meta

mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_kernel_compile.cc

@@ -567,13 +567,13 @@ void TbeKernelCompileManager::GenKernelMod(const std::vector<CNodePtr> &node_lis
   MS_LOG(INFO) << "Gen kernel mod end!";
 }
 
-void TbeKernelCompileManager::LoadPreBuildResult(const std::vector<CNodePtr> &nodes) {
-  // save prebuild result as fusion_type, output_data_desc
+void TbeKernelCompileManager::UpdateFusionTypeAndOutputDataDesc(const std::vector<CNodePtr> &nodes) {
+  // save prebuild result: fusion_type, output_data_desc
   MS_LOG(INFO) << "Start update fusion type after pre build";
   for (auto &node : nodes) {
     MS_EXCEPTION_IF_NULL(node);
     auto full_name = node->fullname_with_scope();
-    auto kernel_name = full_name_to_json_name_[full_name];
+    auto kernel_name = pre_build_full_name_to_json_name_[full_name];
     auto pre_res = prebuild_res_map_[kernel_name];
     auto fusion_type = pre_res.fusion_type;
     auto output_data_desc = pre_res.output_data_desc;
@@ -607,13 +607,11 @@ std::string TbeKernelCompileManager::ParseSelectAndCheckResult(const nlohmann::j
     if (res != kFullySupported) {
       PrintProcessLog(json, WARNING);
     }
-  } else {
-    if (json.at(kStatus) == kFailed) {
-      auto all_logs = GetJsonValue<std::vector<nlohmann::json>>(json, kProcessInfo);
-      auto except_msg = FilterExceptionMessage(all_logs);
-      MS_LOG(EXCEPTION) << job_type << " running failed, op: " << json_name << "\nexception message:" << except_msg
-                        << trace::DumpSourceLines(node);
-    }
+  } else if (json.at(kStatus) == kFailed) {
+    auto all_logs = GetJsonValue<std::vector<nlohmann::json>>(json, kProcessInfo);
+    auto except_msg = FilterExceptionMessage(all_logs);
+    MS_LOG(EXCEPTION) << job_type << " running failed, op: " << json_name << "\nexception message:" << except_msg
+                      << trace::DumpSourceLines(node);
   }
   MS_LOG(DEBUG) << json_name << " " << job_type << " success, get: " << res;
   return res;
@@ -633,8 +631,48 @@ JsonNameMap TbeKernelCompileManager::GetAllSuccessFusion() {
   return success_fusion_ops_;
 }
 
+void TbeKernelCompileManager::DistributePreBuildTask(const std::vector<CNodePtr> &node_list) {
+  auto json_creator = std::make_shared<BuildTbeJsonCreator>();
+  MS_EXCEPTION_IF_NULL(json_creator);
+  nlohmann::json kernel_json;  // for gen json
+  nlohmann::json build_json;   // for assemble json
+  for (const auto &node : node_list) {
+    MS_EXCEPTION_IF_NULL(node);
+    kernel_json.clear();
+    build_json.clear();
+    if (!json_creator->GenJson(node, &kernel_json)) {
+      MS_LOG(EXCEPTION) << "Generate pre build json failed, op: " << node->fullname_with_scope()
+                        << trace::DumpSourceLines(node);
+    }
+    auto json_name = json_creator->GetJsonName();
+    auto full_name = node->fullname_with_scope();
+    pre_build_full_name_to_json_name_[full_name] = json_name;  // cache kernel name
+    if (TbeUtils::IsOneOf(pre_build_single_processed_kernels_, json_name)) {
+      // same op skip prebuild
+      continue;
+    }
+    pre_build_single_processed_kernels_.insert(json_name);
+    JsonAssemble(kPreCompile, kernel_json, &build_json);
+    auto task_id = GetJsonValue<int>(build_json, kJobId);
+    auto is_dynamic = AnfAlgo::IsDynamicShape(node);
+    SaveTaskInfo(is_dynamic, build_json, json_name, full_name, task_id, INT64_MAX);
+
+    // save pair<task_id, node> for exception print and get node trace
+    (void)job_id_to_node_.insert(std::pair<int, CNodePtr>(task_id, node));
+    // start compile
+    auto build_result = DispatchCompileTask(build_json);
+    auto json_obj = TurnStrToJson(build_result);
+    // print message of build
+    PrintCompileResult(json_obj);
+  }
+}
+
 void TbeKernelCompileManager::DistributeCompileTask(const std::vector<CNodePtr> &node_list,
                                                     const std::string &job_type) {
+  if (job_type == kPreCompile) {
+    DistributePreBuildTask(node_list);
+    return;
+  }
   auto json_creator = std::make_shared<BuildTbeJsonCreator>();
   MS_EXCEPTION_IF_NULL(json_creator);
   // for gen json
@@ -644,39 +682,34 @@ void TbeKernelCompileManager::DistributeCompileTask(const std::vector<CNodePtr>
   for (const auto &node : node_list) {
     MS_EXCEPTION_IF_NULL(node);
     kernel_json.clear();
-    auto full_name = node->fullname_with_scope();
-    if ((full_name_to_json_name_.find(full_name) != full_name_to_json_name_.end()) &&
-        tbe::TbeUtils::SearchCache(full_name_to_json_name_[full_name], false, is_tune_flag_, op_debug_level_) !=
-          nullptr) {
-      continue;  // skip fusion ops
-    }
-    if (AnfAlgo::GetKernelMod(node) != nullptr && !is_tune_flag_) {
-      continue;  // kernel mode exist, no need build
+    build_json.clear();
+    if (AnfAlgo::HasNodeAttr(kAttrIsUBFusionOp, node) && AnfAlgo::GetNodeAttr<bool>(node, kAttrIsUBFusionOp)) {
+      // skip fusion op, if node has the attr: kAttrIsUBFusionOp, means already done fusion compile, can not do single
+      // op compile
+      continue;
     }
     if (!json_creator->GenJson(node, &kernel_json)) {
       MS_LOG(EXCEPTION) << "Generate compile json failed, [" << node->fullname_with_scope() << "]"
                         << trace::DumpSourceLines(node);
     }
     auto json_name = json_creator->GetJsonName();
+    auto full_name = node->fullname_with_scope();
     full_name_to_json_name_[full_name] = json_name;
     // save all task io size info for gen kernel mod
     SaveIOSizeInfo(kernel_json, json_name);
-    if (job_type == kCompile || job_type == kTune) {
-      if (single_processed_kernels_.find(json_name) != single_processed_kernels_.end()) {
-        continue;  // same op no need build
-      }
-      (void)single_processed_kernels_.insert(json_name);
-      if (tbe::TbeUtils::SearchCache(json_name, false, is_tune_flag_, op_debug_level_) != nullptr) {
-        continue;  // cache exist, no need build
-      }
+    if (tbe::TbeUtils::SearchCache(json_name, false) != nullptr && !is_need_rebuild_) {
+      // cache exist, no need compile
+      continue;
     }
-    build_json.clear();
+    if (TbeUtils::IsOneOf(single_processed_kernels_, json_name)) {
+      // same op only compile once
+      continue;
+    }
+    single_processed_kernels_.insert(json_name);
     JsonAssemble(job_type, kernel_json, &build_json);
     // save compile json to file; cache io size for gen kernel mod
-    if (job_type == kCompile || job_type == kTune) {
-      auto build_str = build_json.dump(indent);
-      TbeUtils::SaveJsonInfo(json_name, build_str);
-    }
+    auto build_str = build_json.dump(indent);
+    TbeUtils::SaveJsonInfo(json_name, build_str);
     auto task_id = GetJsonValue<int>(build_json, kJobId);
     auto is_dynamic = AnfAlgo::IsDynamicShape(node);
     SaveTaskInfo(is_dynamic, build_json, json_name, full_name, task_id, INT64_MAX);
@@ -703,7 +736,7 @@ void TbeKernelCompileManager::TbePreBuild(const KernelGraphPtr &kernel_graph) {
   }
   DistributeCompileTask(node_list, kPreCompile);
   Query(kPreCompile);
-  LoadPreBuildResult(node_list);
+  UpdateFusionTypeAndOutputDataDesc(node_list);
   (void)gettimeofday(&end_time, nullptr);
   const uint64_t kUSecondInSecond = 1000000;
   uint64_t cost = kUSecondInSecond * static_cast<uint64_t>(end_time.tv_sec - start_time.tv_sec);
@@ -730,6 +763,7 @@ JsonNameMap TbeKernelCompileManager::TbeFusionOpCompile(const std::vector<Fusion
   nlohmann::json build_json;
   for (const auto &fusion_scope_iter : fusion_scopes) {
     fusion_op.clear();
+    build_json.clear();
     if (!json_creator->GenJson(fusion_scope_iter, &fusion_op)) {
       MS_LOG(WARNING) << "Generate fusion json failed, fusion info: " << fusion_scope_iter.full_name;
       continue;
@@ -737,16 +771,18 @@ JsonNameMap TbeKernelCompileManager::TbeFusionOpCompile(const std::vector<Fusion
     auto full_name = fusion_scope_iter.full_name;
     auto json_name = json_creator->GetJsonName();
     full_name_to_json_name_[full_name] = json_name;
+    // save all fusion ops to filter those compile succeed ops
     all_fusion_ops_[fusion_scope_iter.scope_id] = full_name;
     SaveIOSizeInfo(fusion_op, json_name, fusion_scope_iter.output_nodes);
-    if (fusion_processed_kernels_.find(json_name) != fusion_processed_kernels_.end()) {
-      continue;  // same op no need build
+    if (tbe::TbeUtils::SearchCache(json_name, false) != nullptr && !is_need_rebuild_) {
+      // cache exist, no need compile
+      continue;
     }
-    (void)fusion_processed_kernels_.insert(json_name);
-    if (tbe::TbeUtils::SearchCache(json_name, false, is_tune_flag_, op_debug_level_) != nullptr) {
-      continue;  // cache exist, no need build
+    if (TbeUtils::IsOneOf(fusion_processed_kernels_, json_name)) {
+      // same fusion op only compile once
+      continue;
     }
-    build_json.clear();
+    fusion_processed_kernels_.insert(json_name);
     JsonAssemble(job_type, fusion_op, &build_json);
     auto build_str = build_json.dump(indent);
     MS_LOG(DEBUG) << "FusionOp build json file : " << build_str;
@@ -811,6 +847,7 @@ void TbeKernelCompileManager::TbeInitialize() {
   auto auto_tiling_mode = (init_json[kJobContent][kSocInfo]["autoTilingMode"]).get<std::string>();
   tbe_init_flag_ = true;
   is_tune_flag_ = offline_tune || (auto_tiling_mode != "NO_TUNE");
+  is_need_rebuild_ = is_tune_flag_ || TbeUtils::IsOneOf({"1", "2", "4"}, op_debug_level_);
 
   auto init_str = init_json.dump();
   MS_LOG(INFO) << "TbeInitialize json file : " << init_str;
@@ -843,11 +880,14 @@ void TbeKernelCompileManager::ClearOldTask() {
   full_name_to_json_name_.clear();
   single_processed_kernels_.clear();
   fusion_processed_kernels_.clear();
+  pre_build_full_name_to_json_name_.clear();
+  pre_build_single_processed_kernels_.clear();
 }
 
 TbeKernelCompileManager::~TbeKernelCompileManager() { TbeFinalize(); }
 bool TbeKernelCompileManager::tbe_init_flag_ = false;
 bool TbeKernelCompileManager::is_tune_flag_ = false;
+bool TbeKernelCompileManager::is_need_rebuild_ = false;
 }  // namespace ascend
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_kernel_compile.h

@@ -112,20 +112,25 @@ class TbeKernelCompileManager {
   void SaveIOSizeInfo(const nlohmann::json &json, const std::string &json_name,
                       const std::vector<AnfNodePtr> &output_nodes = {});
   void ClearOldTask();
-  void LoadPreBuildResult(const std::vector<CNodePtr> &nodes);
+  void UpdateFusionTypeAndOutputDataDesc(const std::vector<CNodePtr> &nodes);
   JsonNameMap GetAllSuccessFusion();
   void GenKernelMod(const std::vector<CNodePtr> &anf_nodes);
   void DistributeCompileTask(const std::vector<CNodePtr> &anf_nodes, const std::string &job_type);
+  void DistributePreBuildTask(const std::vector<CNodePtr> &node_list);
 
   // init flag
   static bool tbe_init_flag_;
   // tune flag
   static bool is_tune_flag_;
+  // need rebuild when is_tune_flag_ is true, or op_debug_level_ is one of [1, 2, 4]
+  static bool is_need_rebuild_;
   // single op had build
   std::set<std::string> single_processed_kernels_;
+  // single op had pre build
+  std::set<std::string> pre_build_single_processed_kernels_;
   // fusion op had build
   std::set<std::string> fusion_processed_kernels_;
-  // if op_debug_level is "1", skip tbe compile cache and rebuild again.
+  // if op_debug_level is one of [1, 2, 4], skip tbe compile cache and rebuild again.
   std::string op_debug_level_;
   // id_node pair for node trace
   std::map<int, CNodePtr> job_id_to_node_;
@@ -133,6 +138,8 @@ class TbeKernelCompileManager {
   std::map<int, TaskInfo> task_map_;
   // pre build result
   std::map<std::string, PreBuildResult> prebuild_res_map_;
+  // using full_name to find json_name when update fusion type and out data desc
+  std::map<std::string, std::string> pre_build_full_name_to_json_name_;
   // save io size for kernel mod
   std::map<std::string, KernelIOSizeInfo> kernel_io_size_info_;
   // using full_name to find json_name when gen kernel mod

mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_utils.cc

@@ -113,13 +113,15 @@ std::string TbeUtils::GetOpDebugPath() {
 }
 
 std::string GetOpDebugLevel() {
-  const std::set<std::string> exp = {"0", "1", "2", "3", "4"};
-  std::string op_debug_level = "0";
+  static const std::set<int> value_ranges = {OP_DEBUG_LEVEL_0, OP_DEBUG_LEVEL_1, OP_DEBUG_LEVEL_2, OP_DEBUG_LEVEL_3,
+                                             OP_DEBUG_LEVEL_4};
+  std::string op_debug_level = std::to_string(OP_DEBUG_LEVEL_3);
   auto env_level = common::GetEnv(kCOMPILER_OP_LEVEL);
   if (!env_level.empty()) {
-    if (exp.find(env_level) == exp.end()) {
+    if (!TbeUtils::IsOneOf(value_ranges, std::atoi(env_level.c_str()))) {
       MS_LOG(WARNING) << "Invalid environment variable '" << kCOMPILER_OP_LEVEL << "': " << env_level
-                      << ", the value should be in {0, 1, 2, 3, 4}, now using the default value 0";
+                      << ", the value should be in [0, 1, 2, 3, 4], now using the default value 3."
+                         "Get more detail info at https://www.mindspore.cn/docs/note/zh-CN/master/env_var_list.html";
     } else {
       op_debug_level = env_level;
     }
@@ -215,13 +217,8 @@ void TbeUtils::UpdateCache(const std::string &kernel_name) {
   return bin_map->UpdateCache(kernel_name);
 }
 
-KernelPackPtr TbeUtils::SearchCache(const std::string &kernel_name, const bool is_akg, const bool is_tune,
-                                    const std::string op_debug_level) {
+KernelPackPtr TbeUtils::SearchCache(const std::string &kernel_name, const bool is_akg) {
   // search cache.
-  if (is_tune || op_debug_level != "0") {
-    return nullptr;
-  }
-
   KernelMeta *bin_map = KernelMeta::GetInstance();
   if (bin_map == nullptr) {
     MS_LOG(INFO) << "kernel cache is invalid.";

mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_utils.h

@@ -17,6 +17,7 @@
 #ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_TBE_TBE_UTILS_H_
 #define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_TBE_TBE_UTILS_H_
 #include <string>
+#include <set>
 #include <memory>
 #include <vector>
 #include <utility>
@@ -34,6 +35,14 @@ namespace kernel {
 namespace tbe {
 using std::string;
 using std::vector;
+enum OpDebugLevel {
+  OP_DEBUG_LEVEL_0,  // 0: turn off op debug, remove kernel_meta
+  OP_DEBUG_LEVEL_1,  // 1: turn on op debug, gen cce file
+  OP_DEBUG_LEVEL_2,  // 2: turn on op debug, gen cce file, turn off op compile optimization
+  OP_DEBUG_LEVEL_3,  // 3: turn off op debug, keep kernel_meta
+  OP_DEBUG_LEVEL_4   // 4: turn off op debug, gen _compute.json file
+};
+
 class TbeUtils {
  public:
   TbeUtils() = default;
@@ -62,12 +71,21 @@ class TbeUtils {
 
   static bool CheckOfflineTune();
 
-  static KernelPackPtr SearchCache(const std::string &kernel_name, const bool is_akg = false,
-                                   const bool is_tune = false, const std::string op_debug_level = "0");
+  static KernelPackPtr SearchCache(const std::string &kernel_name, const bool is_akg = false);
 
   static KernelPackPtr InsertCache(const std::string &kernel_name, const std::string &processor,
                                    const bool is_akg = false);
   static void UpdateCache(const std::string &kernel_name);
+
+  // check target value is one of the candidates
+  template <typename T>
+  static bool IsOneOf(const std::set<T> &candidates, const T &val) {
+    if (candidates.empty()) {
+      return false;
+    }
+    auto iter = candidates.find(val);
+    return iter != candidates.end();
+  }
 };
 
 struct KernelMetaInfo {

mindspore/ccsrc/backend/optimizer/ascend/buffer_fusion/ub_pattern_fusion.cc

@@ -530,6 +530,7 @@ bool UbPatternFusion::ReplaceFusionOp(mindspore::HashMap<int64_t, BufferFusionIn
     return false;
   }
   AnfAlgo::SetOutputInferTypeAndShape(types, shapes, buffer_fusion.get());
+  AnfAlgo::SetNodeAttr(kAttrIsUBFusionOp, MakeValue(true), buffer_fusion);
   SetFusionOpRefInfos(kernel_graph, buffer_fusion_info.outputs_list, buffer_fusion);
   ReplaceOldNode(buffer_fusion_infos, fusion_id, buffer_fusion, kernel_graph);
   return true;

mindspore/ccsrc/utils/utils.h

@@ -509,6 +509,7 @@ constexpr auto kAttrInputToAttrName = "input_to_attr_name";
 constexpr auto kAttrFuncType = "func_type";
 constexpr auto kAttrCustAicpu = "cust_aicpu";
 constexpr auto kAttrIsInternalOutputNopNode = "is_internal_output_nop_node";
+constexpr auto kAttrIsUBFusionOp = "is_ub_fusion_op";
 
 // custom operator func type
 constexpr auto kCustomTypeAOT = "aot";

mindspore/python/mindspore/_extends/parallel_compile/tbe_compiler/tbe_adapter.py

@@ -15,6 +15,7 @@
 """tbe adapter to adapt te/topi/auto-tune python api """
 import json
 import os
+import shutil
 import sys
 import traceback
 from datetime import datetime
@@ -102,6 +103,22 @@ def _cann_kb_unload(job: TbeJob):
     return res
 
 
+def _remove_cache(job: TbeJob):
+    """
+    :param job: remove cache file:[*.json, *.o, *.info, *.cce] when "op_debug_level" is "0"
+                op_debug_level: representation the env MS_COMPILER_OP_LEVEL
+    :return:
+    """
+    op_debug_level = job.content["SocInfo"]["op_debug_level"]
+    op_debug_dir = job.content["SocInfo"]["op_debug_dir"]
+    if op_debug_level != "0":
+        return
+    root_path = os.path.abspath(op_debug_dir)
+    if os.path.exists(root_path):
+        real_path = os.path.join(root_path, "kernel_meta/")
+        shutil.rmtree(real_path)
+
+
 def __directory_creation(path, concat_path):
     """
     Create directory
@@ -663,4 +680,5 @@ def tbe_finalize(auto_tiling_mode, offline_tune, job: TbeJob):
         job.error("Cann kb unload failed")
         return False
     clear_fusion_params()
+    _remove_cache(job)
     return True