Skip to content

perf: optimize generic tensor conversion #832

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
baominghelly wants to merge 1 commit into
base: master
Choose a base branch
from
Open

perf: optimize generic tensor conversion #832

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 150 additions & 20 deletions src/pybind11_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@
#include <pybind11/stl.h>

#include <algorithm>
#include <limits>
#include <string_view>
#include <type_traits>

#include "data_type.h"
#include "host_range_profiler.h"
#include "tensor.h"
#include "torch/device_.h"
Expand All @@ -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<py::object>(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<py::object>(value);
}

template <typename Integer>
Integer IntegerFromPyObject(PyObject* obj) {
static_assert(std::is_integral_v<Integer>);
if constexpr (std::is_unsigned_v<Integer>) {
const auto value{PyLong_AsUnsignedLongLong(obj)};
if (value == static_cast<unsigned long long>(-1) && PyErr_Occurred()) {
throw py::error_already_set();
}
if (value > std::numeric_limits<Integer>::max()) {
PyErr_SetString(PyExc_OverflowError,
"integer out of range for tensor metadata");
throw py::error_already_set();
}
return static_cast<Integer>(value);
} else {
const auto value{PyLong_AsLongLong(obj)};
if (value == -1 && PyErr_Occurred()) throw py::error_already_set();
if (value < std::numeric_limits<Integer>::min() ||
value > std::numeric_limits<Integer>::max()) {
PyErr_SetString(PyExc_OverflowError,
"integer out of range for tensor metadata");
throw py::error_already_set();
}
return static_cast<Integer>(value);
}
}

template <typename Vector>
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<py::object>(sequence_ptr)};
const auto size{PySequence_Fast_GET_SIZE(sequence.ptr())};

Vector result;
result.reserve(static_cast<std::size_t>(size));
for (Py_ssize_t i = 0; i < size; ++i) {
result.push_back(IntegerFromPyObject<typename Vector::value_type>(
PySequence_Fast_GET_ITEM(sequence.ptr(), i)));
}
return result;
}

template <Device::Type... kDevs>
std::unordered_map<std::string, Device::Type> BuildTorchNameMap(
List<kDevs...>) {
Expand All @@ -26,7 +128,7 @@ std::unordered_map<std::string, Device::Type> 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;
Expand All @@ -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<std::string>()};
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<std::size_t>(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
Expand All @@ -53,11 +159,13 @@ inline DataType DataTypeFromPybind11Handle(py::handle obj) {
}

template <typename T = void>
inline Device::Type DeviceTypeFromString(const std::string& name) {
inline Device::Type DeviceTypeFromString(std::string_view name) {
static const auto kTorchNameToTypes{
detail::BuildTorchNameMap(ActiveDevices<T>{})};

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;
Expand All @@ -66,23 +174,24 @@ inline Device::Type DeviceTypeFromString(const std::string& name) {
std::vector<std::string> 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());
supported_names.erase(
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) {
Expand All @@ -101,11 +210,13 @@ inline Device::Type DeviceTypeFromString(const std::string& name) {
// devices an op may not support, without crashing the process.
template <typename T = void>
inline std::optional<Device::Type> TryDeviceTypeFromString(
const std::string& name) {
std::string_view name) {
static const auto kTorchNameToTypes{
detail::BuildTorchNameMap(ActiveDevices<T>{})};

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;
Expand All @@ -122,7 +233,9 @@ inline std::optional<Device::Type> 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;
Expand All @@ -134,26 +247,43 @@ inline std::optional<Device::Type> TryDeviceTypeFromString(
namespace detail {

inline Device DeviceFromPybind11HandleImpl(py::handle obj) {
auto device_obj{obj.attr("device")};
auto device_type_str{device_obj.attr("type").cast<std::string>()};
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<std::size_t>(device_type_size)};
} else {
device_type_storage = device_type_obj.cast<std::string>();
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<int>()};

return Device{DeviceTypeFromString(device_type_str), device_index};
}

inline Tensor TensorFromPybind11HandleImpl(py::handle obj) {
auto data{
reinterpret_cast<void*>(obj.attr("data_ptr")().cast<std::uintptr_t>())};
auto data{reinterpret_cast<void*>(
detail::CallMethodNoArgs(obj, detail::DataPtrName())
.cast<std::uintptr_t>())};

auto shape{obj.attr("shape").cast<typename Tensor::Shape>()};
auto shape{detail::VectorFromSequence<typename Tensor::Shape>(
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<typename Tensor::Strides>()};
auto strides{detail::VectorFromSequence<typename Tensor::Strides>(
detail::CallMethodNoArgs(obj, detail::StrideName()))};

return Tensor{data, std::move(shape), dtype, device, std::move(strides)};
}
Expand Down
Loading