feat: add Kunlun XRE runtime support

CMakeLists.txt

@@ -39,6 +39,7 @@ option(WITH_MARS "Enable Mars backend" OFF)
 option(WITH_CAMBRICON "Enable Cambricon backend" OFF)
 option(WITH_MOORE "Enable Moore backend" OFF)
 option(WITH_ASCEND "Enable Ascend backend" OFF)
+option(WITH_KUNLUN "Enable Kunlun backend" OFF)
 
 option(AUTO_DETECT_DEVICES "Automatically detect available devices" OFF)
 option(INFINI_RT_BUILD_TESTING "Build InfiniRT tests" OFF)
@@ -123,7 +124,7 @@ if(AUTO_DETECT_DEVICES)
 
     if(WITH_NVIDIA)
         set(_non_nvidia_gpu_detected FALSE)
-        foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_CAMBRICON WITH_MOORE WITH_ASCEND)
+        foreach(_gpu_backend WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_CAMBRICON WITH_MOORE WITH_ASCEND WITH_KUNLUN)
             if(${_gpu_backend})
                 set(_non_nvidia_gpu_detected TRUE)
             endif()
@@ -140,7 +141,7 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src)
 
 # Only one CUDA-like GPU backend can be enabled at a time.
 set(_gpu_backend_count 0)
-foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_MOORE WITH_CAMBRICON WITH_ASCEND)
+foreach(_gpu_backend WITH_NVIDIA WITH_ILUVATAR WITH_HYGON WITH_THEAD WITH_METAX WITH_MARS WITH_MOORE WITH_CAMBRICON WITH_ASCEND WITH_KUNLUN)
     if(${_gpu_backend})
         math(EXPR _gpu_backend_count "${_gpu_backend_count} + 1")
     endif()
@@ -348,8 +349,16 @@ if(WITH_ASCEND)
     message(STATUS "Using Ascend from `${ASCEND_HOME}`.")
 endif()
 
+if(WITH_KUNLUN)
+    set(XRE_ROOT "$ENV{XRE_ROOT}" CACHE PATH "Kunlun XRE SDK root")
+    find_path(KUNLUN_INCLUDE_DIR xpu/runtime.h
+        HINTS "${XRE_ROOT}" PATH_SUFFIXES include REQUIRED)
+    find_library(KUNLUN_RUNTIME_LIB xpurt
+        HINTS "${XRE_ROOT}" PATH_SUFFIXES lib lib64 so REQUIRED)
+endif()
+
 # If all other platforms are not enabled, CPU is enabled by default.
-if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD AND NOT WITH_METAX AND NOT WITH_MARS AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND)
+if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_THEAD AND NOT WITH_METAX AND NOT WITH_MARS AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND AND NOT WITH_KUNLUN)
     set(WITH_CPU ON)
     add_compile_definitions(WITH_CPU=1)
 endif()
@@ -386,6 +395,10 @@ if(WITH_ASCEND)
     list(APPEND INFINI_RT_PUBLIC_HEADER_DEVICES ascend)
 endif()
 
