fix bug of dynamic shape

This commit is contained in:
caifubi 2022-01-08 14:18:04 +08:00
parent dbd99eb804
commit c0e47202b9
3 changed files with 17 additions and 8 deletions

View File

@ -695,6 +695,9 @@ void TbeKernelCompileManager::DistributeCompileTask(const std::vector<CNodePtr>
auto json_name = json_creator->GetJsonName(); auto json_name = json_creator->GetJsonName();
auto full_name = node->fullname_with_scope(); auto full_name = node->fullname_with_scope();
full_name_to_json_name_[full_name] = json_name; 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 // save all task io size info for gen kernel mod
SaveIOSizeInfo(kernel_json, json_name); SaveIOSizeInfo(kernel_json, json_name);
if (tbe::TbeUtils::SearchCache(json_name, false) != nullptr && !is_need_rebuild_) { if (tbe::TbeUtils::SearchCache(json_name, false) != nullptr && !is_need_rebuild_) {

View File

@ -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) { void TbeUtils::GetCompileInfo(const AnfNodePtr &node, std::string *compile_info, bool *get_flag) {
MS_EXCEPTION_IF_NULL(node); MS_EXCEPTION_IF_NULL(node);
MS_LOG(INFO) << "Get compile info from json file start. [" << node->fullname_with_scope() << "]"; MS_LOG(INFO) << "Get compile info from json file start. [" << node->fullname_with_scope() << "]";
auto json_creator = std::make_shared<kernel::BuildTbeJsonCreator>(); std::string json_name;
MS_EXCEPTION_IF_NULL(json_creator); if (AnfAlgo::HasNodeAttr(kAttrJsonFileName, node->cast<CNodePtr>())) {
nlohmann::json kernel_json; json_name = AnfAlgo::GetNodeAttr<std::string>(node, kAttrJsonFileName);
if (!json_creator->GenJson(node, &kernel_json)) { } else {
MS_LOG(WARNING) << "Gen kernel json failed [" << node->fullname_with_scope() << "]"; auto json_creator = std::make_shared<kernel::BuildTbeJsonCreator>();
*get_flag = false; MS_EXCEPTION_IF_NULL(json_creator);
return; 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(); auto config_path = TbeUtils::GetOpDebugPath();
std::string path = config_path + kCceKernelMeta + json_name + kJsonSuffix; std::string path = config_path + kCceKernelMeta + json_name + kJsonSuffix;
if (path.size() > PATH_MAX) { if (path.size() > PATH_MAX) {

View File

@ -513,6 +513,7 @@ constexpr auto kAttrCustAicpu = "cust_aicpu";
constexpr auto kAttrIsInternalOutputNopNode = "is_internal_output_nop_node"; constexpr auto kAttrIsInternalOutputNopNode = "is_internal_output_nop_node";
constexpr auto kAttrIsUBFusionOp = "is_ub_fusion_op"; constexpr auto kAttrIsUBFusionOp = "is_ub_fusion_op";
constexpr auto kAttrMicro = "micro"; constexpr auto kAttrMicro = "micro";
constexpr auto kAttrJsonFileName = "json_file_name";
// custom operator func type // custom operator func type
constexpr auto kCustomTypeAOT = "aot"; constexpr auto kCustomTypeAOT = "aot";