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 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_) {

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) {
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) {

View File

@ -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";