issue/87:增加摩尔线程平台的handle issue/89: 添加摩尔线程运行时

src/infiniop/devices/handle.cc

@@ -14,6 +14,9 @@
 #ifdef ENABLE_ASCEND_API
 #include "ascend/ascend_handle.h"
 #endif
+#ifdef ENABLE_MOORE_API
+#include "musa/musa_handle.h"
+#endif
 #ifdef ENABLE_KUNLUN_API
 #include "kunlun/kunlun_handle.h"
 #endif
@@ -47,6 +50,9 @@ __C infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr) {
 #ifdef ENABLE_ASCEND_API
         CREATE(INFINI_DEVICE_ASCEND, ascend);
 #endif
+#ifdef ENABLE_MOORE_API
+        CREATE(INFINI_DEVICE_MOORE, musa);
+#endif
 #ifdef ENABLE_KUNLUN_API
         CREATE(INFINI_DEVICE_KUNLUN, kunlun);
 #endif
@@ -81,6 +87,9 @@ __C infiniStatus_t infiniopDestroyHandle(infiniopHandle_t handle) {
 #ifdef ENABLE_ASCEND_API
         DELETE(INFINI_DEVICE_ASCEND, ascend);
 #endif
+#ifdef ENABLE_MOORE_API
+        DELETE(INFINI_DEVICE_MOORE, musa);
+#endif
 #ifdef ENABLE_KUNLUN_API
         DELETE(INFINI_DEVICE_KUNLUN, kunlun);
 #endif

src/infiniop/devices/musa/common_musa.h

+#include "../../../utils.h"
+#include "musa_handle.h"
+#include "pool.h"
+#include <memory>
+#include <mublas.h>
+#include <mudnn.h>
+#include <musa.h>
+#include <musa_runtime_api.h>
+
+#define CHECK_MUBLAS(API) CHECK_INTERNAL(API, MUBLAS_STATUS_SUCCESS)
+#define CHECK_MUDNN(API) CHECK_INTERNAL_MUDNN(API, ::musa::dnn::Status::SUCCESS, return INFINI_STATUS_INTERNAL_ERROR)
+
+#define CHECK_INTERNAL_MUDNN(API, EXPECT, ACTION)                                    \
+    do {                                                                             \
+        auto api_result_ = (API);                                                    \
+        if (api_result_ != (EXPECT)) {                                               \
+            std::cerr << "Error Code " << (int)api_result_ << " in `" << #API << "`" \
+                      << " from " << __func__                                        \
+                      << " at " << __FILE__ << ":" << __LINE__ << std::endl;         \
+            { ACTION; }                                                              \
+        }                                                                            \
+    } while (0)
+
+namespace device::musa {
+
+class Handle::Internal {
+    Pool<mublasHandle_t> mublas_handles;
+    Pool<::musa::dnn::Handle> mudnn_handles;
+
+    template <typename T>
+    using Fn = std::function<infiniStatus_t(T)>;
+
+public:
+    infiniStatus_t useMublas(MUstream stream, const Fn<mublasHandle_t> &f) const;
+    infiniStatus_t useMudnn(musaStream_t stream, const Fn<::musa::dnn::Handle &> &f) const;
+};
+
+} // namespace device::musa

src/infiniop/devices/musa/musa_handle.cc

+#include "common_musa.h"
+
+namespace device::musa {
+Handle::Handle(infiniDevice_t device, int device_id)
+    : InfiniopHandle{device, device_id},
+      _internal(std::make_shared<Handle::Internal>()) {}
+
+auto Handle::internal() const -> const std::shared_ptr<Internal> & {
+    return _internal;
+}
+
+infiniStatus_t Handle::Internal::useMublas(MUstream stream, const Fn<mublasHandle_t> &f) const {
+    mublasHandle_t *handle = mublas_handles.pop();
+    if (!handle) {
+        CHECK_MUBLAS(mublasCreate(handle));
+    }
+    CHECK_MUBLAS(mublasSetStream(*handle, stream));
+    CHECK_STATUS(f(*handle));
+    mublas_handles.push(handle);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Handle::Internal::useMudnn(musaStream_t stream, const Fn<::musa::dnn::Handle &> &f) const {
+    ::musa::dnn::Handle *handle = mudnn_handles.pop();
+    if (!handle) {
+        handle = new ::musa::dnn::Handle();
+    }
+    CHECK_MUDNN(handle->SetStream(stream));
+    CHECK_STATUS(f(*handle));
+    mudnn_handles.push(handle);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) {
+    *handle_ptr = new Handle(INFINI_DEVICE_MOORE, device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+} // namespace device::musa

src/infiniop/devices/musa/musa_handle.h

+#ifndef __INFINIOP_MUSA_HANDLE_H__
+#define __INFINIOP_MUSA_HANDLE_H__
+
+#include "../../handle.h"
+#include <memory>
+
+namespace device::musa {
+struct Handle : public InfiniopHandle {
+    class Internal;
+    auto internal() const -> const std::shared_ptr<Internal> &;
+
+public:
+    static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id);
+
+protected:
+    Handle(infiniDevice_t device, int device_id);
+
+private:
+    std::shared_ptr<Internal> _internal;
+};
+
+} // namespace device::musa
+
+#endif // __INFINIOP_MUSA_HANDLE_H__

src/infiniop/devices/musa/pool.h

+#ifndef __POOL_MUSA_H__
+#define __POOL_MUSA_H__
+
+#include <atomic>
+#include <mutex>
+#include <optional>
+
+template <class T>
+class Pool {
+public:
+    Pool() : _head(nullptr) {}
+
+    Pool(const Pool &) = delete;
+
+    Pool(Pool &&pool) noexcept : _head(pool._head.exchange(nullptr)) {}
+
+    ~Pool() {
+        while (this->pop()) {}
+    }
+
+    void push(T *val) const {
+        Node<T> *new_node = new Node<T>(val);
+        new_node->next = _head.load();
+        while (!_head.compare_exchange_weak(new_node->next, new_node));
+    }
+
+    T *pop() const {
+        Node<T> *top = _head.load();
+        Node<T> *new_head = nullptr;
+        do {
+            if (!top) {
+                return nullptr;
+            }
+            new_head = top->next;
+        } while (!_head.compare_exchange_weak(top, new_head));
+        return top->data;
+    }
+
+private:
+    template <class U>
+    struct Node {
+        U *data;
+        Node<U> *next;
+        Node(U *data) : data(data), next(nullptr) {}
+    };
+
+    mutable std::atomic<Node<T> *> _head;
+};
+
+#endif // __POOL_MUSA_H__

src/infinirt/musa/infinirt_musa.cc

+#include "infinirt_musa.h"
+#include "../../utils.h"
+#include <musa_runtime.h>
+#include <musa_runtime_api.h>
+
+#define CHECK_MUSART(RT_API) CHECK_INTERNAL(RT_API, musaSuccess)
+
+namespace infinirt::musa {
+infiniStatus_t getDeviceCount(int *count) {
+    CHECK_MUSART(musaGetDeviceCount(count));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t setDevice(int device_id) {
+    CHECK_MUSART(musaSetDevice(device_id));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t deviceSynchronize() {
+    CHECK_MUSART(musaDeviceSynchronize());
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamCreate(infinirtStream_t *stream_ptr) {
+    musaStream_t stream;
+    CHECK_MUSART(musaStreamCreate(&stream));
+    *stream_ptr = stream;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamDestroy(infinirtStream_t stream) {
+    CHECK_MUSART(musaStreamDestroy((musaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamSynchronize(infinirtStream_t stream) {
+    CHECK_MUSART(musaStreamSynchronize((musaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
+    return INFINI_STATUS_NOT_IMPLEMENTED;
+}
+
+infiniStatus_t eventCreate(infinirtEvent_t *event_ptr) {
+    musaEvent_t event;
+    CHECK_MUSART(musaEventCreate(&event));
+    *event_ptr = event;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventRecord(infinirtEvent_t event, infinirtStream_t stream) {
+    CHECK_MUSART(musaEventRecord((musaEvent_t)event, (musaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr) {
+    auto status = musaEventQuery((musaEvent_t)event);
+    if (status == musaSuccess) {
+        *status_ptr = INFINIRT_EVENT_COMPLETE;
+    } else if (status == musaErrorNotReady) {
+        *status_ptr = INFINIRT_EVENT_NOT_READY;
+    } else {
+        CHECK_MUSART(status);
+    }
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventSynchronize(infinirtEvent_t event) {
+    CHECK_MUSART(musaEventSynchronize((musaEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventDestroy(infinirtEvent_t event) {
+    CHECK_MUSART(musaEventDestroy((musaEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocDevice(void **p_ptr, size_t size) {
+    CHECK_MUSART(musaMalloc(p_ptr, size));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocHost(void **p_ptr, size_t size) {
+    CHECK_MUSART(musaMallocHost(p_ptr, size));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeDevice(void *ptr) {
+    CHECK_MUSART(musaFree(ptr));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeHost(void *ptr) {
+    CHECK_MUSART(musaFreeHost(ptr));
+    return INFINI_STATUS_SUCCESS;
+}
+
+musaMemcpyKind toMusaMemcpyKind(infinirtMemcpyKind_t kind) {
+    switch (kind) {
+    case INFINIRT_MEMCPY_H2D:
+        return musaMemcpyHostToDevice;
+    case INFINIRT_MEMCPY_D2H:
+        return musaMemcpyDeviceToHost;
+    case INFINIRT_MEMCPY_D2D:
+        return musaMemcpyDeviceToDevice;
+    case INFINIRT_MEMCPY_H2H:
+        return musaMemcpyHostToHost;
+    default:
+        return musaMemcpyDefault;
+    }
+}
+
+infiniStatus_t memcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind) {
+    CHECK_MUSART(musaMemcpy(dst, src, size, toMusaMemcpyKind(kind)));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t memcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream) {
+    CHECK_MUSART(musaMemcpyAsync(dst, src, size, toMusaMemcpyKind(kind), (musaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocAsync(void **p_ptr, size_t size, infinirtStream_t stream) {
+    return mallocDevice(p_ptr, size, stream);
+}
+
+infiniStatus_t freeAsync(void *ptr, infinirtStream_t stream) {
+    return freeDevice(ptr);
+}
+} // namespace infinirt::musa

src/infinirt/musa/infinirt_musa.h

+#ifndef __INFINIRT_MUSA_H__
+#define __INFINIRT_MUSA_H__
+#include "../infinirt_impl.h"
+
+namespace infinirt::musa {
+#ifdef ENABLE_MOORE_API
+INFINIRT_DEVICE_API_IMPL
+#else
+INFINIRT_DEVICE_API_NOOP
+#endif
+} // namespace infinirt::musa
+
+#endif // __INFINIRT_MUSA_H__

xmake.lua

@@ -89,7 +89,8 @@ option("moore-gpu")
 option_end()
 
 if has_config("moore-gpu") then
-    add_defines("ENABLE_MUSA_API")
+    add_defines("ENABLE_MOORE_API")
+    includes("xmake/musa.lua")
 end
 
 -- 海光
@@ -182,6 +183,9 @@ target("infiniop")
     if has_config("metax-gpu") then
         add_deps("infiniop-metax")
     end
+    if has_config("moore-gpu") then
+        add_deps("infini-musa")
+    end
     if has_config("kunlun-xpu") then
         add_deps("infiniop-kunlun")
     end

xmake/musa.lua

+local MUSA_HOME = os.getenv("MUSA_INSTALL_PATH")
+add_includedirs(MUSA_HOME .. "/include")
+add_linkdirs(MUSA_HOME .. "/lib")
+add_links("libmusa.so")
+add_links("libmusart.so")
+add_links("libmudnn.so")
+add_links("libmublas.so")
+
+rule("mu")
+    set_extensions(".mu")
+    on_load(function (target)
+        target:add("includedirs", "include")
+    end)
+
+    on_build_file(function (target, sourcefile)
+        local objectfile = target:objectfile(sourcefile)
+        os.mkdir(path.directory(objectfile))
+
+        local mcc = MUSA_HOME .. "/bin/mcc"
+        local includedirs = table.concat(target:get("includedirs"), " ")
+        local args = {"-c", sourcefile, "-o", objectfile, "-I" .. MUSA_HOME .. "/include", "-O3", "-fPIC", "-Wall", "-std=c++17", "-pthread"}
+        for _, includedir in ipairs(target:get("includedirs")) do
+            table.insert(args, "-I" .. includedir)
+        end
+
+        os.execv(mcc, args)
+        table.insert(target:objectfiles(), objectfile)
+    end)
+rule_end()
+
+target("infini-musa")
+    set_kind("static")
+    set_languages("cxx17")
+    add_files("../src/infiniop/devices/musa/*.cc", "../src/infiniop/ops/*/musa/*.cc")
+    add_files("src/ops/*/musa/*.mu", {rule = "mu"})
+    add_cxflags("-lstdc++ -Wall -fPIC")
+target_end()