+if(WITH_KUNLUN)
+    list(APPEND INFINI_RT_PUBLIC_HEADER_DEVICES kunlun)
+endif()
+
 find_package(Python COMPONENTS Interpreter REQUIRED)
 execute_process(
     COMMAND ${Python_EXECUTABLE}
@@ -408,7 +421,7 @@ if(WITH_NVIDIA OR WITH_ILUVATAR OR WITH_HYGON OR WITH_THEAD)
 endif()
 
 set(INFINI_RT_PACKAGE_VENDOR_BACKEND "")
-foreach(_backend METAX MARS MOORE CAMBRICON ASCEND)
+foreach(_backend METAX MARS MOORE CAMBRICON ASCEND KUNLUN)
     if(WITH_${_backend})
         set(INFINI_RT_PACKAGE_VENDOR_BACKEND "${_backend}")
     endif()

cmake/InfiniRTConfig.cmake.in

@@ -15,6 +15,7 @@ set(InfiniRT_WITH_MARS @WITH_MARS@)
 set(InfiniRT_WITH_MOORE @WITH_MOORE@)
 set(InfiniRT_WITH_CAMBRICON @WITH_CAMBRICON@)
 set(InfiniRT_WITH_ASCEND @WITH_ASCEND@)
+set(InfiniRT_WITH_KUNLUN @WITH_KUNLUN@)
 
 if(@INFINI_RT_PACKAGE_NEEDS_CUDATOOLKIT@)
     find_dependency(CUDAToolkit)
@@ -131,6 +132,24 @@ elseif(_InfiniRT_VENDOR_BACKEND STREQUAL "CAMBRICON")
     _infinirt_define_backend_runtime(
         "${_InfiniRT_BACKEND_INCLUDE_DIR}"
         "${_InfiniRT_BACKEND_RUNTIME_LIBRARY}")
+elseif(_InfiniRT_VENDOR_BACKEND STREQUAL "KUNLUN")
+    find_dependency(Threads)
+    find_path(_InfiniRT_BACKEND_INCLUDE_DIR xpu/runtime.h
+        HINTS "${InfiniRT_KUNLUN_ROOT}" "$ENV{XRE_ROOT}"
+        PATH_SUFFIXES include)
+    find_library(_InfiniRT_BACKEND_RUNTIME_LIBRARY xpurt
+        HINTS "${InfiniRT_KUNLUN_ROOT}" "$ENV{XRE_ROOT}"
+        PATH_SUFFIXES lib lib64 so)
+    if(NOT _InfiniRT_BACKEND_INCLUDE_DIR OR
+       NOT _InfiniRT_BACKEND_RUNTIME_LIBRARY)
+        set(InfiniRT_FOUND FALSE)
+        set(InfiniRT_NOT_FOUND_MESSAGE
+            "Kunlun XRE SDK not found. Set InfiniRT_KUNLUN_ROOT or XRE_ROOT.")
+        return()
+    endif()
+    _infinirt_define_backend_runtime(
+        "${_InfiniRT_BACKEND_INCLUDE_DIR}"
+        "${_InfiniRT_BACKEND_RUNTIME_LIBRARY}")
 elseif(_InfiniRT_VENDOR_BACKEND STREQUAL "ASCEND")
     set(_InfiniRT_ROOT_HINTS "${InfiniRT_ASCEND_ROOT}")
     if(DEFINED ENV{ASCEND_HOME_PATH})

scripts/generate_public_headers.py

@@ -6,6 +6,11 @@
 _DETAIL_PREFIX = "infini/rt/detail"
 
 _DEVICE_HEADERS = {
+    "kunlun": (
+        ("kunlun", "data_type_.h", "native/kunlun/data_type_.h"),
+        ("kunlun", "device_.h", "native/kunlun/device_.h"),
+        ("kunlun", "runtime_.h", "native/kunlun/runtime_.h"),
+    ),
     "cpu": (
         ("cpu", "data_type_.h", "native/cpu/data_type_.h"),
         ("cpu", "device_.h", "native/cpu/device_.h"),
@@ -70,6 +75,7 @@
     "moore": "Device::Type::kMoore",
     "cambricon": "Device::Type::kCambricon",
     "ascend": "Device::Type::kAscend",
+    "kunlun": "Device::Type::kKunlun",
 }
 
 _DEFAULT_DEVICE_PRIORITY = (
@@ -82,6 +88,7 @@
     "moore",
     "cambricon",
     "ascend",
+    "kunlun",
     "cpu",
 )
 

src/CMakeLists.txt

@@ -105,6 +105,14 @@ if(WITH_CAMBRICON)
         "${CAMBRICON_RUNTIME_LIB}")
 endif()
 
+if(WITH_KUNLUN)
+    find_package(Threads REQUIRED)
+    target_link_libraries(infinirt PUBLIC Threads::Threads)
+    target_compile_definitions(infinirt PRIVATE WITH_KUNLUN=1)
+    _infinirt_use_backend_runtime(
+        "${KUNLUN_INCLUDE_DIR}" "${KUNLUN_RUNTIME_LIB}")
+endif()
+
 if(WITH_ASCEND)
     # ASCEND_HOME is set by the top-level CMakeLists.txt.
     target_compile_definitions(infinirt PUBLIC WITH_ASCEND=1)

src/device.h

@@ -23,6 +23,7 @@ class Device {
     kHygon = 7,
     kThead = 8,
     kMars = 9,
+    kKunlun = 10,
     kCount
   };
 
@@ -68,6 +69,7 @@ class Device {
           {Type::kHygon, "hygon"},
           {Type::kThead, "thead"},
           {Type::kMars, "mars"},
+          {Type::kKunlun, "kunlun"},
       }}};
 
   static constexpr ConstexprMap<std::string_view, Device::Type,
@@ -83,6 +85,7 @@ class Device {
           {"hygon", Type::kHygon},
           {"thead", Type::kThead},
           {"mars", Type::kMars},
+          {"kunlun", Type::kKunlun},
       }}};
 
   int index_{0};
