use xrtDeviceHandle

nod-ai · Oct 1, 2024 · c1bdeb0 · c1bdeb0
1 parent 9e7d045
commit c1bdeb0
Show file tree

Hide file tree

Showing 14 changed files with 82 additions and 82 deletions.
diff --git a/.github/workflows/ci-linux.yml b/.github/workflows/ci-linux.yml
@@ -150,12 +150,11 @@ jobs:
           source .venv/bin/activate
           pip install -r tests/requirements.txt
 
-
       - name : E2E comparison of AIE to llvm-cpu
         run: |
           source .venv/bin/activate
           source /opt/xilinx/xrt/setup.sh
-          for i in {1..100}; do
+          for i in {1..50}; do
             echo "run $i"
             python build_tools/ci/cpu_comparison/run.py \
               test_aie_vs_cpu \
@@ -196,7 +195,7 @@ jobs:
           sudo prlimit -lunlimited --pid $$
           source .venv/bin/activate
           source /opt/xilinx/xrt/setup.sh
-          for i in {1..100}; do
+          for i in {1..50}; do
             echo "run $i"
             bash build_tools/ci/run_matmul_test.sh \
               test_matmuls \
@@ -210,7 +209,7 @@ jobs:
         run: |
           source .venv/bin/activate
           source /opt/xilinx/xrt/setup.sh
-          for i in {1..100}; do
+          for i in {1..50}; do
             echo "run $i"
             pytest -v tests \
               --capture=tee-sys \

diff --git a/.github/workflows/ci-windows.yml b/.github/workflows/ci-windows.yml
@@ -169,7 +169,7 @@ jobs:
         shell: bash
         run: |
           source .venv/Scripts/activate
-          for i in {1..100}; do
+          for i in {1..50}; do
             echo "run $i"
             bash build_tools/ci/run_matmul_test.sh \
               /c/test_matmuls \
@@ -181,7 +181,7 @@ jobs:
         shell: bash
         run: |
           source .venv/Scripts/activate
-          for i in {1..100}; do
+          for i in {1..50}; do
             echo "run $i"
             python build_tools/ci/cpu_comparison/run.py \
               /c/test_aie_vs_cpu \
@@ -194,7 +194,7 @@ jobs:
           ls $env:XILINX_XRT
           .\.venv\Scripts\Activate.ps1
           mkdir temp
