diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 292a2c00b..315989152 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -5,7 +5,11 @@ #include #include +#include +#include +#include +#include "data_type.h" #include "host_range_profiler.h" #include "tensor.h" #include "torch/device_.h" @@ -16,6 +20,104 @@ namespace infini::ops { namespace detail { +inline PyObject* InternedName(const char* value) { + // Keep one reference for the process lifetime; Python objects must not be + // decref'd by a static destructor after interpreter finalization. + auto* name{PyUnicode_InternFromString(value)}; + if (name == nullptr) throw py::error_already_set(); + return name; +} + +inline PyObject* DataPtrName() { + static PyObject* const name{InternedName("data_ptr")}; + return name; +} + +inline PyObject* ShapeName() { + static PyObject* const name{InternedName("shape")}; + return name; +} + +inline PyObject* DTypeName() { + static PyObject* const name{InternedName("dtype")}; + return name; +} + +inline PyObject* DeviceName() { + static PyObject* const name{InternedName("device")}; + return name; +} + +inline PyObject* TypeName() { + static PyObject* const name{InternedName("type")}; + return name; +} + +inline PyObject* IndexName() { + static PyObject* const name{InternedName("index")}; + return name; +} + +inline PyObject* StrideName() { + static PyObject* const name{InternedName("stride")}; + return name; +} + +inline py::object GetAttr(py::handle obj, PyObject* name) { + auto* value{PyObject_GetAttr(obj.ptr(), name)}; + if (value == nullptr) throw py::error_already_set(); + return py::reinterpret_steal(value); +} + +inline py::object CallMethodNoArgs(py::handle obj, PyObject* name) { + auto* value{PyObject_CallMethodNoArgs(obj.ptr(), name)}; + if (value == nullptr) throw py::error_already_set(); + return py::reinterpret_steal(value); +} + +template +Integer IntegerFromPyObject(PyObject* obj) { + static_assert(std::is_integral_v); + if constexpr (std::is_unsigned_v) { + const auto value{PyLong_AsUnsignedLongLong(obj)}; + if (value == static_cast(-1) && PyErr_Occurred()) { + throw py::error_already_set(); + } + if (value > std::numeric_limits::max()) { + PyErr_SetString(PyExc_OverflowError, + "integer out of range for tensor metadata"); + throw py::error_already_set(); + } + return static_cast(value); + } else { + const auto value{PyLong_AsLongLong(obj)}; + if (value == -1 && PyErr_Occurred()) throw py::error_already_set(); + if (value < std::numeric_limits::min() || + value > std::numeric_limits::max()) { + PyErr_SetString(PyExc_OverflowError, + "integer out of range for tensor metadata"); + throw py::error_already_set(); + } + return static_cast(value); + } +} + +template +Vector VectorFromSequence(py::handle obj) { + auto* sequence_ptr{PySequence_Fast(obj.ptr(), "expected a sequence")}; + if (sequence_ptr == nullptr) throw py::error_already_set(); + auto sequence{py::reinterpret_steal(sequence_ptr)}; + const auto size{PySequence_Fast_GET_SIZE(sequence.ptr())}; + + Vector result; + result.reserve(static_cast(size)); + for (Py_ssize_t i = 0; i < size; ++i) { + result.push_back(IntegerFromPyObject( + PySequence_Fast_GET_ITEM(sequence.ptr(), i))); + } + return result; +} + template std::unordered_map BuildTorchNameMap( List) { @@ -26,7 +128,7 @@ std::unordered_map BuildTorchNameMap( } // namespace detail -inline DataType DataTypeFromString(const std::string& name) { +inline DataType DataTypeFromString(std::string_view name) { // InfiniRT has no bool dtype; carry bool tensor storage as byte data and // restore bool semantics in operators that accept bool tensors. if (name == "bool") return DataType::kUInt8; @@ -37,11 +139,15 @@ inline DataType DataTypeFromString(const std::string& name) { namespace detail { inline DataType DataTypeFromPybind11HandleImpl(py::handle obj) { - auto dtype_str{py::str(obj).cast()}; + py::str dtype_obj{obj}; + Py_ssize_t size{0}; + const char* data{PyUnicode_AsUTF8AndSize(dtype_obj.ptr(), &size)}; + if (data == nullptr) throw py::error_already_set(); + std::string_view dtype_str{data, static_cast(size)}; const auto pos{dtype_str.find_last_of('.')}; return DataTypeFromString( - pos == std::string::npos ? dtype_str : dtype_str.substr(pos + 1)); + pos == std::string_view::npos ? dtype_str : dtype_str.substr(pos + 1)); } } // namespace detail @@ -53,11 +159,13 @@ inline DataType DataTypeFromPybind11Handle(py::handle obj) { } template -inline Device::Type DeviceTypeFromString(const std::string& name) { +inline Device::Type DeviceTypeFromString(std::string_view name) { static const auto kTorchNameToTypes{ detail::BuildTorchNameMap(ActiveDevices{})}; - auto it{kTorchNameToTypes.find(name)}; + auto it{ + std::find_if(kTorchNameToTypes.cbegin(), kTorchNameToTypes.cend(), + [name](const auto& item) { return item.first == name; })}; if (it != kTorchNameToTypes.cend()) { return it->second; @@ -66,14 +174,14 @@ inline Device::Type DeviceTypeFromString(const std::string& name) { std::vector supported_names; for (const auto& [torch_name, device_type] : kTorchNameToTypes) { - const auto internal_name = std::string{Device::StringFromType(device_type)}; + const auto internal_name = Device::StringFromType(device_type); if (name == internal_name) { return device_type; } supported_names.push_back(torch_name); - supported_names.push_back(internal_name); + supported_names.emplace_back(internal_name); } std::sort(supported_names.begin(), supported_names.end()); @@ -81,8 +189,9 @@ inline Device::Type DeviceTypeFromString(const std::string& name) { std::unique(supported_names.begin(), supported_names.end()), supported_names.end()); - std::string message = "Unsupported device type `" + name + - "` for this InfiniOps build. Supported device names: "; + std::string message{"Unsupported device type `"}; + message.append(name.data(), name.size()); + message += "` for this InfiniOps build. Supported device names: "; for (std::size_t i = 0; i < supported_names.size(); ++i) { if (i != 0) { @@ -101,11 +210,13 @@ inline Device::Type DeviceTypeFromString(const std::string& name) { // devices an op may not support, without crashing the process. template inline std::optional TryDeviceTypeFromString( - const std::string& name) { + std::string_view name) { static const auto kTorchNameToTypes{ detail::BuildTorchNameMap(ActiveDevices{})}; - auto it{kTorchNameToTypes.find(name)}; + auto it{ + std::find_if(kTorchNameToTypes.cbegin(), kTorchNameToTypes.cend(), + [name](const auto& item) { return item.first == name; })}; if (it != kTorchNameToTypes.cend()) { return it->second; @@ -122,7 +233,9 @@ inline std::optional TryDeviceTypeFromString( {"hygon", Device::Type::kHygon}, }; - auto platform_it{kPlatformNames.find(name)}; + auto platform_it{ + std::find_if(kPlatformNames.cbegin(), kPlatformNames.cend(), + [name](const auto& item) { return item.first == name; })}; if (platform_it != kPlatformNames.cend()) { return platform_it->second; @@ -134,9 +247,22 @@ inline std::optional TryDeviceTypeFromString( namespace detail { inline Device DeviceFromPybind11HandleImpl(py::handle obj) { - auto device_obj{obj.attr("device")}; - auto device_type_str{device_obj.attr("type").cast()}; - auto device_index_obj{device_obj.attr("index")}; + auto device_obj{detail::GetAttr(obj, detail::DeviceName())}; + auto device_type_obj{detail::GetAttr(device_obj, detail::TypeName())}; + std::string device_type_storage; + std::string_view device_type_str; + if (PyUnicode_Check(device_type_obj.ptr())) { + Py_ssize_t device_type_size{0}; + const char* device_type_data{ + PyUnicode_AsUTF8AndSize(device_type_obj.ptr(), &device_type_size)}; + if (device_type_data == nullptr) throw py::error_already_set(); + device_type_str = {device_type_data, + static_cast(device_type_size)}; + } else { + device_type_storage = device_type_obj.cast(); + device_type_str = device_type_storage; + } + auto device_index_obj{detail::GetAttr(device_obj, detail::IndexName())}; auto device_index{device_index_obj.is_none() ? 0 : device_index_obj.cast()}; @@ -144,16 +270,20 @@ inline Device DeviceFromPybind11HandleImpl(py::handle obj) { } inline Tensor TensorFromPybind11HandleImpl(py::handle obj) { - auto data{ - reinterpret_cast(obj.attr("data_ptr")().cast())}; + auto data{reinterpret_cast( + detail::CallMethodNoArgs(obj, detail::DataPtrName()) + .cast())}; - auto shape{obj.attr("shape").cast()}; + auto shape{detail::VectorFromSequence( + detail::GetAttr(obj, detail::ShapeName()))}; - auto dtype{DataTypeFromPybind11HandleImpl(obj.attr("dtype"))}; + auto dtype{DataTypeFromPybind11HandleImpl( + detail::GetAttr(obj, detail::DTypeName()))}; auto device{DeviceFromPybind11HandleImpl(obj)}; - auto strides{obj.attr("stride")().cast()}; + auto strides{detail::VectorFromSequence( + detail::CallMethodNoArgs(obj, detail::StrideName()))}; return Tensor{data, std::move(shape), dtype, device, std::move(strides)}; }