issue/87: restruct kunlun handle

src/infiniop/devices/handle.cc

@@ -48,9 +48,7 @@ __C infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr) {
         CREATE(INFINI_DEVICE_ASCEND, ascend);
 #endif
 #ifdef ENABLE_KUNLUN_API
-    case INFINI_DEVICE_KUNLUN: {
-        return createKunlunHandle((infiniopKunlunHandle_t *)handle_ptr);
-    }
+        CREATE(INFINI_DEVICE_KUNLUN, kunlun);
 #endif
 #ifdef ENABLE_METAX_API
         CREATE(INFINI_DEVICE_METAX, maca);
@@ -84,9 +82,7 @@ __C infiniStatus_t infiniopDestroyHandle(infiniopHandle_t handle) {
         DELETE(INFINI_DEVICE_ASCEND, ascend);
 #endif
 #ifdef ENABLE_KUNLUN_API
-    case INFINI_DEVICE_KUNLUN: {
-        return destroyKunlunHandle((infiniopKunlunHandle_t)handle);
-    }
+        DELETE(INFINI_DEVICE_KUNLUN, kunlun);
 #endif
 #ifdef ENABLE_METAX_API
         DELETE(INFINI_DEVICE_METAX, maca);

src/infiniop/devices/kunlun/common_kunlun.h

-#ifndef __INFINIOP_COMMON_KUNLUN_H__
-#define __INFINIOP_COMMON_KUNLUN_H__
-
-#include "../../../utils.h"
-#include "../pool.h"
-#include "infinicore.h"
-#include "kunlun_handle.h"
-#include "xpu/runtime.h"
-#include "xpu/runtime_ex.h"
-#include "xpu/xdnn.h"
-#include <memory>
-
-namespace xdnn = baidu::xpu::api;
-typedef xdnn::Context *xdnnHandle_t;
-typedef XPUStream KunlunStream_t;
-
-#define CHECK_KUNLUN(call) CHECK_INTERNAL(call, XPU_SUCCESS)
-
-struct InfiniopKunlunHandle {
-    infiniDevice_t device;
-    int device_id;
-    std::shared_ptr<Pool<xdnnHandle_t>> xdnn_handle_pool;
-};
-
-template <typename T>
-void use_xdnn(std::shared_ptr<Pool<xdnnHandle_t>> &pool, KunlunStream_t stream, const T &f) {
-    auto handle = pool->pop();
-    if (!handle) {
-        *handle = xdnn::create_context();
-    }
-    (*handle)->set_stream(stream);
-    f(*handle);
-    pool->push(std::move(*handle));
-}
-
-#endif //__INFINIOP_COMMON_KUNLUN_H__

src/infiniop/devices/kunlun/kunlun_handle.cc

-#include "common_kunlun.h"
+#include "kunlun_handle.h"
 
-infiniStatus_t createKunlunHandle(infiniopKunlunHandle_t *handle_ptr) {
-    int device_id;
-    CHECK_KUNLUN(xpu_current_device(&device_id));
-    auto pool = std::make_shared<Pool<xdnnHandle_t>>();
-    xdnnHandle_t handle = xdnn::create_context();
-    pool->push(std::move(handle));
+namespace device::kunlun {
 
-    *handle_ptr = new InfiniopKunlunHandle{
-        INFINI_DEVICE_KUNLUN,
-        device_id,
-        std::move(pool),
-    };
+Handle::Handle(infiniDevice_t device, int device_id)
+    : InfiniopHandle{device, device_id},
+      _internal(std::make_shared<Handle::Internal>()) {}
 
-    return INFINI_STATUS_SUCCESS;
+auto Handle::internal() const -> const std::shared_ptr<Internal> & {
+    return _internal;
 }
 
-infiniStatus_t destroyKunlunHandle(infiniopKunlunHandle_t handle_ptr) {
-    handle_ptr->xdnn_handle_pool = nullptr;
-    delete handle_ptr;
+template <typename T>
+using Fn = std::function<void(T)>;
+
+void Handle::Internal::use_xdnn(kunlunStream_t stream, const Fn<xdnnHandle_t> &f) const {
+    auto handle = dnn_handles.pop();
+    if (!handle) {
+        *handle = xdnn::create_context();
+    }
+    (*handle)->set_stream(stream);
+    f(*handle);
+    dnn_handles.push(std::move(*handle));
+}
 
+infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) {
+    *handle_ptr = new Handle(INFINI_DEVICE_KUNLUN, device_id);
     return INFINI_STATUS_SUCCESS;
 }
+
+} // namespace device::kunlun

src/infiniop/devices/kunlun/kunlun_handle.h

@@ -2,12 +2,41 @@
 #define __INFINIOP_KUNLUN_HANDLE_H__
 
 #include "../../handle.h"
+#include "../pool.h"
+#include <functional>
+#include <memory>
+#include <xpu/runtime.h>
+#include <xpu/runtime_ex.h>
+#include <xpu/xdnn.h>
 
-struct InfiniopKunlunHandle;
+namespace xdnn = baidu::xpu::api;
 
-typedef struct InfiniopKunlunHandle *infiniopKunlunHandle_t;
+typedef XPUStream kunlunStream_t;
+typedef XPUEvent kunlunEvent_t;
+typedef xdnn::Context *xdnnHandle_t;
 
-infiniStatus_t createKunlunHandle(infiniopKunlunHandle_t *handle_ptr);
-infiniStatus_t destroyKunlunHandle(infiniopKunlunHandle_t handle);
+namespace device::kunlun {
+
+struct Handle : public InfiniopHandle {
+    class Internal;
+    auto internal() const -> const std::shared_ptr<Internal> &;
+
+    Handle(infiniDevice_t device, int device_id);
+
+private:
+    std::shared_ptr<Internal> _internal;
+
+public:
+    static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id);
+};
+
+class Handle::Internal {
+    Pool<xdnnHandle_t> dnn_handles;
+
+public:
+    void use_xdnn(kunlunStream_t stream, const std::function<void(xdnnHandle_t)> &f) const;
+};
+
+} // namespace device::kunlun
 
 #endif // __INFINIOP_KUNLUN_HANDLE_H__

src/infiniop/ops/matmul/kunlun/matmul_kunlun.cc

 #include "matmul_kunlun.h"
-#include "../../../devices/kunlun/common_kunlun.h"
-#include "../../utils.h"
+#include "../../../../utils.h"
+#include "../../../devices/kunlun/kunlun_handle.h"
 
 namespace op::matmul::kunlun {
 
+typedef device::kunlun::Handle::Internal HandleInternal;
+
 struct Descriptor::Opaque {
-    std::shared_ptr<Pool<xdnnHandle_t>> xdnn_handle_pool;
+    std::shared_ptr<HandleInternal> internal;
 };
 
 Descriptor::~Descriptor() {
     delete _opaque;
 }
 
-infiniStatus_t Descriptor::create(infiniopHandle_t handle_,
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
     Descriptor **desc_ptr,
     infiniopTensorDescriptor_t c_desc,
     infiniopTensorDescriptor_t a_desc,
     infiniopTensorDescriptor_t b_desc) {
-    auto handle = reinterpret_cast<infiniopKunlunHandle_t>(handle_);
-    auto dtype = c_desc->dtype;
+    auto handle = reinterpret_cast<device::kunlun::Handle *>(handle_);
+    auto dtype = c_desc->dtype();
 
     if (dtype != INFINI_DTYPE_F16 && dtype != INFINI_DTYPE_F32) {
         return INFINI_STATUS_BAD_TENSOR_DTYPE;
@@ -32,22 +35,22 @@ infiniStatus_t Descriptor::create(infiniopHandle_t handle_,
 
     *desc_ptr = new Descriptor(
         dtype, info, 0,
-        new Opaque{handle->xdnn_handle_pool},
+        new Opaque{handle->internal()},
         handle->device, handle->device_id);
     return INFINI_STATUS_SUCCESS;
 }
 
 template <class Tdata>
 void calculate(
-    const MatmulInfo &info,
-    std::shared_ptr<Pool<xdnnHandle_t>> &xdnn_handle_pool,
+    MatmulInfo info,
+    std::shared_ptr<HandleInternal> internal,
     infiniDtype_t dtype,
     void *c,
     float beta,
     const void *a,
     const void *b,
     float alpha,
-    KunlunStream_t stream) {
+    kunlunStream_t stream) {
 
     if (info.is_transed) {
         std::swap(a, b);
@@ -58,8 +61,8 @@ void calculate(
 
     auto unit = infiniSizeOf(dtype);
 
-    use_xdnn(xdnn_handle_pool,
-             (KunlunStream_t)stream,
+    internal->use_xdnn(
+        (kunlunStream_t)stream,
         [&](xdnnHandle_t handle) {
             for (size_t i = 0; i < info.batch; i++) {
                 xdnn::fc_fusion<Tdata, Tdata, Tdata, int16_t>(
@@ -98,11 +101,11 @@ infiniStatus_t Descriptor::calculate(
     void *stream) const {
     switch (_dtype) {
     case INFINI_DTYPE_F16:
-        kunlun::calculate<float16>(_info, _opaque->xdnn_handle_pool, _dtype, c, beta, a, b, alpha, (KunlunStream_t)stream);
+        op::matmul::kunlun::calculate<float16>(_info, _opaque->internal, _dtype, c, beta, a, b, alpha, (kunlunStream_t)stream);
         return INFINI_STATUS_SUCCESS;
 
     case INFINI_DTYPE_F32:
-        kunlun::calculate<float>(_info, _opaque->xdnn_handle_pool, _dtype, c, beta, a, b, alpha, (KunlunStream_t)stream);
+        op::matmul::kunlun::calculate<float>(_info, _opaque->internal, _dtype, c, beta, a, b, alpha, (kunlunStream_t)stream);
         return INFINI_STATUS_SUCCESS;
 
     default:

src/infinirt/kunlun/infinirt_kunlun.cc

+#include "infinirt_kunlun.h"
+#include "../../utils.h"
+#include <xpu/runtime.h>
+#include <xpu/runtime_ex.h>
+
+typedef XPUStream kunlunStream_t;
+typedef XPUEvent kunlunEvent_t;
+
+#define CHECK_KUNLUNRT(RT_API) CHECK_INTERNAL(RT_API, XPU_SUCCESS)
+
+namespace infinirt::kunlun {
+infiniStatus_t getDeviceCount(int *count) {
+    CHECK_KUNLUNRT(xpu_device_count(count));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t setDevice(int device_id) {
+    CHECK_KUNLUNRT(xpu_set_device(device_id));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t deviceSynchronize() {
+    CHECK_KUNLUNRT(xpu_wait());
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamCreate(infinirtStream_t *stream_ptr) {
+    kunlunStream_t stream;
+    CHECK_KUNLUNRT(xpu_stream_create(&stream));
+    *stream_ptr = stream;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamDestroy(infinirtStream_t stream) {
+    CHECK_KUNLUNRT(xpu_stream_destroy((kunlunStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamSynchronize(infinirtStream_t stream) {
+    CHECK_KUNLUNRT(xpu_wait((kunlunStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+infiniStatus_t streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
+    CHECK_KUNLUNRT(xpu_stream_wait_event((kunlunStream_t)stream, (kunlunEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventCreate(infinirtEvent_t *event_ptr) {
+    kunlunEvent_t event;
+    CHECK_KUNLUNRT(xpu_event_create(&event));
+    *event_ptr = event;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventRecord(infinirtEvent_t event, infinirtStream_t stream) {
+    CHECK_KUNLUNRT(xpu_event_record((kunlunEvent_t)event, (kunlunStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr) {
+    // no event query in kunlun2
+    return INFINI_STATUS_NOT_IMPLEMENTED;
+}
+
+infiniStatus_t eventSynchronize(infinirtEvent_t event) {
+    CHECK_KUNLUNRT(xpu_event_wait((kunlunEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventDestroy(infinirtEvent_t event) {
+    CHECK_KUNLUNRT(xpu_event_destroy((kunlunEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocDevice(void **p_ptr, size_t size) {
+    CHECK_KUNLUNRT(xpu_malloc(p_ptr, static_cast<uint64_t>(size)));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocHost(void **p_ptr, size_t size) {
+    CHECK_KUNLUNRT(xpu_host_alloc(p_ptr, static_cast<uint64_t>(size), 0));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeDevice(void *ptr) {
+    CHECK_KUNLUNRT(xpu_free(ptr));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeHost(void *ptr) {
+    CHECK_KUNLUNRT(xpu_host_free(ptr));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t memcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind) {
+    switch (kind) {
+    case INFINIRT_MEMCPY_H2D:
+        CHECK_KUNLUNRT(xpu_memcpy(dst, src, static_cast<uint64_t>(size), XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+        return INFINI_STATUS_SUCCESS;
+    case INFINIRT_MEMCPY_D2H:
+        CHECK_KUNLUNRT(xpu_memcpy(dst, src, static_cast<uint64_t>(size), XPUMemcpyKind::XPU_DEVICE_TO_HOST));
+        return INFINI_STATUS_SUCCESS;
+    case INFINIRT_MEMCPY_D2D:
+        CHECK_KUNLUNRT(xpu_memcpy(dst, src, static_cast<uint64_t>(size), XPUMemcpyKind::XPU_DEVICE_TO_DEVICE));
+        return INFINI_STATUS_SUCCESS;
+    default:
+        return INFINI_STATUS_INTERNAL_ERROR;
+    }
+}
+
+infiniStatus_t memcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream) {
+    // no async memcpy func in kunlun2
+    return memcpy(dst, src, size, kind);
+}
+
+infiniStatus_t mallocAsync(void **p_ptr, size_t size, infinirtStream_t stream) {
+    CHECK_KUNLUNRT(xpu_malloc(p_ptr, static_cast<uint64_t>(size)));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeAsync(void *ptr, infinirtStream_t stream) {
+    CHECK_KUNLUNRT(xpu_free(ptr));
+    return INFINI_STATUS_SUCCESS;
+}
+
+} // namespace infinirt::kunlun

src/infinirt/kunlun/infinirt_kunlun.h

+#ifndef __INFINIRT_KUNLUN_H__
+#define __INFINIRT_KUNLUN_H__
+#include "../infinirt_impl.h"
+
+namespace infinirt::kunlun {
+#ifdef ENABLE_KUNLUN_API
+INFINIRT_DEVICE_API_IMPL
+#else
+INFINIRT_DEVICE_API_NOOP
+#endif
+} // namespace infinirt::kunlun
+
+#endif // __INFINIRT_KUNLUN_H__
\ No newline at end of file

xmake.lua

@@ -138,6 +138,9 @@ target("infinirt")
     if has_config("metax-gpu") then
         add_deps("infinirt-metax")
     end
+    if has_config("kunlun-xpu") then
+        add_deps("infinirt-kunlun")
+    end
     set_languages("cxx17")
     set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
     add_files("src/infinirt/*.cc")

xmake/kunlun.lua

@@ -18,3 +18,14 @@ target("infiniop-kunlun")
     set_languages("cxx17")
     add_files("$(projectdir)/src/infiniop/devices/kunlun/*.cc", "$(projectdir)/src/infiniop/ops/*/kunlun/*.cc")
 target_end()
+
+target("infinirt-kunlun")
+    set_kind("static")
+    add_deps("infini-utils")
+    set_languages("cxx17")
+    on_install(function (target) end)
+    -- Add include dirs
+    add_files("../src/infinirt/kunlun/*.cc")
+    add_cxflags("-lstdc++ -Wall -Werror -fPIC")
+
+target_end()