| 171 | } |
| 172 | |
| 173 | void PyCompute(user_op::KernelComputeContext* ctx, const std::string& py_func_name) { |
| 174 | const std::string& op_type_name = ctx->op_type_name(); |
| 175 | const user_op::OpRegistryResult* val = |
| 176 | user_op::UserOpRegistryMgr::Get().GetOpRegistryResult(op_type_name); |
| 177 | CHECK(val) << "Op op_type_name " << op_type_name << " has no definition."; |
| 178 | const UserOpDef& op_def = val->op_def; |
| 179 | |
| 180 | // get GIL |
| 181 | PyGILState_STATE py_gil_st; |
| 182 | py_gil_st = PyGILState_Ensure(); |
| 183 | // prepare for numpy c api |
| 184 | if (PyArray_API == nullptr) { _import_array(); } |
| 185 | |
| 186 | PyObject *py_str, *py_module, *py_func; |
| 187 | PyObject *py_inputs, *py_outputs; |
| 188 | |
| 189 | // get python kernel |
| 190 | static const std::string forward_suffix = "_forward"; |
| 191 | static const std::string backward_suffix = "_backward"; |
| 192 | std::string op_module_name = op_type_name; |
| 193 | if (op_type_name.size() > forward_suffix.size() |
| 194 | && op_type_name.rfind(forward_suffix) == (op_type_name.size() - forward_suffix.size())) { |
| 195 | op_module_name = op_type_name.substr(0, op_type_name.size() - forward_suffix.size()); |
| 196 | } |
| 197 | if (op_type_name.size() > backward_suffix.size() |
| 198 | && op_type_name.rfind(backward_suffix) == (op_type_name.size() - backward_suffix.size())) { |
| 199 | op_module_name = op_type_name.substr(0, op_type_name.size() - backward_suffix.size()); |
| 200 | } |
| 201 | py_str = PyUnicode_DecodeFSDefault(op_module_name.c_str()); |
| 202 | CHECK(py_kernels_dic) << "py_kernels_dic should not be nullptr."; |
| 203 | py_module = PyDict_GetItem(py_kernels_dic, py_str); |
| 204 | CHECK(py_module) << op_module_name << " has no python kernel."; |
| 205 | Py_DECREF(py_str); |
| 206 | |
| 207 | // get func |
| 208 | py_func = PyObject_GetAttrString(py_module, py_func_name.c_str()); |
| 209 | if (py_func == nullptr || !PyCallable_Check(py_func)) { |
| 210 | Py_DECREF(py_module); |
| 211 | PyErr_Print(); |
| 212 | } |
| 213 | |
| 214 | // get numpy input |
| 215 | MakePyInputs(op_def, ctx, &py_inputs); |
| 216 | |
| 217 | // call func |
| 218 | py_outputs = PyEval_CallObject(py_func, py_inputs); |
| 219 | Py_DECREF(py_inputs); |
| 220 | |
| 221 | // get numpy output |
| 222 | GetPyOutputs(op_def, ctx, py_outputs); |
| 223 | |
| 224 | Py_XDECREF(py_func); |
| 225 | Py_DECREF(py_outputs); |
| 226 | |
| 227 | // release GIL |
| 228 | PyGILState_Release(py_gil_st); |
| 229 | } |
| 230 |
no test coverage detected