iree/runtime/instance.c

// Copyright 2021 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 "iree/runtime/instance.h"

#include <stddef.h>
#include <string.h>

#include "iree/base/internal/atomics.h"
#include "iree/base/tracing.h"
#include "iree/hal/api.h"
#include "iree/hal/drivers/init.h"
#include "iree/modules/hal/module.h"
#include "iree/vm/api.h"

//===----------------------------------------------------------------------===//
// iree_runtime_instance_options_t
//===----------------------------------------------------------------------===//

IREE_API_EXPORT void iree_runtime_instance_options_initialize(
    iree_api_version_t api_version,
    iree_runtime_instance_options_t* out_options) {
  memset(out_options, 0, sizeof(*out_options));
  out_options->api_version = api_version;
}

IREE_API_EXPORT void iree_runtime_instance_options_use_all_available_drivers(
    iree_runtime_instance_options_t* options) {
  options->driver_registry = iree_hal_driver_registry_default();
  // TODO(benvanik): remove status result from this; it can't (meaningfully)
  // fail and is just extra bookkeeping.
  iree_status_ignore(
      iree_hal_register_all_available_drivers(options->driver_registry));
}

//===----------------------------------------------------------------------===//
// iree_runtime_instance_t
//===----------------------------------------------------------------------===//

struct iree_runtime_instance_t {
  iree_atomic_ref_count_t ref_count;

  // Allocator used to allocate the instance and all of its resources.
  iree_allocator_t host_allocator;

  // An optional driver registry used to enumerate and create HAL devices.
  iree_hal_driver_registry_t* driver_registry;

  // TODO(#5724): we should have a device cache here so that multiple sessions
  // can find the same devices. This may mean a new HAL type like
  // iree_hal_device_pool_t to prevent too much coupling and make weak
  // references easier.
};

IREE_API_EXPORT iree_status_t iree_runtime_instance_create(
    const iree_runtime_instance_options_t* options,
    iree_allocator_t host_allocator, iree_runtime_instance_t** out_instance) {
  IREE_ASSERT_ARGUMENT(options);
  IREE_ASSERT_ARGUMENT(out_instance);
  *out_instance = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  // Check that the API version matches what the runtime expects. The check here
  // should always succeed when the runtime and the underlying system are linked
  // together into the same binary.
  iree_api_version_t actual_version = IREE_API_VERSION_0;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_api_version_check(options->api_version, &actual_version));

  // Register builtin types.
  // TODO(benvanik): change to per-instance type registries to avoid these
  // global (UNSAFE!) calls. For now hosting applications should really only
  // be using a single instance anyway.
  IREE_RETURN_AND_END_ZONE_IF_ERROR(z0, iree_vm_register_builtin_types());
  IREE_RETURN_AND_END_ZONE_IF_ERROR(z0, iree_hal_module_register_types());

  // Allocate the instance state.
  iree_runtime_instance_t* instance = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_allocator_malloc(host_allocator, sizeof(*instance),
                                (void**)&instance));
  instance->host_allocator = host_allocator;
  iree_atomic_ref_count_init(&instance->ref_count);

  instance->driver_registry = options->driver_registry;
  // TODO(benvanik): driver registry ref counting.

  *out_instance = instance;
  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

static void iree_runtime_instance_destroy(iree_runtime_instance_t* instance) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_allocator_free(instance->host_allocator, instance);

  IREE_TRACE_ZONE_END(z0);
}

IREE_API_EXPORT void iree_runtime_instance_retain(
    iree_runtime_instance_t* instance) {
  if (instance) {
    iree_atomic_ref_count_inc(&instance->ref_count);
  }
}

IREE_API_EXPORT void iree_runtime_instance_release(
    iree_runtime_instance_t* instance) {
  if (instance && iree_atomic_ref_count_dec(&instance->ref_count) == 1) {
    iree_runtime_instance_destroy(instance);
  }
}

IREE_API_EXPORT iree_allocator_t
iree_runtime_instance_host_allocator(const iree_runtime_instance_t* instance) {
  IREE_ASSERT_ARGUMENT(instance);
  return instance->host_allocator;
}

IREE_API_EXPORT iree_hal_driver_registry_t*
iree_runtime_instance_driver_registry(const iree_runtime_instance_t* instance) {
  IREE_ASSERT_ARGUMENT(instance);
  return instance->driver_registry;
}

IREE_API_EXPORT iree_status_t iree_runtime_instance_try_create_default_device(
    iree_runtime_instance_t* instance, iree_string_view_t driver_name,
    iree_hal_device_t** out_device) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_ASSERT_ARGUMENT(out_device);
  *out_device = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_TEXT(z0, driver_name.data, driver_name.size);

  // This is only supported when we have a driver registry we can use to create
  // the drivers.
  iree_hal_driver_registry_t* driver_registry =
      iree_runtime_instance_driver_registry(instance);
  if (!driver_registry) {
    IREE_TRACE_ZONE_END(z0);
    return iree_make_status(IREE_STATUS_FAILED_PRECONDITION,
                            "instance was created without a driver registry "
                            "and cannot perform enumeration");
  }

  // Create a driver with the given name (if one exists).
  iree_allocator_t host_allocator =
      iree_runtime_instance_host_allocator(instance);
  iree_hal_driver_t* driver = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_hal_driver_registry_try_create_by_name(
              driver_registry, driver_name, host_allocator, &driver));

  // Create the default device on that driver.
  iree_status_t status =
      iree_hal_driver_create_default_device(driver, host_allocator, out_device);

  iree_hal_driver_release(driver);
  IREE_TRACE_ZONE_END(z0);
  return status;
}