diff --git a/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt b/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt index 5622013fa27..954402e5c9e 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt +++ b/mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt @@ -35,6 +35,7 @@ if(ENABLE_CPU) "cpu/fl/*.cc" "cpu/ps/*.cc" "cpu/quantum/*.cc" + "cpu/pyfunc/*.cc" ) if(NOT ENABLE_MPI) diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/cpu_kernel_factory.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/cpu_kernel_factory.cc index 12ae560be86..eee6e6f4985 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/cpu_kernel_factory.cc +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/cpu_kernel_factory.cc @@ -21,6 +21,7 @@ #include #include "runtime/device/kernel_info.h" +#include "runtime/device/cpu/kernel_select_cpu.h" namespace mindspore { namespace kernel { @@ -111,6 +112,11 @@ std::pair CPUKernelFactory::CPUKernelAttrCheck(const std::string & MS_LOG(INFO) << "Not registered CPU kernel: op[" << kernel_name << "]!"; return std::make_pair(false, 0); } + + if (device::cpu::IsDynamicParamKernel(kernel_name)) { + return std::make_pair(true, 0); + } + auto kernel_attrs = GetSupportedKernelAttrList(kernel_name); if (kernel_attrs[0].GetInputSize() == 0 && kernel_attrs[0].GetOutputSize() == 0) { auto op_info_ptr = mindspore::kernel::OpLib::FindOp(kernel_name, kernel::OpImplyType::kCPU); diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/pyfunc/py_func_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/pyfunc/py_func_cpu_kernel.cc new file mode 100644 index 00000000000..0d81da83a2e --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/pyfunc/py_func_cpu_kernel.cc @@ -0,0 +1,295 @@ +/** + * 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 "backend/kernel_compiler/cpu/pyfunc/py_func_cpu_kernel.h" + +#include +#include +#include "Eigen/Core" +#include "Eigen/src/Core/arch/CUDA/Half.h" +#include "abstract/utils.h" +#include "runtime/device/cpu/cpu_common.h" +#include "pybind_api/ir/tensor_py.h" +#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h" + +namespace mindspore { +namespace kernel { +namespace { +py::object RawMemoryToScalar(const void *data, const TypePtr &type) { + switch (type->type_id()) { + case kNumberTypeBool: + return py::bool_(*reinterpret_cast(data)); + case kNumberTypeInt16: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeUInt16: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeInt8: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeUInt8: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeInt32: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeUInt32: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeInt64: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeUInt64: + return py::int_(*reinterpret_cast(data)); + case kNumberTypeFloat16: { + const Eigen::half_impl::__half_raw data_half(*reinterpret_cast(data)); + return py::float_(Eigen::half_impl::half_to_float(data_half)); + } + case kNumberTypeFloat32: + return py::float_(*reinterpret_cast(data)); + case kNumberTypeFloat64: + return py::float_(*reinterpret_cast(data)); + default: + MS_LOG(EXCEPTION) << "Type: " << type->type_id() << " not supported."; + } +} + +void ScalarToRawMemory(const py::object &obj, const TypePtr &type, const AddressPtr &address) { + switch (type->type_id()) { + case kNumberTypeBool: { + bool data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(bool)), EOK, "memcpy failed."); + return; + } + // ref: pybind11-src/include/pybind11/pytypes.h + // py::int_ convert py::object to `long`, `unsigned long`, `long long`, `unsigned long long` with Python API + // according to typename T, and then convert to target data type with C style cast. + case kNumberTypeInt8: { + int8_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(int8_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeUInt8: { + uint8_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(uint8_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeInt16: { + int16_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(int16_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeUInt16: { + uint8_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(uint8_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeInt32: { + int32_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(int32_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeUInt32: { + uint32_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(uint32_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeInt64: { + int64_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(int64_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeUInt64: { + uint64_t data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(uint64_t)), EOK, "memcpy failed."); + return; + } + case kNumberTypeFloat16: { + float data = py::cast(obj); + Eigen::half_impl::__half_raw data_half = Eigen::half_impl::float_to_half_rtne(data); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data_half.x, sizeof(data_half.x)), EOK, + "memcpy failed."); + return; + } + case kNumberTypeFloat32: { + float data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(float)), EOK, "memcpy failed."); + return; + } + case kNumberTypeFloat64: { + float data = py::cast(obj); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, &data, sizeof(double)), EOK, "memcpy failed."); + return; + } + default: + MS_LOG(EXCEPTION) << "Type: " << type->type_id() << " not supported."; + } +} + +void ArrayToRawMemory(const py::array &array, const AddressPtr &address) { + if (static_cast(array.flags()) & pybind11::detail::npy_api::NPY_ARRAY_C_CONTIGUOUS_) { + const py::buffer_info &buf_info = array.request(); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, buf_info.ptr, buf_info.size), EOK, "memcpy failed."); + } else { + // Transform numpy array to row major buffer. + Py_buffer pybuf; + if (PyObject_GetBuffer(array.ptr(), &pybuf, PyBUF_ANY_CONTIGUOUS)) { + MS_LOG(EXCEPTION) << "Failed to get buffer from the input!"; + } + + auto buffer = std::make_unique(pybuf.len); + if (PyBuffer_ToContiguous(buffer.get(), &pybuf, pybuf.len, 'C')) { + PyBuffer_Release(&pybuf); + MS_LOG(EXCEPTION) << "Can't copy numpy.ndarray to a contiguous buffer."; + } + PyBuffer_Release(&pybuf); + CHECK_RET_WITH_EXCEPT(memcpy_s(address->addr, address->size, buffer.get(), pybuf.len), EOK, "memcpy failed."); + } +} + +void ObjectToRawMemory(const py::object &object, const PythonOjectType &object_type, const TypePtr &data_type, + const AddressPtr &address) { + switch (object_type) { + case PythonOjectType::kScalar: + return ScalarToRawMemory(object, data_type, address); + case PythonOjectType::kNumpyArray: + return ArrayToRawMemory(object.cast(), address); + default: + MS_LOG(EXCEPTION) << "python object not supported. type: " << object_type; + } +} + +py::tuple RawMemoryToPyObjects(const std::vector &inputs, const PyFuncArgumentInfo &input_infos, + const std::vector &input_tensors) { + py::tuple result(inputs.size()); + for (size_t i = 0; i < inputs.size(); i++) { + switch (input_infos.object_types[i]) { + case PythonOjectType::kScalar: + result[i] = RawMemoryToScalar(inputs[i]->addr, input_infos.dtypes[i]); + break; + case PythonOjectType::kNumpyArray: { + const tensor::TensorPtr &tensor = input_tensors[i]; + CHECK_RET_WITH_EXCEPT(memcpy_s(tensor->data_c(), tensor->Size(), inputs[i]->addr, inputs[i]->size), EOK, + "memcpy failed."); + result[i] = tensor::TensorPy::AsNumpy(*tensor); + break; + } + default: + MS_LOG(EXCEPTION) << "Python args not support. Index: " << i << ", type" << input_infos.object_types[i]; + } + } + return result; +} + +void PyObjectToRawMemorys(const py::object &object, const PyFuncArgumentInfo &output_infos, + const std::vector &outputs) { + // Single output. + if (!py::isinstance(object)) { + if (outputs.size() != 1) { + MS_LOG(EXCEPTION) << "The output num is 1, with " << outputs.size() << " expect."; + } + return ObjectToRawMemory(object, output_infos.object_types[0], output_infos.dtypes[0], outputs[0]); + } + + // Multiply outputs. + auto result_tuple = object.cast(); + if (result_tuple.size() != outputs.size()) { + MS_LOG(EXCEPTION) << "The output num is: " << result_tuple.size() << ", with " << outputs.size() << " expect."; + } + + for (size_t i = 0; i < outputs.size(); i++) { + ObjectToRawMemory(result_tuple[i], output_infos.object_types[i], output_infos.dtypes[i], outputs[i]); + } +} +} // namespace + +void PyFuncCpuKernel::InitKernel(const CNodePtr &kernel_node) { + func_id_ = AnfAlgo::GetNodeAttr(kernel_node, "fn_id"); + BuildFuncInfo(kernel_node); +} + +bool PyFuncCpuKernel::Launch(const std::vector &inputs, const std::vector &, + const std::vector &outputs) { + if (!init_) { + py_func_ = GetPythonFunc(func_id_); + init_ = true; + } + + return ExecuteKernel(inputs, outputs); +} + +void PyFuncCpuKernel::BuildFuncInfo(const CNodePtr &kernel_node) { + const auto &in_shapes = AnfAlgo::GetNodeAttr>>(kernel_node, "in_shapes"); + const auto &in_types = AnfAlgo::GetNodeAttr>(kernel_node, "in_types"); + const auto &out_shapes = AnfAlgo::GetNodeAttr>>(kernel_node, "out_shapes"); + const auto &out_types = AnfAlgo::GetNodeAttr>(kernel_node, "out_types"); + + input_infos_.dtypes = in_types; + input_infos_.shapes = in_shapes; + for (size_t i = 0; i < in_shapes.size(); i++) { + auto tensor = std::make_shared(in_types[i]->type_id(), in_shapes[i]); + input_tensors_.push_back(tensor); + + const auto &object_type = in_shapes[i].empty() ? PythonOjectType::kScalar : PythonOjectType::kNumpyArray; + input_infos_.object_types.emplace_back(object_type); + } + + output_infos_.dtypes = out_types; + output_infos_.shapes = out_shapes; + for (size_t j = 0; j < out_shapes.size(); j++) { + const auto &object_type = out_shapes[j].empty() ? PythonOjectType::kScalar : PythonOjectType::kNumpyArray; + output_infos_.object_types.emplace_back(object_type); + } +} + +bool PyFuncCpuKernel::ExecuteKernel(const std::vector &inputs, const std::vector &outputs) { + if (Py_IsInitialized() != true) { + MS_LOG(ERROR) << "Py_IsInitialized failed."; + return false; + } + + py::gil_scoped_acquire gil_acquire; + py::object result; + if (inputs.size()) { + py::tuple args = RawMemoryToPyObjects(inputs, input_infos_, input_tensors_); + result = py_func_(*args); + } else { + result = py_func_(); + } + + if (output_infos_.shapes.empty()) { + return true; + } + + PyObjectToRawMemorys(result, output_infos_, outputs); + return true; +} + +py::function PyFuncCpuKernel::GetPythonFunc(const int64_t &func_id) { + py::gil_scoped_acquire gil_acquire; + static const std::string &module_name = "mindspore.ops.operations.other_ops"; + static const std::string &func_name = "get_pyfunc"; + py::module module = py::module::import(module_name.c_str()); + py::object get_pyfunc_obj = module.attr(func_name.c_str()); + if (get_pyfunc_obj.is_none()) { + MS_LOG(EXCEPTION) << "Cannot find a python function named " << func_name << "in module" << module_name; + } + + py::function get_pyfunc = get_pyfunc_obj.cast(); + py::object py_func_obj = get_pyfunc(py::int_(func_id)); + if (py_func_obj.is_none()) { + MS_LOG(EXCEPTION) << "Cannot find python func with id: " << func_id; + } + + return py_func_obj.cast(); +} + +MS_REG_CPU_KERNEL(PyFunc, KernelAttr(), PyFuncCpuKernel) +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/pyfunc/py_func_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/pyfunc/py_func_cpu_kernel.h new file mode 100644 index 00000000000..71dd99f9ee1 --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/pyfunc/py_func_cpu_kernel.h @@ -0,0 +1,76 @@ +/** + * 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. + */ + +#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_PYFUNC_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_PYFUNC_KERNEL_H_ + +#include +#include +#include +#include +#include "pybind11/pybind11.h" +#include "pybind11/numpy.h" +#include "backend/kernel_compiler/cpu/cpu_kernel.h" + +namespace py = pybind11; +namespace mindspore { +namespace kernel { +// Indicate Python object type. The input/output of PyFun should be either Scalar or Numpy Array. +enum class PythonOjectType : char { kScalar, kNumpyArray }; +// Indicate PyFunc input/output information +struct PyFuncArgumentInfo { + // Empty vector indicate the Python object is Scalar and non-empty means Numpy Array. + std::vector> shapes; + // Data type as int, float, bool. + std::vector dtypes; + // Python object type + std::vector object_types; +}; + +class PyFuncCpuKernel : public CPUKernel { + public: + PyFuncCpuKernel() : init_(false), func_id_(-1) {} + ~PyFuncCpuKernel() = default; + + // Init kernel including analyse PyFunc input and output info. + void InitKernel(const CNodePtr &kernel_node) override; + // Construct arguments with raw memory, invoke Python function and then convert result to raw memory. + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs) override; + + protected: + // Analyse PyFunc input/output spec. + void BuildFuncInfo(const CNodePtr &kernel_node); + // Get Python function from anchor. + py::function GetPythonFunc(const int64_t &func_id); + bool ExecuteKernel(const std::vector &inputs, const std::vector &outputs); + + bool init_; + // The Python object is not acceptable for `Primitive` attribute. So we pass an unique key instead of Python function. + // ME store the Python function to a dict, and pass the key to backend kernel. + // The kernel get the Python functhon by the key from the dict when the kernel is first invoked. + size_t func_id_; + py::function py_func_; + // Input and output specifications. + PyFuncArgumentInfo input_infos_; + PyFuncArgumentInfo output_infos_; + // The kernel hold the input tensors during execution to avoid dynamic malloc/free host memory. + std::vector> input_tensors_; +}; +} // namespace kernel +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_PYFUNC_KERNEL_H_ diff --git a/mindspore/ccsrc/runtime/device/cpu/cpu_common.h b/mindspore/ccsrc/runtime/device/cpu/cpu_common.h new file mode 100644 index 00000000000..4ea3637a2ec --- /dev/null +++ b/mindspore/ccsrc/runtime/device/cpu/cpu_common.h @@ -0,0 +1,36 @@ +/** + * Copyright 2019 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. + */ + +#ifndef MINDSPORE_CCSRC_RUNTIME_DEVICE_CPU_CPU_COMMON_H_ +#define MINDSPORE_CCSRC_RUNTIME_DEVICE_CPU_CPU_COMMON_H_ + +#include "utils/log_adapter.h" + +namespace mindspore { +namespace device { +namespace cpu { +#define CHECK_RET_WITH_EXCEPT(expression, status, message) \ + { \ + auto ret = (expression); \ + if (ret != status) { \ + MS_LOG(EXCEPTION) << message; \ + } \ + } +} // namespace cpu +} // namespace device +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_RUNTIME_DEVICE_CPU_CPU_COMMON_H_ diff --git a/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.cc b/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.cc index 96c51cd1c34..d096401288f 100644 --- a/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.cc +++ b/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.cc @@ -31,6 +31,8 @@ namespace cpu { using AnfAlgo = mindspore::session::AnfRuntimeAlgorithm; using mindspore::kernel::KernelBuildInfo; namespace { +constexpr auto kParamDynamic = "dynamic"; + bool IsInputNotCNode(const CNodePtr &kernel_node, size_t input_index) { auto input_node = AnfAlgo::VisitKernel(kernel_node->input(input_index + 1), 0).first; MS_EXCEPTION_IF_NULL(input_node); @@ -66,6 +68,13 @@ void GetOutputDtypes(const CNodePtr &kernel_node, std::vector *output_ty } } +void GetOutputFormat(const CNodePtr &kernel_node, std::vector *output_formats) { + size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node); + for (size_t output_index = 0; output_index < output_num; ++output_index) { + output_formats->emplace_back(kOpFormat_DEFAULT); + } +} + void GetInputDtypes(const CNodePtr &kernel_node, std::vector *input_types, std::vector *input_no_cnode_indexes) { size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node); @@ -81,6 +90,13 @@ void GetInputDtypes(const CNodePtr &kernel_node, std::vector *input_type } } +void GetInputFormat(const CNodePtr &kernel_node, std::vector *input_formats) { + size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node); + for (size_t input_index = 0; input_index < input_num; ++input_index) { + input_formats->emplace_back(kOpFormat_DEFAULT); + } +} + void GetOutputFormatsAndDtypes(const CNodePtr &kernel_node, const KernelAttr &kernel_attr, std::vector *output_formats, std::vector *output_types) { size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node); @@ -200,7 +216,57 @@ void KernelNotSupportException(const AnfNodePtr &kernel_node, const std::vector< operator_info << "is not support."; MS_EXCEPTION(TypeError) << operator_info.str() << " Trace: " << trace::DumpSourceLines(kernel_node); } + +void UpdateDynamicKernelBuildInfoAndAttrs(const CNodePtr &kernel_node) { + const std::string &op_name = AnfAlgo::GetCNodeName(kernel_node); + MS_LOG(INFO) << "Operator name: " << op_name; + // Set kernel build info + std::vector input_types; + std::vector input_not_cnode_indexes; + GetInputDtypes(kernel_node, &input_types, &input_not_cnode_indexes); + std::vector output_types; + GetOutputDtypes(kernel_node, &output_types); + std::vector input_formats; + GetInputFormat(kernel_node, &input_formats); + std::vector output_formats; + GetOutputFormat(kernel_node, &output_formats); + SetKernelBuildInfo(input_formats, input_types, output_formats, output_types, kernel_node.get()); + + // Set kernel attrs + KernelAttr attr; + for (size_t i = 0; i < input_types.size(); i++) { + attr.AddInputAttr(input_types[i]); + } + for (size_t j = 0; j < output_types.size(); j++) { + attr.AddInputAttr(output_types[j]); + } + std::vector kernel_attrs = + kernel::CPUKernelFactory::GetInstance().GetSupportedKernelAttrList(AnfAlgo::GetCNodeName(kernel_node)); + kernel_attrs.emplace_back(attr); + kernel::CPUKernelFactory::GetInstance().UpdateKernelAttrs(op_name, kernel_attrs); + return; +} } // namespace + +bool IsDynamicParamKernel(const std::string &op_name) { + const auto &op_info = kernel::OpLib::FindOp(op_name, kernel::OpImplyType::kCPU); + if (op_info == nullptr) { + return false; + } + + const auto &input_io_info = op_info->inputs_ptr(); + if (input_io_info.size() != 1 || input_io_info[0]->param_type() != kParamDynamic) { + return false; + } + + const auto &output_io_info = op_info->outputs_ptr(); + if (output_io_info.size() != 1 || output_io_info[0]->param_type() != kParamDynamic) { + return false; + } + + return true; +} + bool SelectKernel(const CNodePtr &kernel_node, KernelAttr *selected_kernel_attr, const std::vector &kernel_attrs, const std::vector &input_types, const std::vector &input_not_cnode_indexes, const std::vector &output_types, @@ -229,7 +295,14 @@ bool SelectKernel(const CNodePtr &kernel_node, KernelAttr *selected_kernel_attr, } return false; } + void SetKernelInfo(const CNodePtr &kernel_node) { + // Select for dynamic kernel(both the number and data type are undetermined). + const std::string &op_name = AnfAlgo::GetCNodeName(kernel_node); + if (IsDynamicParamKernel(op_name)) { + return UpdateDynamicKernelBuildInfoAndAttrs(kernel_node); + } + std::vector input_formats; std::vector input_types; std::vector input_not_cnode_indexes; @@ -241,7 +314,6 @@ void SetKernelInfo(const CNodePtr &kernel_node) { kernel::CPUKernelFactory::GetInstance().GetSupportedKernelAttrList(AnfAlgo::GetCNodeName(kernel_node)); if (kernel_attrs.empty() || (kernel_attrs[0].GetInputSize() == 0 && kernel_attrs[0].GetOutputSize() == 0)) { MS_LOG(DEBUG) << "Operator[" << AnfAlgo::GetCNodeName(kernel_node) << "] will get ops attr info."; - std::string op_name = AnfAlgo::GetCNodeName(kernel_node); auto op_info_ptr = mindspore::kernel::OpLib::FindOp(op_name, kernel::OpImplyType::kCPU); if (op_info_ptr == nullptr) { MS_LOG(EXCEPTION) << "Not find op[" << op_name << "] in cpu"; diff --git a/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.h b/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.h index 9fd5c55b7d5..867676cd85e 100644 --- a/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.h +++ b/mindspore/ccsrc/runtime/device/cpu/kernel_select_cpu.h @@ -29,6 +29,8 @@ namespace mindspore { namespace device { namespace cpu { void SetKernelInfo(const CNodePtr &apply_kernel_ptr); +// Indicate whether the kernel input/output number are variable. +bool IsDynamicParamKernel(const std::string &op_name); class KernelAttr { public: diff --git a/tests/ut/cpp/stub/runtime/kernel_select_cpu.cc b/tests/ut/cpp/stub/runtime/kernel_select_cpu.cc new file mode 100644 index 00000000000..5741ab4e26f --- /dev/null +++ b/tests/ut/cpp/stub/runtime/kernel_select_cpu.cc @@ -0,0 +1,11 @@ +#include + +namespace mindspore { +namespace device { +namespace cpu { + +bool IsDynamicParamKernel(const std::string &op_name) { return false; } + +} // namespace cpu +} // namespace device +} // namespace mindspore