Address code review issues (#26)

- Clean up constants: move kEpVendor/kEpVersion to .cc file (only used
there), delete unused VendorIds namespace, keep kEpVendorId in header
(used across 3 files).

- Defer device enumeration: move driver/device enumeration from
constructor (CreateIreeHwDevices) into GetSupportedDevicesImpl with lazy
init on first call. Create OrtHardwareDevice and MemoryInfo together in
one pass. Remove CreateIreeHwDevices method.

- Remove dead CPU device type check: replace device_type check (always
GPU) with backend-based validation — require target_arch for hip, cuda,
vulkan backends; make it optional for llvm-cpu.

- Fix CPU detection in data transfer: use MemoryDevice_GetDeviceType
instead of unreliable vendor_id == 0 check, since ORT assigns real
vendor IDs to CPU devices.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: Groverkss

src/iree_data_transfer.cc

@@ -40,24 +40,23 @@ bool ORT_API_CALL IreeDataTransfer::CanCopyImpl(
     const OrtMemoryDevice* dst_memory_device) noexcept {
   const auto* self = static_cast<const IreeDataTransfer*>(this_ptr);
 
-  // Get vendor IDs from both memory devices.
   uint32_t src_vendor_id =
       self->factory_.ep_api.MemoryDevice_GetVendorId(src_memory_device);
   uint32_t dst_vendor_id =
       self->factory_.ep_api.MemoryDevice_GetVendorId(dst_memory_device);
+  OrtMemoryInfoDeviceType src_type =
+      self->factory_.ep_api.MemoryDevice_GetDeviceType(src_memory_device);
+  OrtMemoryInfoDeviceType dst_type =
+      self->factory_.ep_api.MemoryDevice_GetDeviceType(dst_memory_device);
 
-  // We can copy if either source or destination is our IREE device,
-  // and the other is either our IREE device or CPU (vendor_id == 0).
-  // TODO: Is the vendor id check for CPU actually correct? Maybe we can check
-  // the hardware device for and determine if it's HOST/CPU.
   bool src_is_iree = (src_vendor_id == kEpVendorId);
   bool dst_is_iree = (dst_vendor_id == kEpVendorId);
-  bool src_is_cpu = (src_vendor_id == 0);
-  bool dst_is_cpu = (dst_vendor_id == 0);
+  bool src_is_cpu = (src_type == OrtMemoryInfoDeviceType_CPU);
+  bool dst_is_cpu = (dst_type == OrtMemoryInfoDeviceType_CPU);
 
   // Supported transfers:
   // - IREE <-> CPU (H2D, D2H)
-  // - IREE <-> IREE (D2D on same device)
+  // - IREE <-> IREE (D2D)
   return (src_is_iree && (dst_is_iree || dst_is_cpu)) ||
          (dst_is_iree && (src_is_iree || src_is_cpu));
 }

src/iree_ep_factory.cc

@@ -12,7 +12,6 @@
 
 #include "iree_ep_factory.h"
 
-#include <algorithm>
 #include <memory>
 #include <mutex>
 
@@ -24,6 +23,9 @@ namespace onnxruntime::iree {
 
 namespace {
 
+constexpr const char* kEpVendor = "IREE";
+constexpr const char* kEpVersion = "0.1.0";
+
 // Stub kernel — never instantiated or executed. Exists only to satisfy the
 // CustomOpBase<TOp, TKernel> template requirements.
 struct ExternDispatchKernel {
@@ -110,89 +112,6 @@ IreeEpFactory::IreeEpFactory(const char* ep_name, ApiPtrs apis,
     iree_status_ignore(status);
     return;
   }
-
-  CreateIreeHwDevices();
-}
-
-void IreeEpFactory::CreateIreeHwDevices() {
-  iree_allocator_t allocator = iree_allocator_system();
-  iree_hal_driver_registry_t* registry =
-      iree_runtime_instance_driver_registry(instance_.Get());
-
-  iree_host_size_t driver_count = 0;
-  IreeAllocatedPtr<iree_hal_driver_info_t> driver_infos(allocator);
-  iree_status_t status = iree_hal_driver_registry_enumerate(
-      registry, allocator, &driver_count, driver_infos.ForOutput());
-  if (!iree_status_is_ok(status)) {
-    ORT_CXX_LOG_NOEXCEPT(logger_, ORT_LOGGING_LEVEL_WARNING,
-                         "IREE EP: Failed to enumerate drivers");
-    iree_status_ignore(status);
-    return;
-  }
-
-  size_t device_index = 0;
-  for (iree_host_size_t i = 0; i < driver_count; ++i) {
-    iree_string_view_t driver_name = driver_infos.Get()[i].driver_name;
-    std::string driver_name_str(driver_name.data, driver_name.size);
-
-    HalDriverPtr driver;
-    status = iree_hal_driver_registry_try_create(registry, driver_name,
-                                                 allocator, driver.ForOutput());
-    if (!iree_status_is_ok(status)) {
-      iree_status_ignore(status);
-      continue;
-    }
-
-    iree_host_size_t device_count = 0;
-    IreeAllocatedPtr<iree_hal_device_info_t> device_infos(allocator);
-    status = iree_hal_driver_query_available_devices(
-        driver.Get(), allocator, &device_count, device_infos.ForOutput());
-    if (!iree_status_is_ok(status)) {
-      iree_status_ignore(status);
-      continue;
-    }
-
-    // Create OrtHardwareDevice for each IREE device available.
-    //
-    // Given an OrtHardwareDevice, ORT can find which hardware device it is by
-    // checking its vendor id (should be kEpVendorId) and the device id (the
-    // index it's stored in hw_devices_).
-    //
-    // IREE can identify the device using the driver and device path stored in
-    // the hardware device metadata.
-    for (iree_host_size_t j = 0; j < device_count; ++j) {
-      iree_string_view_t device_path = device_infos.Get()[j].path;
-      std::string device_path_str(device_path.data, device_path.size);
-
-      OrtKeyValuePairs* hw_metadata = nullptr;
-      ort_api.CreateKeyValuePairs(&hw_metadata);
-      ort_api.AddKeyValuePair(hw_metadata, "iree.driver",
-                              driver_name_str.c_str());
-      ort_api.AddKeyValuePair(hw_metadata, "iree.device_path",
-                              device_path_str.c_str());
-      // Store device_id for later retrieval in CreateEpImpl.
-      ort_api.AddKeyValuePair(hw_metadata, "iree.device_id",
-                              std::to_string(device_index).c_str());
-
-      // TODO: We are pretending here that all IREE devices are GPUs. This is
-      // because we follow a host <-> device model in IREE. How does this
-      // matter in practice? If we are using IREE APIs anyway, does it matter
-      // if we set this to CPU?
-      OrtHardwareDeviceType device_type = OrtHardwareDeviceType_GPU;
-
-      OrtHardwareDevice* hw_device = nullptr;
-      OrtStatus* ort_status = ep_api.CreateHardwareDevice(
-          device_type, kEpVendorId, static_cast<uint32_t>(device_index++),
-          kEpVendor, hw_metadata, &hw_device);
-      ort_api.ReleaseKeyValuePairs(hw_metadata);
-      if (ort_status != nullptr) {
-        ort_api.ReleaseStatus(ort_status);
-        continue;
-      }
-
-      hw_devices_.push_back(hw_device);
-    }
-  }
 }
 
 IreeEpFactory::~IreeEpFactory() {
@@ -309,44 +228,109 @@ OrtStatus* ORT_API_CALL IreeEpFactory::GetSupportedDevicesImpl(
   size_t& num_ep_devices = *p_num_ep_devices;
   num_ep_devices = 0;
 
-  // Reserve capacity for memory_info objects to prevent reallocation during
-  // the loop. This is critical because EpDevice_AddAllocatorInfo stores
-  // pointers and reallocation would invalidate them.
-  size_t num_devices_to_add =
-      std::min(factory->hw_devices_.size(), max_ep_devices);
-  factory->device_memory_infos_.reserve(num_devices_to_add);
+  // Lazily enumerate IREE devices on first call.
+  if (factory->hw_devices_.empty()) {
+    iree_allocator_t allocator = iree_allocator_system();
+    iree_hal_driver_registry_t* registry =
+        iree_runtime_instance_driver_registry(factory->instance_.Get());
 
-  // Create OrtEpDevice for each hardware device enumerated in constructor.
+    iree_host_size_t driver_count = 0;
+    IreeAllocatedPtr<iree_hal_driver_info_t> driver_infos(allocator);
+    iree_status_t status = iree_hal_driver_registry_enumerate(
+        registry, allocator, &driver_count, driver_infos.ForOutput());
+    if (!iree_status_is_ok(status)) {
+      ORT_CXX_LOG_NOEXCEPT(factory->logger_, ORT_LOGGING_LEVEL_WARNING,
+                           "IREE EP: Failed to enumerate drivers");
+      iree_status_ignore(status);
+      return nullptr;
+    }
+
+    size_t device_index = 0;
+    for (iree_host_size_t i = 0; i < driver_count; ++i) {
+      iree_string_view_t driver_name = driver_infos.Get()[i].driver_name;
+      std::string driver_name_str(driver_name.data, driver_name.size);
+
+      HalDriverPtr driver;
+      status = iree_hal_driver_registry_try_create(
+          registry, driver_name, allocator, driver.ForOutput());
+      if (!iree_status_is_ok(status)) {
+        iree_status_ignore(status);
+        continue;
+      }
+
+      iree_host_size_t device_count = 0;
+      IreeAllocatedPtr<iree_hal_device_info_t> device_infos(allocator);
+      status = iree_hal_driver_query_available_devices(
+          driver.Get(), allocator, &device_count, device_infos.ForOutput());
+      if (!iree_status_is_ok(status)) {
+        iree_status_ignore(status);
+        continue;
+      }
+
+      // Create OrtHardwareDevice and MemoryInfo for each IREE device.
+      //
+      // Given an OrtHardwareDevice, ORT identifies it by vendor_id
+      // (kEpVendorId) and device_id (index in hw_devices_). IREE identifies
+      // the device using the driver and device path in metadata.
+      for (iree_host_size_t j = 0; j < device_count; ++j) {
+        iree_string_view_t device_path = device_infos.Get()[j].path;
+        std::string device_path_str(device_path.data, device_path.size);
+
+        OrtKeyValuePairs* hw_metadata = nullptr;
+        factory->ort_api.CreateKeyValuePairs(&hw_metadata);
+        factory->ort_api.AddKeyValuePair(hw_metadata, "iree.driver",
+                                         driver_name_str.c_str());
+        factory->ort_api.AddKeyValuePair(hw_metadata, "iree.device_path",
+                                         device_path_str.c_str());
+        factory->ort_api.AddKeyValuePair(hw_metadata, "iree.device_id",
+                                         std::to_string(device_index).c_str());
+
+        // TODO: We are pretending here that all IREE devices are GPUs.
+        // This is because we follow a host <-> device model in IREE.
+        OrtHardwareDeviceType device_type = OrtHardwareDeviceType_GPU;
+
+        OrtHardwareDevice* hw_device = nullptr;
+        OrtStatus* ort_status = factory->ep_api.CreateHardwareDevice(
+            device_type, kEpVendorId, static_cast<uint32_t>(device_index),
+            kEpVendor, hw_metadata, &hw_device);
+        factory->ort_api.ReleaseKeyValuePairs(hw_metadata);
+        if (ort_status != nullptr) {
+          factory->ort_api.ReleaseStatus(ort_status);
+          continue;
+        }
+
+        factory->hw_devices_.push_back(hw_device);
+
+        // Create MemoryInfo for device-local memory. These must stay alive
+        // for the factory lifetime since EpDevice_AddAllocatorInfo stores
+        // the pointer without copying.
+        factory->device_memory_infos_.emplace_back(
+            "IREE",                               // name
+            OrtMemoryInfoDeviceType_GPU,          // device_type
+            kEpVendorId,                          // vendor_id
+            static_cast<uint32_t>(device_index),  // device_id
+            OrtDeviceMemoryType_DEFAULT,          // mem_type
+            0,                                    // alignment (default)
+            OrtDeviceAllocator);                  // allocator_type
+
+        ++device_index;
+      }
+    }
+  }
+
+  // Create OrtEpDevice for each enumerated hardware device.
+  // ORT takes ownership of OrtEpDevice objects.
   for (size_t i = 0;
        i < factory->hw_devices_.size() && num_ep_devices < max_ep_devices;
        ++i) {
     OrtEpDevice* ep_device = nullptr;
     ORT_RETURN_IF_ERROR(factory->ep_api.CreateEpDevice(
         factory, factory->hw_devices_[i], nullptr, nullptr, &ep_device));
 
-    // Register allocator info for device-local memory.
-    // This tells ORT that this device has an allocator available for
-    // device-local memory allocations.
-    //
-    // IMPORTANT: EpDevice_AddAllocatorInfo does NOT copy the OrtMemoryInfo,
-    // it just stores the pointer. We must keep the memory_info alive for
-    // the lifetime of the factory by storing it in device_memory_infos_.
-    factory->device_memory_infos_.emplace_back(
-        "IREE",                       // name
-        OrtMemoryInfoDeviceType_GPU,  // device_type
-        kEpVendorId,                  // vendor_id
-        static_cast<uint32_t>(i),     // device_id
-        OrtDeviceMemoryType_DEFAULT,  // mem_type
-        0,                            // alignment (default)
-        OrtDeviceAllocator);          // allocator_type
-
-    // Get raw pointer to pass to ORT.
-    const OrtMemoryInfo* mem_info_ptr = factory->device_memory_infos_.back();
-
+    const OrtMemoryInfo* mem_info_ptr = factory->device_memory_infos_[i];
     OrtStatus* add_alloc_status =
         factory->ep_api.EpDevice_AddAllocatorInfo(ep_device, mem_info_ptr);
     if (add_alloc_status != nullptr) {
-      factory->device_memory_infos_.pop_back();
       factory->ep_api.ReleaseEpDevice(ep_device);
       return add_alloc_status;
     }
@@ -439,13 +423,11 @@ OrtStatus* ORT_API_CALL IreeEpFactory::CreateEpImpl(
         .release();
   }
 
-  // Require target_arch for non-CPU devices.
-  OrtHardwareDeviceType device_type =
-      factory->ort_api.HardwareDevice_Type(&hardware_device);
-  if (device_type != OrtHardwareDeviceType_CPU && config.target_arch.empty()) {
+  // Require target_arch for GPU backends (llvm-cpu defaults to host).
+  if (config.backend != "llvm-cpu" && config.target_arch.empty()) {
     return Ort::Status(
-               "IREE EP: 'target_arch' option must be specified for non-CPU "
-               "devices",
+               "IREE EP: 'target_arch' option must be specified for "
+               "hip, cuda, and vulkan backends",
                ORT_INVALID_ARGUMENT)
         .release();
   }

src/iree_ep_factory.h

@@ -26,18 +26,8 @@ namespace onnxruntime::iree {
 class IreeAllocator;
 class IreeDataTransfer;
 
-// EP configuration constants
-inline constexpr const char* kEpVendor = "IREE";
+// EP vendor ID used across the codebase to identify IREE devices.
 inline constexpr uint32_t kEpVendorId = 0x1EEE;  // "IREE" in hex-ish
-inline constexpr const char* kEpVersion = "0.1.0";
-
-// Hardware vendor IDs for device matching.
-// These match OrtDevice::VendorIds from onnxruntime/core/framework/ortdevice.h.
-namespace VendorIds {
-inline constexpr uint32_t kAmd = 0x1002;     // AMD: ROCm, MIGraphX EPs
-inline constexpr uint32_t kNvidia = 0x10DE;  // NVIDIA: CUDA/TensorRT
-inline constexpr uint32_t kIntel = 0x8086;   // Intel: OpenVINO
-}  // namespace VendorIds
 
 // Helper struct to pass API pointers
 struct ApiPtrs {
@@ -121,9 +111,6 @@ class IreeEpFactory : public OrtEpFactory, public ApiPtrs {
   GetCustomOpDomainsImpl(OrtEpFactory* this_ptr, OrtCustomOpDomain** domains,
                          size_t num_domains) noexcept;
 
-  // Enumerate IREE devices and create OrtHardwareDevice instances.
-  void CreateIreeHwDevices();
-
   // Member variables
   Ort::Logger logger_;
   const std::string ep_name_;