wehub-resource-sync
128 lines
4.8 KiB
C++
128 lines
4.8 KiB
C++
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// 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 "paddle/fluid/pybind/python_callable_registry.h"
|
|
#include "glog/logging.h" // For VLOG()
|
|
|
|
#include "paddle/common/errors.h"
|
|
|
|
namespace paddle {
|
|
namespace pybind {
|
|
|
|
namespace py = ::pybind11;
|
|
|
|
void PirCallPythonFunc(py::object *callable,
|
|
const std::vector<pir::Value> &ins,
|
|
std::vector<pir::Value> *outs) {
|
|
py::gil_scoped_acquire guard;
|
|
py::tuple in_args(ins.size());
|
|
for (size_t i = 0; i < ins.size(); ++i) {
|
|
in_args[i] = py::cast(ins[i]);
|
|
}
|
|
|
|
VLOG(2) << "[PirCallPythonFunc] ins.size() = " << ins.size();
|
|
VLOG(2) << "[PirCallPythonFunc] outs->size() = " << outs->size();
|
|
|
|
auto ret = (*callable)(*in_args);
|
|
auto ret_tuple = py::cast<py::tuple>(ret);
|
|
size_t ret_num = py::len(ret_tuple);
|
|
size_t out_num = outs->size();
|
|
if (UNLIKELY(ret_num != out_num)) {
|
|
// Python function has no return values or returns None
|
|
// In this case, ret_num = 1 && ret[0] == None && out_num should be 0
|
|
// Otherwise, ret_num must be equal to out_num
|
|
PADDLE_ENFORCE_EQ(ret_num == 1,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Python function has no return values or returns "
|
|
"None. In this case, ret_num = 1 && ret[0] == None "
|
|
"&& out_num should be 0. But ret_num is %d",
|
|
ret_num));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
out_num == 0,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Python function has no return values or returns None. In "
|
|
"this case, ret_num = 1 && ret[0] == None && out_num should "
|
|
"be 0. But out_num is %d",
|
|
out_num));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
py::cast<pir::Value *>(ret_tuple[0]) == nullptr,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Python function has no return values or returns None. In "
|
|
"this case, ret_num = 1 && ret[0] == None && out_num should "
|
|
"be 0. But ret[0] is not None"));
|
|
}
|
|
|
|
for (size_t i = 0; i < out_num; ++i) {
|
|
try {
|
|
if (ret_tuple[i].is_none()) {
|
|
VLOG(6) << "Set Output( " << i << " ) value as fake_value";
|
|
(*outs)[i] = pir::Value(nullptr);
|
|
} else {
|
|
auto py_out_value = py::cast<pir::Value>(ret_tuple[i]);
|
|
PADDLE_ENFORCE_NOT_NULL(
|
|
py_out_value.impl(),
|
|
common::errors::InvalidArgument(
|
|
"Output value %d should not be nullptr", i));
|
|
(*outs)[i] = py_out_value;
|
|
}
|
|
} catch (py::cast_error &) {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"pybind11::cast to pir::Value error. The %d-th output exception is "
|
|
"pir::Value",
|
|
i));
|
|
}
|
|
}
|
|
}
|
|
|
|
PythonCallableRegistrar &PythonCallableRegistrar::GetInstance() {
|
|
static thread_local PythonCallableRegistrar instance;
|
|
return instance;
|
|
}
|
|
void PythonCallableRegistrar::Register(uint64_t unique_id,
|
|
const py::object &callable) {
|
|
if (python_callable_registry_.find(unique_id) !=
|
|
python_callable_registry_.end()) {
|
|
LOG(WARNING) << "unique_id " << unique_id
|
|
<< " has already registered python function in "
|
|
"python_callable_registry_. This operation "
|
|
"will override the old one.";
|
|
}
|
|
python_callable_registry_[unique_id] = callable;
|
|
}
|
|
|
|
// Return pybind11::object* instead of pybind11::object
|
|
// Returning pybind11::object would cause reference count increasing
|
|
// but without GIL, reference count in Python may not be safe
|
|
py::object *PythonCallableRegistrar::Get(uint64_t unique_id) {
|
|
PADDLE_ENFORCE_NE(
|
|
python_callable_registry_.find(unique_id),
|
|
python_callable_registry_.end(),
|
|
common::errors::InvalidArgument(
|
|
"Unique_id %d is not found in python_callable_registry_. The "
|
|
"possible "
|
|
"reasons are below:\n"
|
|
"1. The python callable was not registered for `unique_id` by "
|
|
"`PythonCallableRegistrar::Register`\n"
|
|
"2. The python callable was remove from python_callable_registry_",
|
|
unique_id));
|
|
return &(python_callable_registry_[unique_id]);
|
|
}
|
|
|
|
} // namespace pybind
|
|
} // namespace paddle
|