| // Copyright 2019 Google LLC |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // https://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| #include "pyiree/rt/hal.h" |
| |
| #include "absl/container/inlined_vector.h" |
| #include "iree/hal/api.h" |
| |
| namespace iree { |
| namespace python { |
| |
| namespace { |
| |
| class HalMappedMemory { |
| public: |
| HalMappedMemory(iree_hal_mapped_memory_t mapped_memory, |
| iree_hal_buffer_view_t* bv) |
| : mapped_memory_(mapped_memory), bv_(bv) { |
| iree_hal_buffer_view_retain(bv_); |
| } |
| ~HalMappedMemory() { |
| if (bv_) { |
| iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(bv_); |
| IREE_CHECK_OK(iree_hal_buffer_unmap(buffer, &mapped_memory_)); |
| iree_hal_buffer_view_release(bv_); |
| } |
| } |
| HalMappedMemory(HalMappedMemory&& other) |
| : mapped_memory_(other.mapped_memory_), bv_(other.bv_) { |
| other.bv_ = nullptr; |
| } |
| |
| static HalMappedMemory Create(HalBufferView& bv) { |
| iree_hal_buffer_t* buffer = iree_hal_buffer_view_buffer(bv.raw_ptr()); |
| iree_device_size_t byte_length = iree_hal_buffer_byte_length(buffer); |
| iree_hal_mapped_memory_t mapped_memory; |
| CheckApiStatus(iree_hal_buffer_map(buffer, IREE_HAL_MEMORY_ACCESS_READ, |
| 0 /* element_offset */, byte_length, |
| &mapped_memory), |
| "Could not map memory"); |
| return HalMappedMemory(mapped_memory, bv.raw_ptr()); |
| } |
| |
| py::buffer_info ToBufferInfo() { |
| absl::InlinedVector<int32_t, 6> shape(iree_hal_buffer_view_shape_rank(bv_)); |
| CheckApiStatus( |
| iree_hal_buffer_view_shape(bv_, shape.size(), shape.data(), nullptr), |
| "Error getting buffer view shape"); |
| iree_hal_element_type_t element_type = |
| iree_hal_buffer_view_element_type(bv_); |
| int32_t element_size = iree_hal_element_byte_count(element_type); |
| absl::InlinedVector<py::ssize_t, 6> dims(shape.size()); |
| for (int i = 0; i < shape.size(); ++i) { |
| dims[i] = shape[i]; |
| } |
| absl::InlinedVector<py::ssize_t, 8> strides(shape.size()); |
| if (!strides.empty()) { |
| strides[shape.size() - 1] = element_size; |
| for (int i = shape.size() - 2; i >= 0; --i) { |
| strides[i] = strides[i + 1] * shape[i + 1]; |
| } |
| } |
| |
| // TODO(laurenzo): We need to figure out how to propagate dtype in the |
| // buffer view. |
| return py::buffer_info( |
| mapped_memory_.contents.data, element_size, |
| py::format_descriptor<float>::format(), // TODO(laurenzo): DTYPE! |
| shape.size(), dims, strides); |
| } |
| |
| private: |
| iree_hal_mapped_memory_t mapped_memory_; |
| iree_hal_buffer_view_t* bv_; |
| }; |
| |
| } // namespace |
| |
| //------------------------------------------------------------------------------ |
| // HalDriver |
| //------------------------------------------------------------------------------ |
| |
| std::vector<std::string> HalDriver::Query() { |
| iree_hal_driver_info_t* driver_infos = NULL; |
| iree_host_size_t driver_info_count = 0; |
| CheckApiStatus( |
| iree_hal_driver_registry_enumerate(iree_hal_driver_registry_default(), |
| iree_allocator_system(), &driver_infos, |
| &driver_info_count), |
| "Error enumerating HAL drivers"); |
| std::vector<std::string> driver_names(driver_info_count); |
| for (iree_host_size_t i = 0; i < driver_info_count; ++i) { |
| driver_names[i] = std::string(driver_infos[i].driver_name.data, |
| driver_infos[i].driver_name.size); |
| } |
| iree_allocator_system_free(NULL, driver_infos); |
| return driver_names; |
| } |
| |
| HalDriver HalDriver::Create(const std::string& driver_name) { |
| iree_hal_driver_t* driver; |
| CheckApiStatus(iree_hal_driver_registry_try_create_by_name( |
| iree_hal_driver_registry_default(), |
| {driver_name.data(), driver_name.size()}, |
| iree_allocator_system(), &driver), |
| "Error creating driver"); |
| return HalDriver::CreateRetained(driver); |
| } |
| |
| HalDevice HalDriver::CreateDefaultDevice() { |
| iree_hal_device_t* device; |
| CheckApiStatus(iree_hal_driver_create_default_device( |
| raw_ptr(), iree_allocator_system(), &device), |
| "Error creating default device"); |
| return HalDevice::CreateRetained(device); |
| } |
| |
| void SetupHalBindings(pybind11::module m) { |
| // Enums. |
| py::enum_<enum iree_hal_memory_type_e>(m, "MemoryType") |
| .value("NONE", IREE_HAL_MEMORY_TYPE_NONE) |
| .value("TRANSIENT", IREE_HAL_MEMORY_TYPE_TRANSIENT) |
| .value("HOST_VISIBLE", IREE_HAL_MEMORY_TYPE_HOST_VISIBLE) |
| .value("HOST_COHERENT", IREE_HAL_MEMORY_TYPE_HOST_COHERENT) |
| .value("HOST_CACHED", IREE_HAL_MEMORY_TYPE_HOST_CACHED) |
| .value("HOST_LOCAL", IREE_HAL_MEMORY_TYPE_HOST_LOCAL) |
| .value("DEVICE_VISIBLE", IREE_HAL_MEMORY_TYPE_HOST_VISIBLE) |
| .value("DEVICE_LOCAL", IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL) |
| .export_values(); |
| py::enum_<enum iree_hal_buffer_usage_e>(m, "BufferUsage") |
| .value("NONE", IREE_HAL_BUFFER_USAGE_NONE) |
| .value("CONSTANT", IREE_HAL_BUFFER_USAGE_CONSTANT) |
| .value("TRANSFER", IREE_HAL_BUFFER_USAGE_TRANSFER) |
| .value("MAPPING", IREE_HAL_BUFFER_USAGE_MAPPING) |
| .value("DISPATCH", IREE_HAL_BUFFER_USAGE_DISPATCH) |
| .value("ALL", IREE_HAL_BUFFER_USAGE_ALL) |
| .export_values(); |
| py::enum_<enum iree_hal_memory_access_e>(m, "MemoryAccess") |
| .value("NONE", IREE_HAL_MEMORY_ACCESS_NONE) |
| .value("READ", IREE_HAL_MEMORY_ACCESS_READ) |
| .value("WRITE", IREE_HAL_MEMORY_ACCESS_WRITE) |
| .value("DISCARD", IREE_HAL_MEMORY_ACCESS_DISCARD) |
| .value("DISCARD_WRITE", IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE) |
| .value("ALL", IREE_HAL_MEMORY_ACCESS_ALL) |
| .export_values(); |
| |
| py::class_<HalDevice>(m, "HalDevice"); |
| py::class_<HalDriver>(m, "HalDriver") |
| .def_static("query", &HalDriver::Query) |
| .def_static("create", &HalDriver::Create, py::arg("driver_name")) |
| .def("create_default_device", &HalDriver::CreateDefaultDevice); |
| |
| py::class_<HalShape>(m, "Shape").def(py::init(&HalShape::FromIntVector)); |
| py::class_<HalBufferView>(m, "BufferView") |
| .def("map", HalMappedMemory::Create); |
| py::class_<HalMappedMemory>(m, "MappedMemory", py::buffer_protocol()) |
| .def_buffer(&HalMappedMemory::ToBufferInfo); |
| py::class_<HalBuffer>(m, "HalBuffer") |
| .def_static("allocate_heap", &HalBuffer::AllocateHeapBuffer, |
| py::arg("memory_type"), py::arg("usage"), |
| py::arg("allocation_size")) |
| .def("fill_zero", &HalBuffer::FillZero, py::arg("byte_offset"), |
| py::arg("byte_length")) |
| .def("create_view", &HalBuffer::CreateView, py::arg("shape"), |
| py::arg("element_size")); |
| } |
| |
| } // namespace python |
| } // namespace iree |
