#include "module.h"
#include <sstream>
#include <ctranslate2/storage_view.h>
#include "utils.h"
using namespace pybind11::literals;
namespace ctranslate2 {
namespace python {
static DataType typestr_to_dtype(const std::string& typestr) {
const auto type_code = typestr[1];
const auto num_bytes = typestr[2];
if (type_code == 'i') {
if (num_bytes == '1')
return DataType::INT8;
if (num_bytes == '2')
return DataType::INT16;
if (num_bytes == '4')
return DataType::INT32;
} else if (type_code == 'f') {
if (num_bytes == '2')
return DataType::FLOAT16;
if (num_bytes == '4')
return DataType::FLOAT32;
}
throw std::invalid_argument("Unsupported type: " + typestr);
}
static std::string dtype_to_typestr(const DataType dtype) {
switch (dtype) {
case DataType::FLOAT32:
return "<f4";
case DataType::FLOAT16:
return "<f2";
case DataType::INT8:
return "|i1";
case DataType::INT16:
return "<i2";
case DataType::INT32:
return "<i4";
default:
throw std::runtime_error("Data type " + dtype_name(dtype) + " is not supported in the array "
"interface. You should first convert this storage to a supported "
"type, for example with "
"`storage = storage.to(ctranslate2.DataType.float32)`");
}
}
static StorageView create_view_from_array(py::object array) {
auto device = Device::CPU;
py::object interface_obj = py::getattr(array, "__array_interface__", py::none());
if (interface_obj.is_none()) {
interface_obj = py::getattr(array, "__cuda_array_interface__", py::none());
if (interface_obj.is_none())
throw std::invalid_argument("Object does not implement the array interface");
device = Device::CUDA;
}
py::dict interface = interface_obj.cast<py::dict>();
if (interface_obj.contains("strides") && !interface_obj["strides"].is_none())
throw std::invalid_argument("StorageView does not support arrays with non contiguous memory");
auto shape = interface["shape"].cast<Shape>();
auto dtype = typestr_to_dtype(interface["typestr"].cast<std::string>());
auto data = interface["data"].cast<py::tuple>();
auto ptr = data[0].cast<uintptr_t>();
auto read_only = data[1].cast<bool>();
if (read_only)
throw std::invalid_argument("StorageView does not support read-only arrays");
StorageView view(dtype, device);
view.view((void*)ptr, std::move(shape));
return view;
}
static py::dict get_array_interface(const StorageView& view) {
py::tuple shape(view.rank());
for (size_t i = 0; i < shape.size(); ++i)
shape[i] = view.dim(i);
return py::dict(
"shape"_a=shape,
"typestr"_a=dtype_to_typestr(view.dtype()),
"data"_a=py::make_tuple((uintptr_t)view.buffer(), false),
"version"_a=3);
}
void register_storage_view(py::module& m) {
py::enum_<DataType>(m, "DataType")
.value("float32", DataType::FLOAT32)
.value("float16", DataType::FLOAT16)
.value("bfloat16", DataType::BFLOAT16)
.value("int8", DataType::INT8)
.value("int16", DataType::INT16)
.value("int32", DataType::INT32)
;
py::enum_<Device>(m, "Device")
.value("cpu", Device::CPU)
.value("cuda", Device::CUDA)
;
py::class_<StorageView>(
m, "StorageView",
R"pbdoc(
An allocated buffer with shape information.
The object implements the
`Array Interface <https://numpy.org/doc/stable/reference/arrays.interface.html>`_
and the
`CUDA Array Interface <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_
so that it can be passed to Numpy or PyTorch without copy.
Example:
>>> x = np.ones((2, 4), dtype=np.int32)
>>> y = ctranslate2.StorageView.from_array(x)
>>> print(y)
1 1 1 ... 1 1 1
[cpu:0 int32 storage viewed as 2x4]
>>> z = np.array(y)
...
>>> x = torch.ones((2, 4), dtype=torch.int32, device="cuda")
>>> y = ctranslate2.StorageView.from_array(x)
>>> print(y)
1 1 1 ... 1 1 1
[cuda:0 int32 storage viewed as 2x4]
>>> z = torch.as_tensor(y, device="cuda")
)pbdoc")
.def_static("from_array", &create_view_from_array, py::arg("array"),
py::keep_alive<0, 1>(),
R"pbdoc(
Creates a ``StorageView`` from an object implementing the array interface.
Arguments:
array: An object implementing the array interface (e.g. a Numpy array
or a PyTorch Tensor).
Returns:
A new ``StorageView`` instance sharing the same data as the input array.
Raises:
ValueError: if the object does not implement the array interface or
uses an unsupported array specification.
)pbdoc")
.def_property_readonly("dtype", &StorageView::dtype,
"Data type used by the storage.")
.def_property_readonly("shape", &StorageView::shape,
"Shape of the storage view.")
.def_property_readonly("device_index", &StorageView::device_index,
"Device index.")
.def_property_readonly("device",
[](const StorageView& view) {
return device_to_str(view.device());
},
"Device where the storage is allocated (\"cpu\" or \"cuda\").")
.def_property_readonly("__array_interface__", [](const StorageView& view) {
if (view.device() == Device::CUDA)
throw py::attribute_error("Cannot get __array_interface__ when the StorageView "
"is viewing a CUDA array");
return get_array_interface(view);
})
.def_property_readonly("__cuda_array_interface__", [](const StorageView& view) {
if (view.device() == Device::CPU)
throw py::attribute_error("Cannot get __cuda_array_interface__ when the StorageView "
"is viewing a CPU array");
return get_array_interface(view);
})
.def("__str__", [](const StorageView& view) {
std::ostringstream stream;
stream << view;
return stream.str();
})
.def("__repr__", [](const StorageView& view) {
std::ostringstream stream;
stream << view;
return stream.str();
})
.def("to",
[](const StorageView& view, DataType dtype) {
ScopedDeviceSetter device_setter(view.device(), view.device_index());
StorageView converted = view.to(dtype);
synchronize_stream(view.device());
return converted;
},
py::arg("dtype"),
py::call_guard<py::gil_scoped_release>(),
R"pbdoc(
Converts the storage to another type.
Arguments:
dtype: The data type to convert to.
Returns:
A new ``StorageView`` instance.
)pbdoc")
.def("to_device",
[](const StorageView& view, Device device) {
ScopedDeviceSetter device_setter(view.device(), view.device_index());
StorageView converted = view.to(device);
synchronize_stream(view.device());
return converted;
},
py::arg("device"),
py::call_guard<py::gil_scoped_release>(),
R"pbdoc(
Converts the storage to another device.
Arguments:
device: The device to copy the data to.
Returns:
A new ``StorageView`` instance.
)pbdoc")
;
}
}
}