| // Copyright 2023 The IREE Authors |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| #ifndef IREE_BINDINGS_PYTHON_NUMPY_INTEROP_H_ |
| #define IREE_BINDINGS_PYTHON_NUMPY_INTEROP_H_ |
| #include "iree/hal/api.h" |
| namespace iree::python::numpy { |
| // Must be called in init of extension module. |
| void InitializeNumPyInterop(); |
| // Converts an IREE element type to a NumPy NPY_TYPES value. |
| int ConvertHalElementTypeToNumPyTypeNum(iree_hal_element_type_t t); |
| // Wraps a call to PyArray_DescrNewFromType(int). |
| py::object DescrNewFromType(int typenum); |
| // Extracts a typenum from a dtype (descriptor) object. |
| int TypenumFromDescr(py::handle dtype); |
| // Delegates to PyArray_SimpleNewFromData and sets the base_object. |
| py::object SimpleNewFromData(int nd, intptr_t const *dims, int typenum, |
| void *data, py::handle base_object); |
| } // namespace iree::python::numpy |
| #endif // IREE_BINDINGS_PYTHON_NUMPY_INTEROP_H_ |
