IREE (OpenXLA): Ahead‑of‑Time and JIT Compilation for ML

python -m pip install --upgrade iree-compiler iree-runtime

# Sanity check: these should print help/version text
iree-compile --version
iree-run-module --help

iree-compile model.stablehlo \
  --iree-hal-target-backends=vulkan \
  --o=model_vulkan.vmfb

# List exported functions and inputs/outputs
iree-run-module --module=model_vulkan.vmfb --list-functions

# Run a function with randomly generated inputs (useful smoke test)
iree-run-module \
  --module=model_vulkan.vmfb \
  --function=main \
  --input=1x3x224x224xf32=0  # shape and dtype; values auto-filled

#include <iree/base/api.h>
#include <iree/hal/api.h>
#include <iree/hal/drivers/vulkan/api.h>
#include <iree/vm/api.h>

// Utility that checks an iree_status_t and aborts on failure.
static void IREE_CHECK_OK(iree_status_t status) {
  if (IREE_UNLIKELY(!iree_status_is_ok(status))) {
    iree_status_fprint(stderr, status);
    abort();
  }
}

int main(int argc, char** argv) {
  iree_allocator_t host_allocator = iree_allocator_system();

  // 1) VM instance
  iree_vm_instance_t* instance = NULL;
  IREE_CHECK_OK(iree_vm_instance_create(IREE_VM_TYPE_CAPACITY_DEFAULT, host_allocator, &instance));

  // 2) Create Vulkan HAL driver and device
  iree_hal_driver_t* driver = NULL;
  IREE_CHECK_OK(iree_hal_vulkan_driver_create(IREE_SV("vulkan"), host_allocator, &driver));

  iree_hal_device_t* device = NULL;
  IREE_CHECK_OK(iree_hal_driver_create_default_device(driver, host_allocator, &device));

  // 3) Load compiled module (.vmfb)
  const char* vmfb_path = (argc > 1) ? argv[1] : "model_vulkan.vmfb";
  iree_vm_module_t* bytecode_module = NULL;
  IREE_CHECK_OK(iree_vm_bytecode_module_load_from_file(
      iree_make_cstring_view(vmfb_path), IREE_VM_BYTECODE_MODULE_FLAG_NONE,
      host_allocator, host_allocator, &bytecode_module));

  // 4) HAL module gives programs access to the device
  iree_vm_module_t* hal_module = NULL;
  IREE_CHECK_OK(iree_hal_module_create(device, host_allocator, &hal_module));

  // 5) Create a context with HAL + your module
  iree_vm_module_t* modules[] = {hal_module, bytecode_module};
  iree_vm_context_t* context = NULL;
  IREE_CHECK_OK(iree_vm_context_create_with_modules(instance, IREE_ARRAYSIZE(modules),
                                                    modules, host_allocator, &context));

  // 6) Look up the exported function
  // Replace "main" with your model’s exported entry (e.g., module.function)
  iree_string_view_t entry_name = iree_make_cstring_view("main");
  iree_vm_function_t function;
  IREE_CHECK_OK(iree_vm_context_resolve_function(context, entry_name, &function));

  // 7) Prepare inputs (simple example: one 1D float tensor)
  const float input_data[4] = {1, 2, 3, 4};
  iree_hal_buffer_params_t params = {
      .type = IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE,
      .usage = IREE_HAL_BUFFER_USAGE_DEFAULT | IREE_HAL_BUFFER_USAGE_TRANSFER,
  };
  iree_hal_buffer_view_t* input = NULL;
  IREE_CHECK_OK(iree_hal_buffer_view_allocate_and_assign(
      device, /*shape_rank=*/1, (iree_hal_dim_t[]){4}, IREE_HAL_ELEMENT_TYPE_FLOAT_32,
      IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, params, sizeof(input_data),
      (const void*)input_data, &input));

  iree_vm_list_t* inputs = NULL;
  IREE_CHECK_OK(iree_vm_list_create(/*element_type=*/NULL, 1, host_allocator, &inputs));
  iree_vm_ref_t input_ref = iree_hal_buffer_view_move_ref(input);
  IREE_CHECK_OK(iree_vm_list_push_ref_move(inputs, &input_ref));

  // 8) Invoke
  iree_vm_list_t* outputs = NULL;
  IREE_CHECK_OK(iree_vm_invoke(context, function, IREE_VM_INVOCATION_FLAG_NONE,
                               /*policy=*/NULL, inputs, &outputs, host_allocator));

  // 9) Read back the first output as a buffer view
  iree_hal_buffer_view_t* out_view = NULL;
  iree_vm_ref_t out_ref = iree_vm_list_get_ref_deref(outputs, 0);
  out_view = (iree_hal_buffer_view_t*)out_ref.ptr;

  // (Optional) Map or copy to host to inspect results
  iree_host_size_t out_element_count = 0;
  IREE_CHECK_OK(iree_hal_buffer_view_compute_num_elements(out_view, &out_element_count));
  std::vector<float> host_out(out_element_count);

  iree_hal_buffer_t* out_buffer = iree_hal_buffer_view_buffer(out_view);
  IREE_CHECK_OK(iree_hal_device_transfer_d2h(
      device, out_buffer, /*source_offset=*/0,
      host_out.data(), host_out.size() * sizeof(float), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
      host_allocator));

  // Print first few values (for sanity)
  for (int i = 0; i < std::min<size_t>(host_out.size(), 4); ++i) {
    fprintf(stdout, "out[%d]=%f\n", i, host_out[i]);
  }

  // Cleanup (omitted for brevity)
  return 0;
}