host: Allow to use numpy array slices for send/recv calls

The current implementation used numpy's macros to get access to the
underlying data when doing send and recv calls. Because the memory
layout of numpy arrays and UHD data arrays is different we need to
create a vector pointing into the numpy array for each channel.
The marcos used to get access to the numpy array made copies of
the array data so modifications by the recv call did not made it
into the numpy array.

The change now uses the pybind11 array wrapper for numpy arrays.
We therefore can now use class methods to query the properties
of the array instead of macros. The array class also keeps track
of the reference counter for us. Finally it provides a method to
access the mutable data of the numpy array. Therefore also passing
slices to the receive call works.
This commit is contained in:
Lars Amsel 2025-05-27 16:58:26 +02:00 committed by Jörg Hofrichter
parent 2ad396f9c3
commit 70664d552b

View File

@ -12,31 +12,30 @@
#include <uhd/types/metadata.hpp>
#include <boost/format.hpp>
#include <numpy/ndarrayobject.h>
#include <numpy/ndarraytypes.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
static size_t wrap_recv(uhd::rx_streamer* rx_stream,
py::object& np_array,
py::array& array,
uhd::rx_metadata_t& metadata,
const double timeout = 0.1)
{
// Get a numpy array object from given python object
// No sanity checking possible!
PyObject* array_obj = PyArray_FROM_OF(np_array.ptr(), NPY_ARRAY_CARRAY);
PyArrayObject* array_type_obj = reinterpret_cast<PyArrayObject*>(array_obj);
// Get dimensions of the numpy array
const size_t dims = PyArray_NDIM(array_type_obj);
const npy_intp* shape = PyArray_SHAPE(array_type_obj);
// How many bytes to jump to get to the next element of this stride
// (next row)
const npy_intp* strides = PyArray_STRIDES(array_type_obj);
const size_t channels = rx_stream->get_num_channels();
// Using the pybind11's numpy array wrapper provides easy
// access instead of using numpy's C API macros and also keeps
// track of the reference count for the numpy array.
if (!array.writeable()) {
throw std::runtime_error("Array is not writable");
}
char* data = static_cast<char*>(array.mutable_data());
const size_t dims = array.ndim();
const size_t channels = rx_stream->get_num_channels();
// Check if numpy array sizes are okay
if (((channels > 1) && (dims != 2)) or ((size_t)shape[0] < channels)) {
// Manually decrement the ref count
Py_DECREF(array_obj);
if (((channels > 1) && (dims != 2)) or ((size_t)array.shape(0) < channels)) {
// If we don't have a 2D NumPy array, assume we have a 1D array
size_t input_channels = (dims != 2) ? 1 : shape[0];
size_t input_channels = (dims != 2) ? 1 : array.shape(0);
throw uhd::runtime_error(
str(boost::format("Number of RX channels (%d) does not match the dimensions "
"of the data array (%d)")
@ -45,18 +44,12 @@ static size_t wrap_recv(uhd::rx_streamer* rx_stream,
// Get a pointer to the storage
std::vector<void*> channel_storage;
char* data = PyArray_BYTES(array_type_obj);
for (size_t i = 0; i < channels; ++i) {
channel_storage.push_back((void*)(data + i * strides[0]));
channel_storage.push_back((void*)(data + i * array.strides(0)));
}
// Get data buffer and size of the array
size_t nsamps_per_buff;
if (dims > 1) {
nsamps_per_buff = (size_t)shape[1];
} else {
nsamps_per_buff = PyArray_SIZE(array_type_obj);
}
// determine how many samples we can receive
size_t nsamps_per_buff = dims > 1 ? array.shape(1) : array.size();
// Release the GIL only for the recv() call
const size_t result = [&]() {
@ -65,38 +58,25 @@ static size_t wrap_recv(uhd::rx_streamer* rx_stream,
return rx_stream->recv(channel_storage, nsamps_per_buff, metadata, timeout);
}();
// Manually decrement the ref count
Py_DECREF(array_obj);
return result;
}
static size_t wrap_send(uhd::tx_streamer* tx_stream,
py::object& np_array,
py::array& array,
uhd::tx_metadata_t& metadata,
const double timeout = 0.1)
{
// Get a numpy array object from given python object
// No sanity checking possible!
// Note: this increases the ref count, which we'll need to manually decrease at the
// end
PyObject* array_obj = PyArray_FROM_OF(np_array.ptr(), NPY_ARRAY_CARRAY);
PyArrayObject* array_type_obj = reinterpret_cast<PyArrayObject*>(array_obj);
// Get dimensions of the numpy array
const size_t dims = PyArray_NDIM(array_type_obj);
const npy_intp* shape = PyArray_SHAPE(array_type_obj);
// How many bytes to jump to get to the next element of the stride
// (next row)
const npy_intp* strides = PyArray_STRIDES(array_type_obj);
const size_t channels = tx_stream->get_num_channels();
// Using pybind11's numpy array wrapper provides easy
// access instead of using numpy's C API macros and also keeps
// track of the reference count for the numpy array.
const char* data = static_cast<const char*>(array.data());
const size_t dims = array.ndim();
const size_t channels = tx_stream->get_num_channels();
// Check if numpy array sizes are ok
if (((channels > 1) && (dims != 2)) or ((size_t)shape[0] < channels)) {
// Manually decrement the ref count
Py_DECREF(array_obj);
if (((channels > 1) && (dims != 2)) or ((size_t)array.shape(0) < channels)) {
// If we don't have a 2D NumPy array, assume we have a 1D array
size_t input_channels = (dims != 2) ? 1 : shape[0];
size_t input_channels = (dims != 2) ? 1 : array.shape(0);
throw uhd::runtime_error(
str(boost::format("Number of TX channels (%d) does not match the dimensions "
"of the data array (%d)")
@ -105,13 +85,12 @@ static size_t wrap_send(uhd::tx_streamer* tx_stream,
// Get a pointer to the storage
std::vector<void*> channel_storage;
char* data = PyArray_BYTES(array_type_obj);
for (size_t i = 0; i < channels; ++i) {
channel_storage.push_back((void*)(data + i * strides[0]));
channel_storage.push_back((void*)(data + i * array.strides(0)));
}
// Get data buffer and size of the array
size_t nsamps_per_buff = (dims > 1) ? (size_t)shape[1] : PyArray_SIZE(array_type_obj);
// determine how many samples we have to send
size_t nsamps_per_buff = (dims > 1) ? (size_t)array.shape(1) : array.size();
// Release the GIL only for the send() call
const size_t result = [&]() {
@ -120,8 +99,6 @@ static size_t wrap_send(uhd::tx_streamer* tx_stream,
return tx_stream->send(channel_storage, nsamps_per_buff, metadata, timeout);
}();
// Manually decrement the ref count
Py_DECREF(array_obj);
return result;
}