diff --git a/mindspore/_extends/graph_kernel/expanders/__init__.py b/mindspore/_extends/graph_kernel/expanders/__init__.py index fb229621c89..bef9989cd6b 100644 --- a/mindspore/_extends/graph_kernel/expanders/__init__.py +++ b/mindspore/_extends/graph_kernel/expanders/__init__.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ -"""expanders init""" +"""expanders init. Deprecated, please add the new operators in the c++ file""" + from .addn import AddN from .assign_add import AssignAdd diff --git a/mindspore/_extends/graph_kernel/model/graph_split.py b/mindspore/_extends/graph_kernel/model/graph_split.py index af92662e5df..cf83042af4c 100644 --- a/mindspore/_extends/graph_kernel/model/graph_split.py +++ b/mindspore/_extends/graph_kernel/model/graph_split.py @@ -1198,8 +1198,8 @@ class GraphSplitCpu(GraphSplitByPattern): def get_default_mode(self, op): """Get default mode in CPU""" - pattern = PrimLib.iter_type(op) - return self.Area.MODE_BASIC if pattern == PrimLib.RESHAPE else self.Area.MODE_COMPOSITE + del op + return self.Area.MODE_COMPOSITE def pattern_fuse(self, fuse_func=None): """fuse Areas by pattern""" diff --git a/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt b/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt index 8f5a0c1f030..3dbe1a48665 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt +++ b/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt @@ -114,6 +114,7 @@ if(ENABLE_AKG AND ${CMAKE_SYSTEM_NAME} MATCHES "Linux") list(APPEND AKG_SRC_LIST ${AKG_D_SRC_LIST}) endif() if(ENABLE_CPU) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp") file(GLOB_RECURSE AKG_CPU_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "akg/cpu/*.cc" ) diff --git a/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.cc b/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.cc index 5cbed01b01d..078236eca12 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.cc @@ -17,6 +17,8 @@ #include "backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.h" #include +#include +#include #include #include #include @@ -28,24 +30,29 @@ namespace mindspore { namespace kernel { -namespace { -using AkgParallelLambda = int (*)(int task_id, int num_task, void *cdata); -int AkgLaunchFunc(AkgParallelLambda flambda, void *cdata, int num_task) { - size_t num_workers = - std::min(mindspore::common::ThreadPool::GetInstance().GetSyncRunThreadNum(), static_cast(num_task)); - std::vector tasks; - size_t thread_index = 0; - while (thread_index < num_workers) { - auto block = [&, thread_index]() { - flambda(thread_index, num_workers, cdata); - return mindspore::common::SUCCESS; - }; - tasks.emplace_back(block); - thread_index++; +class AkgParallelLaunch { + public: + using AkgParallelLambda = int (*)(int task_id, int num_task, void *cdata); + static size_t num_workers; + static int AkgLaunchFunc(AkgParallelLambda flambda, void *cdata, int num_task) { +#pragma omp parallel num_threads(num_workers) + { flambda(omp_get_thread_num(), num_workers, cdata); } + return 0; } - mindspore::common::ThreadPool::GetInstance().SyncRun(tasks); - return 0; -} + // the GetFunc should be called only once. + static void *GetFunc() { + const char *omp_num_threads = getenv("OMP_NUM_THREADS"); + if (omp_num_threads != nullptr) { + auto env_thread = std::stoi(std::string(omp_num_threads)); + if (env_thread > 0) { + AkgParallelLaunch::num_workers = static_cast(env_thread); + } + } + MS_LOG(INFO) << "AKG threads is : " << AkgParallelLaunch::num_workers; + return reinterpret_cast(&AkgParallelLaunch::AkgLaunchFunc); + } +}; +size_t AkgParallelLaunch::num_workers = 1; struct AkgCallBack { void *parallel_launch_func; @@ -53,13 +60,13 @@ struct AkgCallBack { void (*free_func)(void *); AkgCallBack() { - parallel_launch_func = reinterpret_cast(&AkgLaunchFunc); + parallel_launch_func = AkgParallelLaunch::GetFunc(); malloc_func = &malloc; free_func = &free; } ~AkgCallBack() = default; }; -} // namespace + CpuKernelManagerPtr CpuKernelMod::kernelmanager_ = std::make_shared(); CpuKernelManager::~CpuKernelManager() { @@ -120,14 +127,16 @@ void *CpuKernelManager::GetFunction(const std::string &kernel_name) { return launch_func; } +CpuKernelMod::CpuKernelMod(const KernelPackPtr &kp) { + auto js = nlohmann::json::parse(kp->GetJson()->contents, kp->GetJson()->contents + kp->GetJson()->len); + kernel_name_ = js["kernelName"]; + launch_func_ = kernelmanager_->GetFunction(kernel_name_); +} + bool CpuKernelMod::Launch(const std::vector &inputs, const std::vector &, const std::vector &outputs, void *stream_ptr) { - auto js = nlohmann::json::parse(kernel_pack_->GetJson()->contents, - kernel_pack_->GetJson()->contents + kernel_pack_->GetJson()->len); - std::string kernel_name = js["kernelName"]; - auto launch_func = kernelmanager_->GetFunction(kernel_name); - if (launch_func == nullptr) { - MS_LOG(ERROR) << "GetFunction failed. kernel: " << kernel_name; + if (launch_func_ == nullptr) { + MS_LOG(ERROR) << "GetFunction failed. kernel: " << kernel_name_; return false; } std::vector runtimeargs; @@ -135,10 +144,10 @@ bool CpuKernelMod::Launch(const std::vector &inputs, const std::vect [](const AddressPtr &input) -> void * { return input->addr; }); (void)std::transform(std::begin(outputs), std::end(outputs), std::back_inserter(runtimeargs), [](const AddressPtr &output) -> void * { return output->addr; }); - AkgCallBack akg_callback; + static AkgCallBack akg_callback; runtimeargs.emplace_back(reinterpret_cast(&akg_callback)); using AkgCpuKernelFunction = void (*)(void *); - reinterpret_cast(launch_func)(reinterpret_cast(runtimeargs.data())); + reinterpret_cast(launch_func_)(reinterpret_cast(runtimeargs.data())); return true; } } // namespace kernel diff --git a/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.h b/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.h index 71cf82405f9..8c834fd8b28 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.h +++ b/mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.h @@ -43,10 +43,9 @@ class CpuKernelManager { mutable std::shared_mutex mutex_; }; using CpuKernelManagerPtr = std::shared_ptr; - class CpuKernelMod : public KernelMod { public: - explicit CpuKernelMod(const KernelPackPtr &kp) : kernel_pack_(kp) {} + explicit CpuKernelMod(const KernelPackPtr &kp); ~CpuKernelMod() = default; void SetInputSizeList(const std::vector &size_list) { input_size_list_ = size_list; } @@ -61,10 +60,11 @@ class CpuKernelMod : public KernelMod { static CpuKernelManagerPtr kernelmanager_; private: - KernelPackPtr kernel_pack_; std::vector input_size_list_; std::vector output_size_list_; std::vector workspace_size_list_; // workspace is not used in cpu kernel. + void *launch_func_; + std::string kernel_name_; }; using CpuKernelModPtr = std::shared_ptr; diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/adam.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/adam.cc new file mode 100644 index 00000000000..8eadb3a2b83 --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/adam.cc @@ -0,0 +1,99 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "backend/optimizer/graph_kernel/expanders/expander_factory.h" +#include "ir/dtype.h" + +namespace mindspore::graphkernel::expanders { +class Adam : public OpExpander { + public: + Adam() { + std::initializer_list attrs{"use_nesterov"}; + (void)validators_.emplace_back(std::make_unique(attrs)); + } + ~Adam() = default; + + protected: + bool CheckInputs() override { + const auto &var = inputs_info_[0]; + if (var.type != kNumberTypeFloat32 && var.type != kNumberTypeFloat16) { + MS_LOG(INFO) << "In Adam, var's dtype must be float16 or float32"; + return false; + } + return true; + } + + NodePtrList Expand() override { + const auto &inputs = gb.Get()->inputs(); + const auto &var = inputs[0]; + const auto &m = inputs[1]; + const auto &v = inputs[2]; + const auto &beta1_power = inputs[3]; + const auto &beta2_power = inputs[4]; + const auto &lr = inputs[5]; + const auto &beta1 = inputs[6]; + const auto &beta2 = inputs[7]; + const auto &epsilon = inputs[8]; + const auto &grad = inputs[9]; + + // m_new <- beta1 * m + (1 - beta1) * grad + auto m_b = gb.Emit("Mul", {beta1, m}); + tensor::TensorPtr data = std::make_shared(static_cast(1.0), TypeIdToType(var->type)); + auto const_one = gb.Value(data); + auto m1_beta1 = gb.Emit("Sub", {const_one, beta1}); + auto m_g = gb.Emit("Mul", {m1_beta1, grad}); + auto m_new = gb.Emit("Add", {m_b, m_g}); + + // v_new <- beta2 * v + (1 - beta2) * grad * grad + auto v_b = gb.Emit("Mul", {beta2, v}); + auto m1_beta2 = gb.Emit("Sub", {const_one, beta2}); + auto grad_mul = gb.Emit("Mul", {grad, grad}); + auto v_g = gb.Emit("Mul", {m1_beta2, grad_mul}); + auto v_new = gb.Emit("Add", {v_b, v_g}); + // lr_t <- lr * sqrt(1 - beta2_power) / (1 - beta1_power); + auto m1_beta2_power = gb.Emit("Sub", {const_one, beta2_power}); + auto m1_beta2_power_sqrt = gb.Emit("Sqrt", {m1_beta2_power}); + auto m1_beta1_power = gb.Emit("Sub", {const_one, beta1_power}); + auto power_div = gb.Emit("RealDiv", {m1_beta2_power_sqrt, m1_beta1_power}); + auto lr_t = gb.Emit("Mul", {lr, power_div}); + + // if use_nesterov: var_new <- var - lr_t * (m_new * beta1 + (1 - beta1) * grad) / (epsilon + sqrt(v_new)) + // if not use_nesterov: var_new <- var - lr_t * m_new / (epsilon + sqrt(v_new)) + auto v_new_sqrt = gb.Emit("Sqrt", {v_new}); + auto v_new_sqrt_e = gb.Emit("Add", {epsilon, v_new_sqrt}); + auto lr_t_div = gb.Emit("RealDiv", {lr_t, v_new_sqrt_e}); + mindspore::graphkernel::inner::NodePtr var_sub; + if (GetValue(attrs_["use_nesterov"])) { + auto m_new_mul = gb.Emit("Mul", {m_new, beta1}); + auto m_new_mul_add = gb.Emit("Add", {m_new_mul, m_g}); + var_sub = gb.Emit("Mul", {lr_t_div, m_new_mul_add}); + } else { + var_sub = gb.Emit("Mul", {lr_t_div, m_new}); + } + + auto var_new = gb.Emit("Sub", {var, var_sub}); + auto var_result = gb.Emit("InplaceAssign", {var, var_new, var_new}, {{"fake_output", MakeValue(true)}}); + auto m_result = gb.Emit("InplaceAssign", {m, m_new, var_result}, {{"fake_output", MakeValue(true)}}); + auto v_result = gb.Emit("InplaceAssign", {v, v_new, m_result}, {{"fake_output", MakeValue(true)}}); + auto result = {var_result, m_result, v_result}; + return result; + } +}; +OP_EXPANDER_REGISTER("Adam", Adam); +} // namespace mindspore::graphkernel::expanders diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/softplus.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/softplus.cc new file mode 100644 index 00000000000..ddafb53d1e3 --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/softplus.cc @@ -0,0 +1,51 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "backend/optimizer/graph_kernel/expanders/expander_factory.h" +#include "ir/dtype.h" + +namespace mindspore::graphkernel::expanders { +class Softplus : public OpExpander { + public: + Softplus() {} + ~Softplus() = default; + + protected: + bool CheckInputs() override { + const auto &input_x = inputs_info_[0]; + if (input_x.type != kNumberTypeFloat32 && input_x.type != kNumberTypeFloat16) { + MS_LOG(INFO) << "In Softplus, input_x's dtype must be float16 or float32"; + return false; + } + return true; + } + + NodePtrList Expand() override { + const auto &inputs = gb.Get()->inputs(); + const auto &input_x = inputs[0]; + auto exp_x = gb.Emit("Exp", {input_x}); + tensor::TensorPtr data = std::make_shared(static_cast(1.0), TypeIdToType(input_x->type)); + auto const_one = gb.Value(data); + auto exp_x_add_one = gb.Emit("Add", {exp_x, const_one}); + auto result = gb.Emit("Log", {exp_x_add_one}); + return {result}; + } +}; +OP_EXPANDER_REGISTER("Softplus", Softplus); +} // namespace mindspore::graphkernel::expanders diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/softplus_grad.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/softplus_grad.cc new file mode 100644 index 00000000000..61fc9e28782 --- /dev/null +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/softplus_grad.cc @@ -0,0 +1,53 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "backend/optimizer/graph_kernel/expanders/expander_factory.h" +#include "ir/dtype.h" + +namespace mindspore::graphkernel::expanders { +class SoftplusGrad : public OpExpander { + public: + SoftplusGrad() {} + ~SoftplusGrad() = default; + + protected: + bool CheckInputs() override { + const auto &input_x = inputs_info_[1]; + if (input_x.type != kNumberTypeFloat32 && input_x.type != kNumberTypeFloat16) { + MS_LOG(INFO) << "In SoftplusGrad, input_x's dtype must be float16 or float32"; + return false; + } + return true; + } + + NodePtrList Expand() override { + const auto &inputs = gb.Get()->inputs(); + const auto &input_dy = inputs[0]; + const auto &input_x = inputs[1]; + auto exp_x = gb.Emit("Exp", {input_x}); + tensor::TensorPtr data = std::make_shared(static_cast(1.0), TypeIdToType(input_x->type)); + auto const_one = gb.Value(data); + auto exp_x_add_one = gb.Emit("Add", {exp_x, const_one}); + auto dy_mul_exp_x = gb.Emit("Mul", {input_dy, exp_x}); + auto result = gb.Emit("RealDiv", {dy_mul_exp_x, exp_x_add_one}); + return {result}; + } +}; +OP_EXPANDER_REGISTER("SoftplusGrad", SoftplusGrad); +} // namespace mindspore::graphkernel::expanders diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/standardnormal.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/standardnormal.cc index 6f7a6115f7c..784763481b8 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/standardnormal.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/expanders/standardnormal.cc @@ -27,6 +27,8 @@ class StandardNormal : public OpExpander { (void)validators_.emplace_back(std::make_unique(attrs)); } ~StandardNormal() {} + + protected: NodePtrList Expand() override { const auto &inputs = gb.Get()->inputs(); const auto &input_x = inputs[0]; diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_cluster.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_cluster.cc index 7d62de1ca5f..20715d82b9d 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_cluster.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_cluster.cc @@ -100,6 +100,13 @@ std::vector GraphKernelCluster::GetClusterableOpList() { {kGPUDevice, OpLevel_0, prim::kPrimSign}, {kGPUDevice, OpLevel_0, prim::kPrimSin}, {kGPUDevice, OpLevel_0, prim::kPrimStridedSlice}, + // cpu + {kCPUDevice, OpLevel_0, prim::kPrimLogicalNot}, + {kCPUDevice, OpLevel_0, prim::kPrimMod}, + {kCPUDevice, OpLevel_1, prim::kPrimReduceMax}, + {kCPUDevice, OpLevel_0, prim::kPrimSelect}, + {kCPUDevice, OpLevel_0, prim::kPrimLess}, + {kCPUDevice, OpLevel_0, prim::kPrimLessEqual}, }; const auto &flags = GraphKernelFlags::GetInstance(); return GkUtils::GetValidOps(clusterable_ops_with_level, flags.fusion_ops_level, flags.enable_cluster_ops_only, diff --git a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_expander.cc b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_expander.cc index c96d62865f3..6a020746265 100644 --- a/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_expander.cc +++ b/mindspore/ccsrc/backend/optimizer/graph_kernel/graph_kernel_expander.cc @@ -46,6 +46,7 @@ constexpr size_t kAssignInputIdx = 1; constexpr size_t kLambOptimizerInputIdx = 12; constexpr size_t kLambWeightInputIdx = 4; constexpr size_t kRandomInputIdx = 1; +constexpr size_t kAdamInputIdx = 10; std::vector GetExpandOps() { std::vector expand_ops_with_level = { @@ -94,6 +95,13 @@ std::vector GetExpandOps() { {kGPUDevice, OpLevel_0, prim::kPrimIdentityMath}, {kGPUDevice, OpLevel_0, prim::kPrimOnesLike}, {kGPUDevice, OpLevel_0, prim::kPrimStandardNormal}, + {kCPUDevice, OpLevel_0, prim::kPrimOnesLike}, + {kCPUDevice, OpLevel_0, prim::kPrimBiasAdd}, + {kCPUDevice, OpLevel_1, prim::kPrimBiasAddGrad}, + {kCPUDevice, OpLevel_0, prim::kPrimRelu}, + {kCPUDevice, OpLevel_1, prim::kPrimMaximumGrad}, + {kCPUDevice, OpLevel_1, prim::kPrimMinimumGrad}, + {kCPUDevice, OpLevel_1, prim::kPrimAdam}, }; const auto &flags = GraphKernelFlags::GetInstance(); return GkUtils::GetValidOps(expand_ops_with_level, flags.fusion_ops_level, flags.enable_expand_ops_only, @@ -199,6 +207,7 @@ ExpanderPtr GraphKernelExpander::GetExpander(const AnfNodePtr &node) { {prim::kLambApplyOptimizerAssign, std::make_shared(kLambOptimizerInputIdx)}, {prim::kLambApplyWeightAssign, std::make_shared(kLambWeightInputIdx)}, {prim::kPrimStandardNormal, std::make_shared(kRandomInputIdx)}, + {prim::kPrimAdam, std::make_shared(kAdamInputIdx)}, }; for (auto &e : expanders) { diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 3ec0de251e2..70151e170e1 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -204,6 +204,8 @@ list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/backend/optimizer/ list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/backend/optimizer/graph_kernel/lite_adapter/callback_impl.cc") list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/backend/kernel_compiler/tbe/tbe_kernel_compile.cc") +list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_mod.cc") +list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/backend/kernel_compiler/akg/cpu/akg_cpu_kernel_build.cc") if(ENABLE_SECURITY) list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/profiler/device/profiling.cc") list(REMOVE_ITEM MINDSPORE_SRC_LIST "../../../mindspore/ccsrc/profiler/device/ascend/memory_profiling.cc")