@@ -98,7 +101,7 @@ using AllDeviceTypes =
     List<Device::Type::kCpu, Device::Type::kNvidia, Device::Type::kCambricon,
          Device::Type::kAscend, Device::Type::kMetax, Device::Type::kMoore,
          Device::Type::kIluvatar, Device::Type::kHygon, Device::Type::kThead,
-         Device::Type::kMars>;
+         Device::Type::kMars, Device::Type::kKunlun>;
 
 // Deferred computation of active devices. The `Filter` and `FilterList`
 // evaluation are nested inside a class template so that `DeviceEnabled`

src/native/kunlun/data_type_.h

@@ -0,0 +1,22 @@
+#ifndef INFINI_RT_KUNLUN_DATA_TYPE__H_
+#define INFINI_RT_KUNLUN_DATA_TYPE__H_
+
+#include "data_type.h"
+#include "native/kunlun/device_.h"
+
+namespace infini::rt {
+
+// Host storage types; device kernels use the corresponding XTDK intrinsics.
+template <>
+struct TypeMap<Device::Type::kKunlun, DataType::kFloat16> {
+  using type = Float16;
+};
+
+template <>
+struct TypeMap<Device::Type::kKunlun, DataType::kBFloat16> {
+  using type = BFloat16;
+};
+
+}  // namespace infini::rt
+
+#endif

src/native/kunlun/device_.h

@@ -0,0 +1,13 @@
+#ifndef INFINI_RT_KUNLUN_DEVICE__H_
+#define INFINI_RT_KUNLUN_DEVICE__H_
+
+#include "device.h"
+
+namespace infini::rt {
+
+template <>
+struct DeviceEnabled<Device::Type::kKunlun> : std::true_type {};
+
+}  // namespace infini::rt
+
+#endif

src/native/kunlun/runtime_.h