-          for ($i = 1; $i -le 100; $i++) {
+          for ($i = 1; $i -le 50; $i++) {
             echo "run $i"
             pytest tests -sv `
               --basetemp=$PWD\temp `

diff --git a/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc b/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.cc
@@ -28,7 +28,7 @@ typedef struct iree_hal_xrt_allocator_t {
   // The device that this allocator is attached to.
   iree_hal_device_t* base_device;
 
-  xrt::device device;
+  xrtDeviceHandle device_hdl;
 
   iree_allocator_t host_allocator;
 
@@ -46,7 +46,7 @@ static iree_hal_xrt_allocator_t* iree_hal_xrt_allocator_cast(
 }
 
 iree_status_t iree_hal_xrt_allocator_create(
-    iree_hal_device_t* base_device, xrt::device device,
+    iree_hal_device_t* base_device, xrtDeviceHandle device_hdl,
     iree_allocator_t host_allocator, iree_hal_allocator_t** out_allocator) {
   IREE_ASSERT_ARGUMENT(base_device);
   IREE_ASSERT_ARGUMENT(out_allocator);
@@ -61,7 +61,7 @@ iree_status_t iree_hal_xrt_allocator_create(
                                &allocator->resource);
   allocator->base_device = base_device;
   iree_hal_device_retain(base_device);
-  allocator->device = device;
+  allocator->device_hdl = device_hdl;
   allocator->host_allocator = host_allocator;
 
   *out_allocator = (iree_hal_allocator_t*)allocator;
@@ -171,7 +171,8 @@ static iree_status_t iree_hal_xrt_allocator_allocate_buffer(
   std::unique_ptr<xrt::bo> xrt_buffer;
 
   try {
-    xrt_buffer = std::make_unique<xrt::bo>(allocator->device, allocation_size,
+    xrt::device device(xrtDeviceToXclDevice(allocator->device_hdl));
+    xrt_buffer = std::make_unique<xrt::bo>(device, allocation_size,
                                            XRT_BO_FLAGS_HOST_ONLY, group_id);
   } catch (...) {
     IREE_TRACE_ZONE_END(z0);

diff --git a/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.h b/runtime/src/iree-amd-aie/driver/xrt/direct_allocator.h
@@ -17,7 +17,7 @@ extern "C" {
 
 // Creates an XRT memory allocator.
 iree_status_t iree_hal_xrt_allocator_create(
-    iree_hal_device_t* base_device, xrt::device device,
+    iree_hal_device_t* base_device, xrtDeviceHandle device_hdl,
     iree_allocator_t host_allocator, iree_hal_allocator_t** out_allocator);
 
 #ifdef __cplusplus

diff --git a/runtime/src/iree-amd-aie/driver/xrt/direct_command_buffer.cc b/runtime/src/iree-amd-aie/driver/xrt/direct_command_buffer.cc
@@ -292,7 +292,6 @@ static iree_status_t iree_hal_xrt_direct_command_buffer_push_descriptor_set(
     IREE_RETURN_AND_END_ZONE_IF_ERROR(
         z0, iree_hal_resource_set_insert(command_buffer->resource_set, 1,
                                          &binding->buffer));
-    std::unique_ptr<xrt::bo> sub_buffer;
     current_bindings[i] = iree_hal_xrt_buffer_handle(
         iree_hal_buffer_allocated_buffer(binding->buffer));
     current_offsets[i] =

diff --git a/runtime/src/iree-amd-aie/driver/xrt/native_executable.cc b/runtime/src/iree-amd-aie/driver/xrt/native_executable.cc
@@ -99,7 +99,8 @@ static iree_status_t iree_amd_aie_hal_xrt_native_executable_flatbuffer_verify(
 }
 
 iree_status_t iree_hal_xrt_native_executable_create(
-    xrt::device device, const iree_hal_executable_params_t* executable_params,
+    xrtDeviceHandle device_hdl,
+    const iree_hal_executable_params_t* executable_params,
     iree_allocator_t host_allocator, iree_hal_executable_t** out_executable) {
   IREE_ASSERT_ARGUMENT(executable_params);
   IREE_ASSERT_ARGUMENT(out_executable);
@@ -180,24 +181,27 @@ iree_status_t iree_hal_xrt_native_executable_create(
     // XRT API needs this vector and cant actually read a void*.
     std::vector<char> xclbinVector(
         xclbin_fb, xclbin_fb + flatbuffers_string_len(xclbin_fb));
+    xrt::xclbin xclbin;
     try {
-      params->xclbin = xrt::xclbin(xclbinVector);
+      xclbin = xrt::xclbin(xclbinVector);
     } catch (std::exception& e) {
       return iree_make_status(IREE_STATUS_INTERNAL, "XCLBIN load error: %s",
                               e.what());
     }
 
+    xrt::device device(xrtDeviceToXclDevice(device_hdl));
+    IREE_ASSERT(device, "failed to find device");
+
     try {
-      device.register_xclbin(params->xclbin);
+      device.register_xclbin(xclbin);
     } catch (std::exception& e) {
       return iree_make_status(IREE_STATUS_INTERNAL, "XCLBIN register error: %s",
                               e.what());
     }
 
     try {
-      params->context =
-          xrt::hw_context(device, params->xclbin.get_uuid(),
-                          xrt::hw_context::access_mode::exclusive);
+      params->context = xrt::hw_context(
+          device, xclbin.get_uuid(), xrt::hw_context::access_mode::exclusive);
     } catch (std::exception& e) {
       return iree_make_status(IREE_STATUS_INTERNAL,
                               "xrt::hw_context context: %s", e.what());

diff --git a/runtime/src/iree-amd-aie/driver/xrt/native_executable.h b/runtime/src/iree-amd-aie/driver/xrt/native_executable.h
@@ -7,7 +7,7 @@
 #ifndef IREE_AMD_AIE_DRIVER_XRT_NATIVE_EXECUTABLE_H_
 #define IREE_AMD_AIE_DRIVER_XRT_NATIVE_EXECUTABLE_H_
 
-#include <stdint.h>
+#include <cstdint>
 
 #include "iree/base/api.h"
 #include "iree/base/tracing.h"
@@ -22,7 +22,6 @@ extern "C" {
 // Object and launch parameters for a compute kernel.
 typedef struct iree_hal_xrt_kernel_params_t {
   xrt::hw_context context;
-  xrt::xclbin xclbin;
   // The kernel code object.
   xrt::kernel kernel;
   // Instruction buffer argument to the kernel.
@@ -37,7 +36,8 @@ typedef struct iree_hal_xrt_kernel_params_t {
 // |out_executable| must be released by the caller (see
 // iree_hal_executable_release).
 iree_status_t iree_hal_xrt_native_executable_create(
-    xrt::device device, const iree_hal_executable_params_t* executable_params,
+    xrtDeviceHandle device_hdl,
+    const iree_hal_executable_params_t* executable_params,
     iree_allocator_t host_allocator, iree_hal_executable_t** out_executable);
 
 // Returns the kernel launch parameters for the given |entry_point|.

diff --git a/runtime/src/iree-amd-aie/driver/xrt/nop_executable_cache.cc b/runtime/src/iree-amd-aie/driver/xrt/nop_executable_cache.cc
@@ -17,7 +17,7 @@ typedef struct iree_hal_xrt_nop_executable_cache_t {
   // at offset 0.
   iree_hal_resource_t resource;
 
-  xrt::device device;
+  xrtDeviceHandle device_hdl;
 
   iree_allocator_t host_allocator;
 } iree_hal_xrt_nop_executable_cache_t;
@@ -35,7 +35,7 @@ iree_hal_xrt_nop_executable_cache_cast(
 }
 
 iree_status_t iree_hal_xrt_nop_executable_cache_create(
-    xrt::device device, iree_string_view_t identifier,
+    xrtDeviceHandle device_hdl, iree_string_view_t identifier,
     iree_allocator_t host_allocator,
     iree_hal_executable_cache_t** out_executable_cache) {
   IREE_ASSERT_ARGUMENT(out_executable_cache);
@@ -49,7 +49,7 @@ iree_status_t iree_hal_xrt_nop_executable_cache_create(
   iree_hal_resource_initialize(&iree_hal_xrt_nop_executable_cache_vtable,
                                &executable_cache->resource);
   executable_cache->host_allocator = host_allocator;
-  executable_cache->device = device;
+  executable_cache->device_hdl = device_hdl;
 
   *out_executable_cache = (iree_hal_executable_cache_t*)executable_cache;
   IREE_TRACE_ZONE_END(z0);
@@ -82,7 +82,7 @@ static iree_status_t iree_hal_xrt_nop_executable_cache_prepare_executable(
   iree_hal_xrt_nop_executable_cache_t* executable_cache =
       iree_hal_xrt_nop_executable_cache_cast(base_executable_cache);
   return iree_hal_xrt_native_executable_create(
-      executable_cache->device, executable_params,
+      executable_cache->device_hdl, executable_params,
       executable_cache->host_allocator, out_executable);
 }
 

diff --git a/runtime/src/iree-amd-aie/driver/xrt/nop_executable_cache.h b/runtime/src/iree-amd-aie/driver/xrt/nop_executable_cache.h
@@ -22,7 +22,7 @@ extern "C" {
 // |out_executable_cache| must be released by the caller (see
 // iree_hal_executable_cache_release).
 iree_status_t iree_hal_xrt_nop_executable_cache_create(
-    xrt::device device, iree_string_view_t identifier,
+    xrtDeviceHandle device_hdl, iree_string_view_t identifier,
     iree_allocator_t host_allocator,
     iree_hal_executable_cache_t** out_executable_cache);
 

diff --git a/runtime/src/iree-amd-aie/driver/xrt/xrt_device.cc b/runtime/src/iree-amd-aie/driver/xrt/xrt_device.cc
@@ -6,6 +6,7 @@
 
 #include "iree-amd-aie/driver/xrt/xrt_device.h"
 
+#include "experimental/xrt_system.h"
 #include "iree-amd-aie/driver/xrt/direct_allocator.h"
 #include "iree-amd-aie/driver/xrt/direct_command_buffer.h"
 #include "iree-amd-aie/driver/xrt/nop_executable_cache.h"
@@ -32,7 +33,7 @@ typedef struct iree_hal_xrt_device_t {
   iree_allocator_t host_allocator;
   iree_hal_allocator_t* device_allocator;
 
-  xrt::device device;
+  xrtDeviceHandle device_hdl;
 } iree_hal_xrt_device_t;
 
 namespace {
@@ -52,17 +53,30 @@ void iree_hal_xrt_device_params_initialize(
 }
 
 static iree_status_t iree_hal_xrt_device_create_internal(
-    iree_string_view_t identifier, xrt::device xrt_device,
-    const iree_hal_xrt_device_params_t* params, iree_allocator_t host_allocator,
-    iree_hal_device_t** out_device) {
+    iree_string_view_t identifier, const iree_hal_xrt_device_params_t* params,
+    iree_allocator_t host_allocator, iree_hal_device_t** out_device) {
   iree_hal_xrt_device_t* device = nullptr;
 
   iree_host_size_t total_size = iree_sizeof_struct(*device) + identifier.size;
   IREE_RETURN_IF_ERROR(
       iree_allocator_malloc(host_allocator, total_size, (void**)&device));
 
+  try {
+    if (IREE_UNLIKELY(xrt::system::enumerate_devices() == 0)) {
+      return iree_make_status(IREE_STATUS_FAILED_PRECONDITION,
+                              "No XRT devices found");
+    }
+  } catch (std::exception& e) {
+    return iree_make_status(IREE_STATUS_INTERNAL,
+                            "xrt::system::enumerate_devices failed: %s",
+                            e.what());
+  }
+
+  xrtDeviceHandle device_hdl = xrtDeviceOpen(0);
+  IREE_ASSERT(device_hdl, "failed to open xrt device");
+
   iree_status_t status =
-      iree_hal_xrt_allocator_create((iree_hal_device_t*)device, xrt_device,
+      iree_hal_xrt_allocator_create((iree_hal_device_t*)device, device_hdl,
                                     host_allocator, &device->device_allocator);
   if (iree_status_is_ok(status)) {
     iree_hal_resource_initialize(&iree_hal_xrt_device_vtable,
@@ -74,7 +88,7 @@ static iree_status_t iree_hal_xrt_device_create_internal(
                                      &device->block_pool);
 
     device->host_allocator = host_allocator;
-    device->device = xrt_device;
+    device->device_hdl = device_hdl;
     device->params = *params;
     *out_device = (iree_hal_device_t*)device;
   } else {
@@ -85,13 +99,12 @@ static iree_status_t iree_hal_xrt_device_create_internal(
 
 iree_status_t iree_hal_xrt_device_create(
     iree_string_view_t identifier, const iree_hal_xrt_device_params_t* params,
-    xrt::device device, iree_allocator_t host_allocator,
-    iree_hal_device_t** out_device) {
+    iree_allocator_t host_allocator, iree_hal_device_t** out_device) {
   IREE_ASSERT_ARGUMENT(out_device);
   IREE_TRACE_ZONE_BEGIN(z0);
 
   iree_status_t status = iree_hal_xrt_device_create_internal(
-      identifier, device, params, host_allocator, out_device);
+      identifier, params, host_allocator, out_device);
 
   IREE_TRACE_ZONE_END(z0);
   return status;
@@ -104,7 +117,9 @@ static void iree_hal_xrt_device_destroy(iree_hal_device_t* base_device) {
 
   iree_hal_allocator_release(device->device_allocator);
   iree_arena_block_pool_deinitialize(&device->block_pool);
+  xrtDeviceHandle device_hdl = device->device_hdl;
   iree_allocator_free(host_allocator, device);
+  (void)xrtDeviceClose(device_hdl);
 
   IREE_TRACE_ZONE_END(z0);
 }
@@ -201,7 +216,8 @@ static iree_status_t iree_hal_xrt_device_create_executable_cache(
     iree_loop_t loop, iree_hal_executable_cache_t** out_executable_cache) {
   iree_hal_xrt_device_t* device = iree_hal_xrt_device_cast(base_device);
   return iree_hal_xrt_nop_executable_cache_create(
-      device->device, identifier, device->host_allocator, out_executable_cache);
+      device->device_hdl, identifier, device->host_allocator,
+      out_executable_cache);
 }
 
 static iree_status_t iree_hal_xrt_device_import_file(

diff --git a/runtime/src/iree-amd-aie/driver/xrt/xrt_device.h b/runtime/src/iree-amd-aie/driver/xrt/xrt_device.h
@@ -21,8 +21,7 @@ extern "C" {
 // |out_device| must be released by the caller (see iree_hal_device_release).
 iree_status_t iree_hal_xrt_device_create(
     iree_string_view_t identifier, const iree_hal_xrt_device_params_t* params,
-    xrt::device device, iree_allocator_t host_allocator,
-    iree_hal_device_t** out_device);
+    iree_allocator_t host_allocator, iree_hal_device_t** out_device);
 
 #ifdef __cplusplus
 }  // extern "C"