// Copyright 2019 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

#include "./hal.h"

#include "./numpy_interop.h"
#include "./vm.h"
#include "iree/base/internal/path.h"
#include "iree/hal/api.h"
#include "iree/hal/utils/allocators.h"
#include "iree/modules/hal/module.h"

namespace iree {
namespace python {

namespace {

// RAII wrapper for a Py_buffer which calls PyBuffer_Release when it goes
// out of scope.
class PyBufferReleaser {
 public:
  PyBufferReleaser(Py_buffer& b) : b_(b) {}
  ~PyBufferReleaser() { PyBuffer_Release(&b_); }

 private:
  Py_buffer& b_;
};

static std::string ToHexString(const uint8_t* data, size_t length) {
  static constexpr char kHexChars[] = {'0', '1', '2', '3', '4', '5', '6', '7',
                                       '8', '9', 'a', 'b', 'c', 'd', 'e', 'f'};
  std::string s(length * 2, ' ');
  for (size_t i = 0; i < length; ++i) {
    s[2 * i + 0] = kHexChars[(data[i] & 0xF0) >> 4];
    s[2 * i + 1] = kHexChars[(data[i] & 0x0F) >> 0];
  }
  return s;
}
static std::string ToHexString(uint32_t value) {
  return ToHexString((const uint8_t*)&value, sizeof(value));
}

}  // namespace

//------------------------------------------------------------------------------
// HalAllocator
//------------------------------------------------------------------------------

py::dict HalAllocator::QueryStatistics() {
  py::dict items;
  iree_hal_allocator_statistics_t stats;
  iree_hal_allocator_query_statistics(raw_ptr(), &stats);
#if IREE_STATISTICS_ENABLE
  items["host_bytes_peak"] = stats.host_bytes_peak;
  items["host_bytes_allocated"] = stats.host_bytes_allocated;
  items["host_bytes_freed"] = stats.host_bytes_freed;
  items["device_bytes_peak"] = stats.device_bytes_peak;
  items["device_bytes_allocated"] = stats.device_bytes_allocated;
  items["device_bytes_freed"] = stats.device_bytes_freed;
#endif
  return items;
}

py::str HalAllocator::FormattedStatistics() {
  // Perform all allocating string manipulation without early exit.
  iree_string_builder_t builder;
  iree_string_builder_initialize(iree_allocator_system(), &builder);
  iree_hal_allocator_statistics_t stats;
  iree_hal_allocator_query_statistics(raw_ptr(), &stats);
  auto status = iree_hal_allocator_statistics_format(&stats, &builder);
  iree_string_view_t view = iree_string_builder_view(&builder);
  py::str result = py::str(view.data, view.size);
  iree_string_builder_deinitialize(&builder);

  // Check/raise after all memory alloc/dealloc.
  CheckApiStatus(status, "unable to format statistics");
  return result;
}

py::object HalAllocator::AllocateBufferCopy(
    int memory_type, int allowed_usage, py::object buffer,
    std::optional<iree_hal_element_types_t> element_type) {
  IREE_TRACE_SCOPE_NAMED("HalAllocator::AllocateBufferCopy");
  // Request a view of the buffer (use the raw python C API to avoid
  // some allocation and copying at the pybind level).
  Py_buffer py_view;
  // Note that only C-Contiguous ND-arrays are presently supported, so
  // only request that via PyBUF_ND. Long term, we should consult an
  // "oracle" in the runtime to determine the precise required format
  // and set flags accordingly (and fallback/copy on failure).
  int flags = PyBUF_FORMAT | PyBUF_ND;

  // Acquire the backing buffer and setup RAII release.
  if (PyObject_GetBuffer(buffer.ptr(), &py_view, flags) != 0) {
    // The GetBuffer call is required to set an appropriate error.
    throw py::python_error();
  }
  PyBufferReleaser py_view_releaser(py_view);

  iree_hal_buffer_params_t params = {0};
  // TODO: Should not require host visible :(
  params.type = memory_type | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE;
  params.usage = allowed_usage;

  iree_hal_buffer_t* hal_buffer = nullptr;
  iree_status_t status = iree_ok_status();
  {
    py::gil_scoped_release release;
    status = iree_hal_allocator_allocate_buffer(
        raw_ptr(), params, py_view.len,
        iree_make_const_byte_span(py_view.buf, py_view.len), &hal_buffer);
  }
  CheckApiStatus(status, "Failed to allocate device visible buffer");

  if (!element_type) {
    return py::cast(HalBuffer::StealFromRawPtr(hal_buffer),
                    py::rv_policy::move);
  }

  // Create the buffer_view. (note that numpy shape is ssize_t, so we need to
  // copy).
  iree_hal_encoding_type_t encoding_type =
      IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR;
  std::vector<iree_hal_dim_t> dims(py_view.ndim);
  std::copy(py_view.shape, py_view.shape + py_view.ndim, dims.begin());
  iree_hal_buffer_view_t* hal_buffer_view;
  CheckApiStatus(
      iree_hal_buffer_view_create(
          hal_buffer, dims.size(), dims.data(), *element_type, encoding_type,
          iree_hal_allocator_host_allocator(raw_ptr()), &hal_buffer_view),
      "Error allocating buffer_view");
  iree_hal_buffer_release(hal_buffer);

  return py::cast(HalBufferView::StealFromRawPtr(hal_buffer_view),
                  py::rv_policy::move);
}

//------------------------------------------------------------------------------
// HalBuffer
//------------------------------------------------------------------------------

namespace {

void AppendHalBufferRepr(iree_hal_buffer_t* buffer, std::string& repr) {
  repr.append(std::to_string(iree_hal_buffer_byte_length(buffer)));
  repr.append(" bytes (at offset ");
  repr.append(std::to_string(iree_hal_buffer_byte_offset(buffer)));
  repr.append(" into ");
  repr.append(std::to_string(iree_hal_buffer_allocation_size(buffer)));
  repr.append("), memory_type=");

  // Memory type.
  iree_bitfield_string_temp_t tmp;
  iree_string_view_t sv;
  sv = iree_hal_memory_type_format(iree_hal_buffer_memory_type(buffer), &tmp);
  repr.append(sv.data, sv.size);

  // Allowed access.
  repr.append(", allowed_access=");
  sv = iree_hal_memory_access_format(iree_hal_buffer_allowed_access(buffer),
                                     &tmp);
  repr.append(sv.data, sv.size);

  // Allowed usage.
  repr.append(", allowed_usage=");
  sv =
      iree_hal_buffer_usage_format(iree_hal_buffer_allowed_usage(buffer), &tmp);
  repr.append(sv.data, sv.size);
}

}  // namespace

py::str HalBuffer::Repr() {
  std::string repr("<HalBuffer ");
  AppendHalBufferRepr(raw_ptr(), repr);
  repr.append(">");
  return py::str(py::cast(repr));
}

//------------------------------------------------------------------------------
// HalBufferView
//------------------------------------------------------------------------------

py::str HalBufferView::Repr() {
  std::string repr("<HalBufferView (");

  // Shape.
  iree_host_size_t rank = iree_hal_buffer_view_shape_rank(raw_ptr());
  for (iree_host_size_t i = 0; i < rank; ++i) {
    if (i > 0) {
      repr.append(", ");
    }
    repr.append(std::to_string(iree_hal_buffer_view_shape_dim(raw_ptr(), i)));
  }
  repr.append(")");

  // Element type.
  repr.append(", element_type=0x");
  auto element_type = iree_hal_buffer_view_element_type(raw_ptr());
  repr.append(ToHexString(static_cast<uint32_t>(element_type)));

  repr.append(", ");
  AppendHalBufferRepr(iree_hal_buffer_view_buffer(raw_ptr()), repr);
  repr.append(">");
  return py::str(py::cast(repr));
}

//------------------------------------------------------------------------------
// HalDevice
//------------------------------------------------------------------------------

void HalDevice::BeginProfiling(std::optional<std::string> mode,
                               std::optional<std::string> file_path) {
  iree_hal_device_profiling_options_t options;
  memset(&options, 0, sizeof(options));

  options.mode = IREE_HAL_DEVICE_PROFILING_MODE_QUEUE_OPERATIONS;
  if (mode) {
    if (*mode == "queue") {
      options.mode = IREE_HAL_DEVICE_PROFILING_MODE_QUEUE_OPERATIONS;
    } else if (*mode == "dispatch") {
      options.mode = IREE_HAL_DEVICE_PROFILING_MODE_DISPATCH_COUNTERS;
    } else if (*mode == "executable") {
      options.mode = IREE_HAL_DEVICE_PROFILING_MODE_EXECUTABLE_COUNTERS;
    } else {
      throw RaiseValueError("unrecognized profiling mode");
    }
  }

  options.file_path = file_path ? file_path->c_str() : nullptr;
  CheckApiStatus(iree_hal_device_profiling_begin(raw_ptr(), &options),
                 "starting device profiling");
}

void HalDevice::EndProfiling() {
  CheckApiStatus(iree_hal_device_profiling_end(raw_ptr()),
                 "ending device profiling");
}

//------------------------------------------------------------------------------
// HalDriver
//------------------------------------------------------------------------------

std::vector<std::string> HalDriver::Query() {
  iree_host_size_t driver_info_count = 0;
  iree_hal_driver_info_t* driver_infos = NULL;
  CheckApiStatus(
      iree_hal_driver_registry_enumerate(iree_hal_driver_registry_default(),
                                         iree_allocator_system(),
                                         &driver_info_count, &driver_infos),
      "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_free(iree_allocator_system(), driver_infos);
  return driver_names;
}

py::object HalDriver::Create(const std::string& device_uri,
                             py::dict& driver_cache) {
  iree_string_view_t driver_name, device_path, params_str;
  iree_string_view_t device_uri_sv{
      device_uri.data(), static_cast<iree_host_size_t>(device_uri.size())};
  iree_uri_split(device_uri_sv, &driver_name, &device_path, &params_str);

  // Check cache.
  py::str cache_key(driver_name.data, driver_name.size);
  py::object cached = driver_cache.attr("get")(cache_key);
  if (!cached.is_none()) {
    return cached;
  }

  // Create.
  iree_hal_driver_t* driver;
  CheckApiStatus(iree_hal_driver_registry_try_create(
                     iree_hal_driver_registry_default(), driver_name,
                     iree_allocator_system(), &driver),
                 "Error creating driver");

  // Cache.
  py::object driver_obj = py::cast(HalDriver::StealFromRawPtr(driver));
  driver_cache[cache_key] = driver_obj;
  return driver_obj;
}

py::list HalDriver::QueryAvailableDevices() {
  iree_hal_device_info_t* device_infos;
  iree_host_size_t count;
  CheckApiStatus(iree_hal_driver_query_available_devices(
                     raw_ptr(), iree_allocator_system(), &count, &device_infos),
                 "Error querying devices");
  py::list results;
  for (iree_host_size_t i = 0; i < count; ++i) {
    py::dict device_data;
    device_data["device_id"] = py::cast(device_infos[i].device_id);
    device_data["path"] =
        py::str(device_infos[i].path.data, device_infos[i].path.size);
    device_data["name"] =
        py::str(device_infos[i].name.data, device_infos[i].name.size);
    results.append(device_data);
  }

  iree_allocator_free(iree_allocator_system(), device_infos);
  return results;
}

// Configures |device| based on flags before returning it to the user.
static iree_status_t ConfigureDevice(iree_hal_device_t* device,
                                     std::optional<py::list> allocators) {
  // Optionally wrap the base device allocator with caching/pooling.
  // Doing this here satisfies the requirement that no buffers have been
  // allocated yet - if we returned the device without doing this the caller
  // can more easily break the rules.
  if (allocators) {
    // NOTE: we need to pass string views that point to the std::string storage.
    // We do that in two passes because as we grow spec_storage it may
    // reallocate itself and invalidate the pointers - only after we're done
    // can we capture them in views.
    auto& spec_list = *allocators;
    std::vector<std::string> spec_storage;
    spec_storage.reserve(spec_list.size());
    for (auto item : spec_list) {
      auto spec = py::cast<std::string>(item);
      spec_storage.push_back(std::move(spec));
    }
    std::vector<iree_string_view_t> spec_views;
    spec_views.reserve(spec_list.size());
    for (const auto& spec : spec_storage) {
      spec_views.push_back(iree_make_string_view(spec.data(), spec.size()));
    }
    IREE_RETURN_IF_ERROR(iree_hal_configure_allocator_from_specs(
        spec_views.size(), spec_views.data(), device));
  }
  return iree_ok_status();
}

HalDevice HalDriver::CreateDefaultDevice(std::optional<py::list> allocators) {
  iree_hal_device_t* device;
  CheckApiStatus(iree_hal_driver_create_default_device(
                     raw_ptr(), iree_allocator_system(), &device),
                 "Error creating default device");
  CheckApiStatus(ConfigureDevice(device, allocators),
                 "Error configuring the device");
  return HalDevice::StealFromRawPtr(device);
}

HalDevice HalDriver::CreateDevice(iree_hal_device_id_t device_id,
                                  std::optional<py::list> allocators) {
  // Since the device ids are supposed to be opaque, we need to verify
  // them by querying available devices.
  py::list available_devices = QueryAvailableDevices();
  bool found = false;
  py::object compare_device_id = py::cast(device_id);
  for (auto record : available_devices) {
    // Each record is a dict:
    // {"device_id": obj, "path": str, "name": str}.
    auto record_dict = py::cast<py::dict>(record);
    py::object found_device_id = record_dict["device_id"];
    if (found_device_id.is(compare_device_id)) {
      found = true;
      break;
    }
  }

  if (!found) {
    std::string msg;
    msg.append("Device id ");
    msg.append(std::to_string(device_id));
    msg.append(" not found. Available devices: ");
    msg.append(py::cast<std::string>(py::repr(available_devices)));
    throw std::invalid_argument(std::move(msg));
  }

  std::vector<iree_string_pair_t> params;
  iree_hal_device_t* device;
  CheckApiStatus(iree_hal_driver_create_device_by_id(
                     raw_ptr(), device_id, params.size(), &params.front(),
                     iree_allocator_system(), &device),
                 "Error creating default device");
  CheckApiStatus(ConfigureDevice(device, allocators),
                 "Error configuring the device");
  return HalDevice::StealFromRawPtr(device);
}

HalDevice HalDriver::CreateDeviceByURI(std::string& device_uri,
                                       std::optional<py::list> allocators) {
  iree_hal_device_t* device;
  iree_string_view_t device_uri_sv{
      device_uri.data(), static_cast<iree_host_size_t>(device_uri.size())};
  CheckApiStatus(
      iree_hal_driver_create_device_by_uri(raw_ptr(), device_uri_sv,
                                           iree_allocator_system(), &device),
      "Error creating device");
  CheckApiStatus(ConfigureDevice(device, allocators),
                 "Error configuring the device");
  return HalDevice::StealFromRawPtr(device);
}

//------------------------------------------------------------------------------
// HAL module
//------------------------------------------------------------------------------

VmModule CreateHalModule(VmInstance* instance, HalDevice* device) {
  iree_vm_module_t* module = NULL;
  CheckApiStatus(iree_hal_module_create(instance->raw_ptr(), device->raw_ptr(),
                                        IREE_HAL_MODULE_FLAG_NONE,
                                        iree_allocator_system(), &module),
                 "Error creating hal module");
  return VmModule::StealFromRawPtr(module);
}

//------------------------------------------------------------------------------
// Bindings
//------------------------------------------------------------------------------

void SetupHalBindings(nanobind::module_ m) {
  py::dict driver_cache;

  // Built-in module creation.
  m.def("create_hal_module", &CreateHalModule);

  // Enums.
  py::enum_<enum iree_hal_memory_type_bits_t>(m, "MemoryType")
      .value("NONE", IREE_HAL_MEMORY_TYPE_NONE)
      .value("OPTIMAL", IREE_HAL_MEMORY_TYPE_OPTIMAL)
      .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_DEVICE_VISIBLE)
      .value("DEVICE_LOCAL", IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)
      .export_values()
      .def("__or__",
           [](enum iree_hal_memory_type_bits_t self,
              enum iree_hal_memory_type_bits_t other) { return self | other; })
      .def("__and__",
           [](enum iree_hal_memory_type_bits_t self,
              enum iree_hal_memory_type_bits_t other) { return self & other; });

  py::enum_<enum iree_hal_buffer_compatibility_bits_t>(m, "BufferCompatibility")
      .value("NONE", IREE_HAL_BUFFER_COMPATIBILITY_NONE)
      .value("ALLOCATABLE", IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE)
      .value("IMPORTABLE", IREE_HAL_BUFFER_COMPATIBILITY_IMPORTABLE)
      .value("EXPORTABLE", IREE_HAL_BUFFER_COMPATIBILITY_EXPORTABLE)
      .value("QUEUE_TRANSFER", IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER)
      .value("QUEUE_DISPATCH", IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_DISPATCH)
      .export_values()
      .def("__or__",
           [](enum iree_hal_buffer_compatibility_bits_t self,
              enum iree_hal_buffer_compatibility_bits_t other) {
             return self | other;
           })
      .def("__and__", [](enum iree_hal_buffer_compatibility_bits_t self,
                         enum iree_hal_buffer_compatibility_bits_t other) {
        return self & other;
      });

  py::enum_<enum iree_hal_buffer_usage_bits_t>(m, "BufferUsage")
      .value("NONE", IREE_HAL_BUFFER_USAGE_NONE)
      .value("TRANSFER_SOURCE", IREE_HAL_BUFFER_USAGE_TRANSFER_SOURCE)
      .value("TRANSFER_TARGET", IREE_HAL_BUFFER_USAGE_TRANSFER_TARGET)
      .value("TRANSFER", IREE_HAL_BUFFER_USAGE_TRANSFER)
      .value("DISPATCH_INDIRECT_PARAMS",
             IREE_HAL_BUFFER_USAGE_DISPATCH_INDIRECT_PARAMS)
      .value("DISPATCH_UNIFORM_READ",
             IREE_HAL_BUFFER_USAGE_DISPATCH_UNIFORM_READ)
      .value("DISPATCH_STORAGE_READ",
             IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE_READ)
      .value("DISPATCH_STORAGE_WRITE",
             IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE_WRITE)
      .value("DISPATCH_STORAGE", IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE)
      .value("DISPATCH_IMAGE_READ", IREE_HAL_BUFFER_USAGE_DISPATCH_IMAGE_READ)
      .value("DISPATCH_IMAGE_WRITE", IREE_HAL_BUFFER_USAGE_DISPATCH_IMAGE_WRITE)
      .value("DISPATCH_IMAGE", IREE_HAL_BUFFER_USAGE_DISPATCH_IMAGE)
      .value("SHARING_EXPORT", IREE_HAL_BUFFER_USAGE_SHARING_EXPORT)
      .value("SHARING_REPLICATE", IREE_HAL_BUFFER_USAGE_SHARING_REPLICATE)
      .value("SHARING_CONCURRENT", IREE_HAL_BUFFER_USAGE_SHARING_CONCURRENT)
      .value("SHARING_IMMUTABLE", IREE_HAL_BUFFER_USAGE_SHARING_IMMUTABLE)
      .value("MAPPING_SCOPED", IREE_HAL_BUFFER_USAGE_MAPPING_SCOPED)
      .value("MAPPING_PERSISTENT", IREE_HAL_BUFFER_USAGE_MAPPING_PERSISTENT)
      .value("MAPPING_OPTIONAL", IREE_HAL_BUFFER_USAGE_MAPPING_OPTIONAL)
      .value("MAPPING_ACCESS_RANDOM",
             IREE_HAL_BUFFER_USAGE_MAPPING_ACCESS_RANDOM)
      .value("MAPPING_ACCESS_SEQUENTIAL_WRITE",
             IREE_HAL_BUFFER_USAGE_MAPPING_ACCESS_SEQUENTIAL_WRITE)
      .value("MAPPING", IREE_HAL_BUFFER_USAGE_MAPPING)
      .value("DEFAULT", IREE_HAL_BUFFER_USAGE_DEFAULT)
      .export_values()
      .def("__or__",
           [](enum iree_hal_buffer_usage_bits_t self,
              enum iree_hal_buffer_usage_bits_t other) {
             return (enum iree_hal_buffer_usage_bits_t)(self | other);
           })
      .def("__and__", [](enum iree_hal_buffer_usage_bits_t self,
                         enum iree_hal_buffer_usage_bits_t other) {
        return (enum iree_hal_buffer_usage_bits_t)(self & other);
      });

  py::enum_<enum iree_hal_memory_access_bits_t>(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()
      .def(
          "__or__",
          [](enum iree_hal_memory_access_bits_t self,
             enum iree_hal_memory_access_bits_t other) { return self | other; })
      .def("__and__", [](enum iree_hal_memory_access_bits_t self,
                         enum iree_hal_memory_access_bits_t other) {
        return self & other;
      });

  py::enum_<enum iree_hal_element_types_t>(m, "HalElementType")
      .value("NONE", IREE_HAL_ELEMENT_TYPE_NONE)
      .value("OPAQUE_8", IREE_HAL_ELEMENT_TYPE_OPAQUE_8)
      .value("OPAQUE_16", IREE_HAL_ELEMENT_TYPE_OPAQUE_16)
      .value("OPAQUE_32", IREE_HAL_ELEMENT_TYPE_OPAQUE_32)
      .value("OPAQUE_64", IREE_HAL_ELEMENT_TYPE_OPAQUE_64)
      .value("BOOL_8", IREE_HAL_ELEMENT_TYPE_BOOL_8)
      .value("INT_4", IREE_HAL_ELEMENT_TYPE_INT_4)
      .value("INT_8", IREE_HAL_ELEMENT_TYPE_INT_8)
      .value("INT_16", IREE_HAL_ELEMENT_TYPE_INT_16)
      .value("INT_32", IREE_HAL_ELEMENT_TYPE_INT_32)
      .value("INT_64", IREE_HAL_ELEMENT_TYPE_INT_64)
      .value("SINT_4", IREE_HAL_ELEMENT_TYPE_SINT_4)
      .value("SINT_8", IREE_HAL_ELEMENT_TYPE_SINT_8)
      .value("SINT_16", IREE_HAL_ELEMENT_TYPE_SINT_16)
      .value("SINT_32", IREE_HAL_ELEMENT_TYPE_SINT_32)
      .value("SINT_64", IREE_HAL_ELEMENT_TYPE_SINT_64)
      .value("UINT_4", IREE_HAL_ELEMENT_TYPE_UINT_4)
      .value("UINT_8", IREE_HAL_ELEMENT_TYPE_UINT_8)
      .value("UINT_16", IREE_HAL_ELEMENT_TYPE_UINT_16)
      .value("UINT_32", IREE_HAL_ELEMENT_TYPE_UINT_32)
      .value("UINT_64", IREE_HAL_ELEMENT_TYPE_UINT_64)
      .value("FLOAT_16", IREE_HAL_ELEMENT_TYPE_FLOAT_16)
      .value("FLOAT_32", IREE_HAL_ELEMENT_TYPE_FLOAT_32)
      .value("FLOAT_64", IREE_HAL_ELEMENT_TYPE_FLOAT_64)
      .value("BFLOAT_16", IREE_HAL_ELEMENT_TYPE_BFLOAT_16)
      .value("COMPLEX_64", IREE_HAL_ELEMENT_TYPE_COMPLEX_FLOAT_64)
      .value("COMPLEX_128", IREE_HAL_ELEMENT_TYPE_COMPLEX_FLOAT_128)
      .export_values()
      .def_static("map_to_dtype", [](iree_hal_element_type_t element_type) {
        int typenum = numpy::ConvertHalElementTypeToNumPyTypeNum(element_type);
        return numpy::DescrNewFromType(typenum);
      });

  py::class_<HalDevice>(m, "HalDevice")
      .def_prop_ro(
          "allocator",
          [](HalDevice& self) {
            return HalAllocator::BorrowFromRawPtr(self.allocator());
          },
          py::keep_alive<0, 1>())
      .def("begin_profiling", &HalDevice::BeginProfiling,
           py::arg("mode") = py::none(), py::arg("file_path") = py::none())
      .def("end_profiling", &HalDevice::EndProfiling);

  py::class_<HalDriver>(m, "HalDriver")
      .def_static("query", &HalDriver::Query)

      // All 'create_device' functions take optional kwargs that should be kept
      // in sync.
      .def("create_default_device", &HalDriver::CreateDefaultDevice,
           py::keep_alive<0, 1>(), py::arg("allocators") = py::none())
      .def("create_device", &HalDriver::CreateDevice, py::keep_alive<0, 1>(),
           py::arg("device_id"), py::arg("allocators") = py::none())
      .def("create_device_by_uri", &HalDriver::CreateDeviceByURI,
           py::keep_alive<0, 1>(), py::arg("device_uri"),
           py::arg("allocators") = py::none())
      .def(
          "create_device",
          [](HalDriver& self, py::dict device_info,
             std::optional<py::list> allocators) -> HalDevice {
            // Alias of create_device that takes a dict as returned from
            // query_available_devices for convenience.
            auto device_id =
                py::cast<iree_hal_device_id_t>(device_info["device_id"]);
            return self.CreateDevice(device_id, allocators);
          },
          py::keep_alive<0, 1>(), py::arg("device_info"),
          py::arg("allocators") = py::none())
      .def("query_available_devices", &HalDriver::QueryAvailableDevices);

  m.def(
      "get_cached_hal_driver",
      [driver_cache](std::string device_uri) {
        return HalDriver::Create(device_uri,
                                 const_cast<py::dict&>(driver_cache));
      },
      py::arg("device_uri"));

  py::class_<HalAllocator>(m, "HalAllocator")
      .def("trim",
           [](HalAllocator& self) {
             CheckApiStatus(iree_hal_allocator_trim(self.raw_ptr()),
                            "Error trim()'ing HAL allocator");
           })
      .def_prop_ro(
          "has_statistics",
          [](HalAllocator& self) -> bool { return IREE_STATISTICS_ENABLE; })
      .def_prop_ro("statistics", &HalAllocator::QueryStatistics)
      .def_prop_ro("formatted_statistics", &HalAllocator::FormattedStatistics)
      .def(
          "query_buffer_compatibility",
          [](HalAllocator& self, int memory_type, int allowed_usage,
             int intended_usage, iree_device_size_t allocation_size) -> int {
            iree_hal_buffer_params_t params = {0};
            params.type = memory_type;
            params.usage = allowed_usage & intended_usage;
            return iree_hal_allocator_query_buffer_compatibility(
                self.raw_ptr(), params, allocation_size,
                /*out_params=*/nullptr, /*out_allocation_size=*/0);
          },
          py::arg("memory_type"), py::arg("allowed_usage"),
          py::arg("intended_usage"), py::arg("allocation_size"))
      .def(
          "allocate_buffer",
          [](HalAllocator& self, int memory_type, int allowed_usage,
             iree_device_size_t allocation_size) {
            iree_hal_buffer_params_t params = {0};
            params.type = memory_type;
            params.usage = allowed_usage;
            iree_hal_buffer_t* buffer = nullptr;
            iree_const_byte_span_t empty_initial_data{nullptr, 0};
            CheckApiStatus(iree_hal_allocator_allocate_buffer(
                               self.raw_ptr(), params, allocation_size,
                               empty_initial_data, &buffer),
                           "could not allocate buffer");
            return HalBuffer::StealFromRawPtr(buffer);
          },
          py::arg("memory_type"), py::arg("allowed_usage"),
          py::arg("allocation_size"), py::keep_alive<0, 1>(),
          "Allocates a new buffer with requested characteristics (does not "
          "initialize with specific data).")
      .def("allocate_buffer_copy", &HalAllocator::AllocateBufferCopy,
           py::arg("memory_type"), py::arg("allowed_usage"), py::arg("buffer"),
           py::arg("element_type") = py::none(), py::keep_alive<0, 1>(),
           "Allocates a new buffer and initializes it from a Python buffer "
           "object. If an element type is specified, wraps in a BufferView "
           "matching the characteristics of the Python buffer. The format is "
           "requested as ND/C-Contiguous, which may incur copies if not "
           "already in that format.");

  py::class_<HalBuffer>(m, "HalBuffer")
      .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"), py::keep_alive<0, 1>())
      .def("__repr__", &HalBuffer::Repr);

  auto hal_buffer_view = py::class_<HalBufferView>(m, "HalBufferView");
  VmRef::BindRefProtocol(hal_buffer_view, iree_hal_buffer_view_type,
                         iree_hal_buffer_view_retain_ref,
                         iree_hal_buffer_view_deref, iree_hal_buffer_view_isa);
  hal_buffer_view.def("map", HalMappedMemory::Create, py::keep_alive<0, 1>())
      .def_prop_ro("shape",
                   [](HalBufferView& self) {
                     iree_host_size_t rank =
                         iree_hal_buffer_view_shape_rank(self.raw_ptr());
                     auto* dims =
                         iree_hal_buffer_view_shape_dims(self.raw_ptr());
                     py::list result;
                     for (iree_host_size_t i = 0; i < rank; ++i) {
                       result.append(dims[i]);
                     }
                     return result;
                   })
      .def_prop_ro("element_type",
                   [](HalBufferView& self) {
                     return iree_hal_buffer_view_element_type(self.raw_ptr());
                   })
      .def("__repr__", &HalBufferView::Repr);

  py::class_<HalMappedMemory>(m, "MappedMemory")
      .def(
          "asarray",
          [](HalMappedMemory* self, std::vector<iree_host_size_t> shape,
             py::object dtype) {
            py::object py_mapped_memory = py::cast(self);
            static_assert(sizeof(shape[0]) == sizeof(intptr_t),
                          "size_t not of same size as intptr_t");
            int typenum = numpy::TypenumFromDescr(dtype);
            return numpy::SimpleNewFromData(
                shape.size(), reinterpret_cast<intptr_t const*>(shape.data()),
                typenum, self->mapped_memory().contents.data, py_mapped_memory);
          },
          py::arg("shape"), py::arg("element_type"));

  py::class_<HalShape>(m, "Shape")
      .def("__init__", [](HalShape* self, std::vector<iree_hal_dim_t> indices) {
        new (self) HalShape(indices);
      });
}

}  // namespace python
}  // namespace iree