@@ -0,0 +1,169 @@
+#ifndef INFINI_RT_KUNLUN_RUNTIME__H_
+#define INFINI_RT_KUNLUN_RUNTIME__H_
+
+#include <xpu/runtime.h>
+
+#include <algorithm>
+#include <array>
+#include <cstring>
+#include <mutex>
+#include <utility>
+#include <vector>
+
+#include "native/kunlun/device_.h"
+#include "runtime.h"
+
+namespace infini::rt::runtime {
+
+template <>
+struct Runtime<Device::Type::kKunlun>
+    : DeviceRuntime<Runtime<Device::Type::kKunlun>> {
+  using Error = int;
+  using Stream = XPUStream;
+  using Event = XPUEvent;
+  using Graph = void*;
+  using GraphExec = void*;
+  using StreamCaptureMode = int;
+
+  static constexpr Device::Type kDeviceType = Device::Type::kKunlun;
+  static constexpr Error kSuccess = XPU_SUCCESS;
+  static constexpr auto SetDevice = xpu_set_device;
+  static constexpr auto GetDevice = xpu_current_device;
+  static constexpr auto GetDeviceCount = xpu_device_count;
+
+  static Error DeviceSynchronize() {
+    int device;
+    auto status = GetDevice(&device);
+    if (status != kSuccess) return status;
+    std::lock_guard<std::mutex> lock(stream_mutex_);
+    status = xpu_wait(nullptr);
+    if (status != kSuccess) return status;
+    // XRE only synchronizes individual streams. Include all streams created
+    // by this runtime on the current device, not just the default stream.
+    for (const auto& entry : streams_) {
+      if (entry.first == device) {
+        status = xpu_wait(entry.second);
+        if (status != kSuccess) return status;
+      }
+    }
+    return kSuccess;
+  }
+
+  static Error Malloc(void** ptr, std::size_t size) {
+    return xpu_malloc(ptr, size, XPU_MEM_MAIN);
+  }
+  static Error MallocHost(void** ptr, std::size_t size) {
+    return xpu_host_alloc(ptr, size, 0);
+  }
+  static Error MallocAsync(void**, std::size_t, Stream) {
+    return XPUERR_NOSUPPORT;
+  }
+  static constexpr auto Free = xpu_free;
+  static constexpr auto FreeHost = xpu_host_free;
+  static Error FreeAsync(void*, Stream) { return XPUERR_NOSUPPORT; }
+  static Error MemGetInfo(std::size_t*, std::size_t*) {
+    return XPUERR_NOSUPPORT;
+  }
+
+  // XRE has no host-to-host enumerator.
+  static constexpr int kMemcpyHostToHost = 3;
+  static constexpr int kMemcpyHostToDevice = XPU_HOST_TO_DEVICE;
+  static constexpr int kMemcpyDeviceToHost = XPU_DEVICE_TO_HOST;
+  static constexpr int kMemcpyDeviceToDevice = XPU_DEVICE_TO_DEVICE;
+
+  static Error Memcpy(void* dst, const void* src, std::size_t size, int kind) {
+    if (kind == kMemcpyHostToHost) {
+      std::memcpy(dst, src, size);
+      return kSuccess;
+    }
+    return xpu_memcpy(dst, src, size, static_cast<XPUMemcpyKind>(kind));
+  }
+  static Error MemcpyAsync(void* dst, const void* src, std::size_t size,
+                           int kind, Stream stream) {
+    if (kind == kMemcpyHostToHost) {
+      const auto status = xpu_wait(stream);
+      if (status != kSuccess) return status;
+      return Memcpy(dst, src, size, kind);
+    }
+    return xpu_memcpy_async(dst, src, size, static_cast<XPUMemcpyKind>(kind),
+                            stream);
+  }
+
+  static Error Memset(void* ptr, int value, std::size_t size) {
+    // XRE has no memset primitive. Bound the synchronous staging memory.
+    std::array<unsigned char, 4096> staging;
+    staging.fill(static_cast<unsigned char>(value));
+    auto* dst = static_cast<unsigned char*>(ptr);
+    for (std::size_t done = 0; done < size;) {
+      const auto bytes = std::min(staging.size(), size - done);
+      const auto status =
+          xpu_memcpy(dst + done, staging.data(), bytes, XPU_HOST_TO_DEVICE);
+      if (status != kSuccess) return status;
+      done += bytes;
+    }
+    return kSuccess;
+  }
+  static Error MemsetAsync(void*, int, std::size_t, Stream) {
+    return XPUERR_NOSUPPORT;
+  }
+
+  static Error StreamCreate(Stream* stream) {
+    int device;
+    auto status = GetDevice(&device);
+    if (status != kSuccess) return status;
+    std::lock_guard<std::mutex> lock(stream_mutex_);
+    status = xpu_stream_create(stream);
+    if (status == kSuccess) streams_.emplace_back(device, *stream);
+    return status;
+  }
+  static Error StreamDestroy(Stream stream) {
+    std::lock_guard<std::mutex> lock(stream_mutex_);
+    const auto status = xpu_stream_destroy(stream);
+    if (status == kSuccess) {
+      streams_.erase(std::remove_if(streams_.begin(), streams_.end(),
+                                    [stream](const auto& entry) {
+                                      return entry.second == stream;
+                                    }),
+                     streams_.end());
+    }
+    return status;
+  }
+  static constexpr auto StreamSynchronize = xpu_wait;
+  static Error StreamWaitEvent(Stream stream, Event event, unsigned int flags) {
+    return flags == 0 ? xpu_stream_wait_event(stream, event) : XPUERR_NOSUPPORT;
+  }
+
+  static constexpr auto EventCreate = xpu_event_create;
+  static Error EventCreateWithFlags(Event* event, unsigned int flags) {
+    return flags == 0 ? EventCreate(event) : XPUERR_NOSUPPORT;
+  }
+  static constexpr auto EventRecord = xpu_event_record;
+  static constexpr auto EventQuery = xpu_event_query;
+  static constexpr auto EventSynchronize = xpu_event_wait;
+  static constexpr auto EventDestroy = xpu_event_destroy;
+  static Error EventElapsedTime(float*, Event, Event) {
+    return XPUERR_NOSUPPORT;
+  }
+
+  static constexpr StreamCaptureMode kStreamCaptureModeGlobal = 0;
+  static constexpr StreamCaptureMode kStreamCaptureModeThreadLocal = 1;
+  static constexpr StreamCaptureMode kStreamCaptureModeRelaxed = 2;
+  static Error StreamBeginCapture(Stream, StreamCaptureMode) {
+    return XPUERR_NOSUPPORT;
+  }
+  static Error StreamEndCapture(Stream, Graph*) { return XPUERR_NOSUPPORT; }
+  static Error GraphDestroy(Graph) { return XPUERR_NOSUPPORT; }
+  static Error GraphInstantiate(GraphExec*, Graph) { return XPUERR_NOSUPPORT; }
+  static Error GraphExecDestroy(GraphExec) { return XPUERR_NOSUPPORT; }
+  static Error GraphLaunch(GraphExec, Stream) { return XPUERR_NOSUPPORT; }
+
+ private:
+  inline static std::mutex stream_mutex_;
+  inline static std::vector<std::pair<int, Stream>> streams_;
+};
+
+static_assert(Runtime<Device::Type::kKunlun>::Validate());
+
+}  // namespace infini::rt::runtime
+
+#endif