diff --git a/akg b/akg index c9a2e4edfd1..751fc1693d0 160000 --- a/akg +++ b/akg @@ -1 +1 @@ -Subproject commit c9a2e4edfd1bc70f2498b0bc6228377c45a11951 +Subproject commit 751fc1693d06d732266c063867f4acfc79c6fe63 diff --git a/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.cc b/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.cc index 1acbc3aea59..cef4c826b46 100644 --- a/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.cc +++ b/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.cc @@ -282,6 +282,10 @@ void GraphKernelFlags::RegisterFlags(std::map *flag_ma // String flags reg.AddFlag("repository_path", &repository_path); + reg.AddFlag("target_os", &target_os); + reg.AddFlag("cpu_arch", &cpu_arch); + reg.AddFlag("cpu_feature", &cpu_feature); + reg.AddFlag("cpu_type", &cpu_type); // String list flags reg.AddFlag("enable_expand_ops", &enable_expand_ops); @@ -320,6 +324,10 @@ std::string GraphKernelFlags::DumpAllFlags() const { json["recompute_peak_threshold"] = recompute_peak_threshold; json["repository_path"] = repository_path; + json["target_os"] = target_os; + json["cpu_arch"] = cpu_arch; + json["cpu_feature"] = cpu_feature; + json["cpu_type"] = cpu_type; json["enable_expand_ops"] = enable_expand_ops; json["enable_expand_ops_only"] = enable_expand_ops_only; diff --git a/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.h b/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.h index 24e8d2e067c..9590aa4a154 100644 --- a/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.h +++ b/mindspore/ccsrc/common/graph_kernel/graph_kernel_flags.h @@ -169,6 +169,19 @@ class GraphKernelFlags { */ std::string repository_path; + /** + * Target info. + * These flags can be used for cross-compiling. Available when the device target is cpu. + * target_os: the operating system to run kernels. + * cpu_arch: the architecture, default value is related to the building environment (e.g. "arm" or "x86_64") + * cpu_feature: the instruction set to be used. (e.g. "avx" or "avx512") + * cpu_type: the cpu processor type. (e.g. "core-avx2" or "skylake-avx512") + */ + std::string target_os{"linux"}; + std::string cpu_arch; + std::string cpu_feature; + std::string cpu_type; + /** * Additional expanding operators (case sensitive). * The operators to be added into the default expanding operator list. diff --git a/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.cc b/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.cc index 5fe68dac06a..5765389e32f 100644 --- a/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.cc +++ b/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.cc @@ -26,6 +26,7 @@ #include "common/graph_kernel/core/graph_builder.h" #include "common/graph_kernel/core/graph_kernel_utils.h" #include "runtime/hardware/device_context_manager.h" +#include "common/graph_kernel/graph_kernel_flags.h" namespace mindspore::graphkernel { using kernel::OpAttr; @@ -1115,20 +1116,61 @@ bool AkgKernelJsonGenerator::CollectFusedJsonWithSingleKernel(const CNodePtr &c_ namespace { bool GetCpuInfo(nlohmann::json *target_info) { - (*target_info)[kJsonKeySystem] = "linux"; + const auto &flags = GraphKernelFlags::GetInstance(); + std::string target_os = flags.target_os; + std::string arch = flags.cpu_arch; + std::string feature = flags.cpu_feature; + std::string type = flags.cpu_type; + std::set valid_os = {"linux", "windows"}; + std::set valid_arch = {"arm", "x86_64"}; + // arch: <{supported-features}, default-feature> + std::map, std::string>> valid_features = { + {"arm", {{"neon"}, "neon"}}, + {"x86_64", {{"sse", "avx", "avx512"}, "avx"}}, + }; + std::set valid_cpu_types = {"core-avx2", "skylake-avx512", "core-avx-i", "haswell", "skylake"}; + + if (valid_os.count(target_os) == 0) { + MS_LOG(WARNING) << "GraphKernelFlag: unsupported \"target_os\": " << target_os; + target_os = "linux"; + } + if (valid_arch.count(arch) == 0) { + if (!arch.empty()) { + MS_LOG(WARNING) << "GraphKernelFlag: unsupported \"cpu_arch\": " << arch; + } #if defined(ENABLE_ARM) || defined(ENABLE_ARM64) - (*target_info)[kJsonKeyArch] = "arm"; - (*target_info)[kJsonKeyFeature] = "neon"; + arch = "arm"; #else - (*target_info)[kJsonKeyArch] = "x86"; -#ifdef ENABLE_AVX512 - (*target_info)[kJsonKeyFeature] = "avx512"; -#elif defined(ENABLE_AVX) - (*target_info)[kJsonKeyFeature] = "avx"; -#else - (*target_info)[kJsonKeyFeature] = "sse"; -#endif + arch = "x86_64"; #endif + } + + auto &features = valid_features[arch]; + if (features.first.count(feature) == 0) { + if (!feature.empty()) { + MS_LOG(WARNING) << "GraphKernelFlag: unsupported \"cpu_feature\": " << feature; + } + feature = features.second; + } + + if (valid_cpu_types.count(type) == 0) { + if (!type.empty()) { + MS_LOG(WARNING) << "GraphKernelFlag: unsupported \"cpu_type\": " << type; + type = ""; + } + if (feature == "avx512") { + type = "skylake-avx512"; + } else if (feature == "avx") { + type = "core-avx2"; + } + } + + (*target_info)[kJsonKeySystem] = target_os; + (*target_info)[kJsonKeyCpuArch] = arch; + (*target_info)[kJsonKeyCpuFeature] = feature; + if (!type.empty()) { + (*target_info)[kJsonKeyCpuType] = type; + } return true; } } // namespace diff --git a/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.h b/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.h index 69f4e3e7d96..01047f1fad0 100644 --- a/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.h +++ b/mindspore/ccsrc/kernel/akg/akg_kernel_json_generator.h @@ -68,8 +68,9 @@ constexpr auto kJsonKeyTargetInfo = "target_info"; constexpr auto kJsonKeyComputeCapability = "compute_capability"; constexpr auto kJsonKeySmCount = "sm_count"; constexpr auto kJsonKeySystem = "system"; -constexpr auto kJsonKeyArch = "arch"; -constexpr auto kJsonKeyFeature = "feature"; +constexpr auto kJsonKeyCpuArch = "arch"; +constexpr auto kJsonKeyCpuFeature = "feature"; +constexpr auto kJsonKeyCpuType = "cpu"; constexpr auto kJsonKeyNodeName = "node_name"; // dump option