#include <api_impl.h>

class DeviceInstance

inline DeviceInstance(int client_id, ClientInstance & client, iree_hal_driver_t * driver, iree_hal_device_info_t * info)

~DeviceInstance()

inline operator PJRT_Device *()

inline ClientInstance & client()

inline bool is_addressable()

inline int local_hardware_id()

inline void set_memory(MemoryInstance * memory)

inline MemoryInstance * memory()

iree_status_t HostBufferToDeviceZeroDim(PJRT_Buffer_Type type, const int64_t * dims, size_t num_dims, EventInstance ** out_done_with_host_buffer_event, BufferInstance ** out_buffer)

iree_status_t HostBufferToDeviceSplat(const void * data, PJRT_Buffer_Type type, const int64_t * dims, size_t num_dims, EventInstance ** out_done_with_host_buffer_event, BufferInstance ** out_buffer)

iree_status_t TransposeBroadcastDeviceBuffer(BufferInstance * buffer, iree_hal_element_type_t type, const iree_hal_dim_t * input_dims, const iree_hal_dim_t * output_dims, const int64_t * perms, size_t num_dims, PJRT_HostBufferSemantics host_buffer_semantics, EventInstance ** out_done_with_host_buffer_event, BufferInstance ** out_buffer)

iree_status_t HostBufferToDevice(const void * data, PJRT_Buffer_Type type, const int64_t * dims, size_t num_dims, const int64_t * byte_strides, size_t num_byte_strides, PJRT_HostBufferSemantics host_buffer_semantics, EventInstance ** out_done_with_host_buffer_event, BufferInstance ** out_buffer)

iree_status_t GetHalDevice(iree_hal_device_t ** out_device)

inline DeviceDescription * device_description()

inline iree_hal_semaphore_t * main_timeline()

inline iree_hal_device_t * device()

inline iree_hal_allocator_t * device_allocator()

iree_status_t CreateFence(iree_hal_fence_t ** out_fence)

static void BindApi(PJRT_Api * api)

static inline DeviceInstance * Unwrap(PJRT_Device * device)

static inline DeviceInstance * Unwrap(PJRT_DeviceDescription * device_description)

ClientInstance & client_

iree_hal_driver_t * driver_

iree::vm::ref< iree_hal_device_t > device_

iree::vm::ref< iree_hal_semaphore_t > main_timeline_

iree::vm::ref< iree_hal_semaphore_t > transfer_timeline_

iree::vm::ref< iree_hal_fence_t > transfer_now_fence_

uint64_t last_transfer_timepoint_ = 0

DeviceDescription info_

std::vector< MemoryInstance * > memories_

iree_status_t OpenDevice()

iree_status_t AcquireHostStagingBuffer(iree_const_byte_span_t initial_contents, bool snapshot_initial_contents_now, bool * initial_contents_snapshotted, iree_hal_buffer_t ** out_buffer)