Merge pull request #77 from PanZezhong1725/issue/24

issue/24 cuda以及部分cpu运行时实现

include/infinirt.h

@@ -6,7 +6,7 @@
 typedef void *infinirtStream_t;
 typedef void *infinirtEvent_t;
 
-__C __export infiniStatus_t infinirtInit(infiniDevice_t device);
+__C __export infiniStatus_t infinirtInit();
 
 // Device
 __C __export infiniStatus_t infinirtGetAllDeviceCount(int *count_array);
@@ -47,7 +47,7 @@ __C __export infiniStatus_t infinirtFree(void *ptr);
 __C __export infiniStatus_t infinirtFreeHost(void *ptr);
 
 __C __export infiniStatus_t infinirtMemcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind);
-__C __export infiniStatus_t infinirtMemcpyAsync(void *dst, const void *src, infinirtMemcpyKind_t kind, infinirtStream_t stream);
+__C __export infiniStatus_t infinirtMemcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream);
 
 // Stream-ordered memory
 __C __export infiniStatus_t infinirtMallocAsync(void **p_ptr, size_t size, infinirtStream_t stream);

src/check.h

+#ifndef INFINI_CHECK_H
+#define INFINI_CHECK_H
+#include <iostream>
+
+#define CHECK_API_OR(API, EXPECT, ACTION)                                       \
+    do {                                                                        \
+        auto api_result_ = (API);                                               \
+        if (api_result_ != (EXPECT)) {                                          \
+            std::cerr << "Error Code " << api_result_ << " in `" << #API << "`" \
+                      << " from " << __func__                                   \
+                      << " at " << __FILE__ << ":" << __LINE__ << std::endl;    \
+            { ACTION; }                                                         \
+        }                                                                       \
+    } while (0)
+
+#define CHECK_INTERNAL(API, EXPECT) CHECK_API_OR(API, EXPECT, return INFINI_STATUS_INTERNAL_ERROR)
+
+#endif // INFINI_CHECK_H

src/infinirt/cpu/infinirt_cpu.cc

+#include "infinirt_cpu.h"
+#include <cstdlib>
+#include <cstring>
+
+namespace infinirt::cpu {
+infiniStatus_t getDeviceCount(int *count) {
+    *count = 1;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t setDevice(int device_id) {
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t deviceSynchronize() {
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamCreate(infinirtStream_t *stream_ptr) {
+    *stream_ptr = nullptr;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamDestroy(infinirtStream_t stream) {
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamSynchronize(infinirtStream_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) {
+    return INFINI_STATUS_NOT_IMPLEMENTED;
+}
+
+infiniStatus_t eventRecord(infinirtEvent_t event, infinirtStream_t stream) {
+    return INFINI_STATUS_NOT_IMPLEMENTED;
+}
+
+infiniStatus_t eventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr) {
+    return INFINI_STATUS_NOT_IMPLEMENTED;
+}
+
+infiniStatus_t eventSynchronize(infinirtEvent_t event) {
+    return INFINI_STATUS_NOT_IMPLEMENTED;
+}
+
+infiniStatus_t eventDestroy(infinirtEvent_t event) {
+    return INFINI_STATUS_NOT_IMPLEMENTED;
+}
+
+infiniStatus_t mallocDevice(void **p_ptr, size_t size) {
+    *p_ptr = std::malloc(size);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocHost(void **p_ptr, size_t size) {
+    return mallocDevice(p_ptr, size);
+}
+
+infiniStatus_t freeDevice(void *ptr) {
+    std::free(ptr);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeHost(void *ptr) {
+    return freeDevice(ptr);
+}
+
+infiniStatus_t memcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind) {
+    std::memcpy(dst, src, size);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t memcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream) {
+    return memcpy(dst, src, size, kind);
+}
+
+infiniStatus_t mallocAsync(void **p_ptr, size_t size, infinirtStream_t stream) {
+    return mallocDevice(p_ptr, size);
+}
+
+infiniStatus_t freeAsync(void *ptr, infinirtStream_t stream) {
+    return freeDevice(ptr);
+}
+
+} // namespace infinirt::cpu

src/infinirt/cpu/infinirt_cpu.h

+#ifndef __INFINIRT_CPU_H__
+#define __INFINIRT_CPU_H__
+#include "../infinirt_api.h"
+
+namespace infinirt::cpu {
+#ifdef ENABLE_CPU_API
+INFINIRT_DEVICE_API_IMPL
+#else
+INFINIRT_DEVICE_API_NOOP
+#endif
+} // namespace infinirt::cpu
+
+#endif // __INFINIRT_CPU_H__

src/infinirt/cuda/infinirt_cuda.cu

+#include "../../check.h"
+#include "infinirt_cuda.cuh"
+#include <cuda_runtime.h>
+
+#define CHECK_CUDART(RT_API) CHECK_INTERNAL(RT_API, cudaSuccess)
+
+namespace infinirt::cuda {
+infiniStatus_t getDeviceCount(int *count) {
+    CHECK_CUDART(cudaGetDeviceCount(count));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t setDevice(int device_id) {
+    CHECK_CUDART(cudaSetDevice(device_id));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t deviceSynchronize() {
+    CHECK_CUDART(cudaDeviceSynchronize());
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamCreate(infinirtStream_t *stream_ptr) {
+    cudaStream_t stream;
+    CHECK_CUDART(cudaStreamCreate(&stream));
+    *stream_ptr = stream;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamDestroy(infinirtStream_t stream) {
+    CHECK_CUDART(cudaStreamDestroy((cudaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamSynchronize(infinirtStream_t stream) {
+    CHECK_CUDART(cudaStreamSynchronize((cudaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
+    CHECK_CUDART(cudaStreamWaitEvent((cudaStream_t)stream, (cudaEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventCreate(infinirtEvent_t *event_ptr) {
+    cudaEvent_t event;
+    CHECK_CUDART(cudaEventCreate(&event));
+    *event_ptr = event;
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventRecord(infinirtEvent_t event, infinirtStream_t stream) {
+    CHECK_CUDART(cudaEventRecord((cudaEvent_t)event, (cudaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr) {
+    auto status = cudaEventQuery((cudaEvent_t)event);
+    if (status == cudaSuccess) {
+        *status_ptr = INFINIRT_EVENT_COMPLETE;
+    } else if (status == cudaErrorNotReady) {
+        *status_ptr = INFINIRT_EVENT_NOT_READY;
+    } else {
+        CHECK_CUDART(status);
+    }
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventSynchronize(infinirtEvent_t event) {
+    CHECK_CUDART(cudaEventSynchronize((cudaEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t eventDestroy(infinirtEvent_t event) {
+    CHECK_CUDART(cudaEventDestroy((cudaEvent_t)event));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocDevice(void **p_ptr, size_t size) {
+    CHECK_CUDART(cudaMalloc(p_ptr, size));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocHost(void **p_ptr, size_t size) {
+    CHECK_CUDART(cudaMallocHost(p_ptr, size));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeDevice(void *ptr) {
+    CHECK_CUDART(cudaFree(ptr));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeHost(void *ptr) {
+    CHECK_CUDART(cudaFreeHost(ptr));
+    return INFINI_STATUS_SUCCESS;
+}
+
+cudaMemcpyKind toCudaMemcpyKind(infinirtMemcpyKind_t kind) {
+    switch (kind) {
+    case INFINIRT_MEMCPY_H2D:
+        return cudaMemcpyHostToDevice;
+    case INFINIRT_MEMCPY_D2H:
+        return cudaMemcpyDeviceToHost;
+    case INFINIRT_MEMCPY_D2D:
+        return cudaMemcpyDeviceToDevice;
+    case INFINIRT_MEMCPY_H2H:
+        return cudaMemcpyHostToHost;
+    default:
+        return cudaMemcpyDefault;
+    }
+}
+
+infiniStatus_t memcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind) {
+    CHECK_CUDART(cudaMemcpy(dst, src, size, toCudaMemcpyKind(kind)));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t memcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream) {
+    CHECK_CUDART(cudaMemcpyAsync(dst, src, size, toCudaMemcpyKind(kind), (cudaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t mallocAsync(void **p_ptr, size_t size, infinirtStream_t stream) {
+    CHECK_CUDART(cudaMallocAsync(p_ptr, size, (cudaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t freeAsync(void *ptr, infinirtStream_t stream) {
+    CHECK_CUDART(cudaFreeAsync(ptr, (cudaStream_t)stream));
+    return INFINI_STATUS_SUCCESS;
+}
+} // namespace infinirt::cuda

src/infinirt/cuda/infinirt_cuda.cuh

+#ifndef __INFINIRT_CUDA_H__
+#define __INFINIRT_CUDA_H__
+#include "../infinirt_api.h"
+
+namespace infinirt::cuda {
+#ifdef ENABLE_CUDA_API
+INFINIRT_DEVICE_API_IMPL
+#else
+INFINIRT_DEVICE_API_NOOP
+#endif
+} // namespace infinirt::cuda
+
+#endif // __INFINIRT_CUDA_H__

src/infinirt/infinirt.cc

+#include "infinirt.h"
+#include "cpu/infinirt_cpu.h"
+#include "cuda/infinirt_cuda.cuh"
+
+thread_local infiniDevice_t CURRENT_DEVICE_TYPE = INFINI_DEVICE_CPU;
+thread_local int CURRENT_DEVICE_ID = 0;
+
+__C infiniStatus_t infinirtInit() {
+    return INFINI_STATUS_SUCCESS;
+}
+
+__C infiniStatus_t infinirtGetAllDeviceCount(int *count_array) {
+    if (count_array == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    for (size_t i = 0; i < INFINI_DEVICE_TYPE_COUNT; i++) {
+        auto status = infinirtGetDeviceCount(static_cast<infiniDevice_t>(i), &count_array[i]);
+        if (status != INFINI_STATUS_SUCCESS) {
+            return status;
+        }
+    }
+    return INFINI_STATUS_SUCCESS;
+}
+
+__C infiniStatus_t infinirtGetDevice(infiniDevice_t *device_ptr, int *device_id_ptr) {
+    if (device_ptr == nullptr && device_id_ptr == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    if (device_ptr != nullptr) {
+        *device_ptr = CURRENT_DEVICE_TYPE;
+    }
+    if (device_id_ptr != nullptr) {
+        *device_id_ptr = CURRENT_DEVICE_ID;
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+#define INFINIRT_CALL_DEVICE_API_AND(API, PARAMS, ACTION)   \
+    {                                                       \
+        infiniStatus_t _status;                             \
+        switch (CURRENT_DEVICE_TYPE) {                      \
+        case INFINI_DEVICE_CPU:                             \
+            _status = infinirt::cpu::API PARAMS;            \
+            break;                                          \
+        case INFINI_DEVICE_NVIDIA:                          \
+            _status = infinirt::cuda::API PARAMS;           \
+            break;                                          \
+        default:                                            \
+            return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; \
+        }                                                   \
+        { ACTION; }                                         \
+        return _status;                                     \
+    }
+
+#define INFINIRT_CALL_DEVICE_API(API, PARAMS) INFINIRT_CALL_DEVICE_API_AND(API, PARAMS, )
+
+__C infiniStatus_t infinirtGetDeviceCount(infiniDevice_t device, int *count) {
+    if (count == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    INFINIRT_CALL_DEVICE_API(getDeviceCount, (count));
+}
+
+__C infiniStatus_t infinirtSetDevice(infiniDevice_t device, int device_id) {
+    INFINIRT_CALL_DEVICE_API_AND(setDevice, (device_id),
+                                 { CURRENT_DEVICE_TYPE = device;
+                                   CURRENT_DEVICE_ID = device_id; });
+}
+
+__C infiniStatus_t infinirtDeviceSynchronize() {
+    INFINIRT_CALL_DEVICE_API(deviceSynchronize, ());
+}
+
+__C infiniStatus_t infinirtStreamCreate(infinirtStream_t *stream_ptr) {
+    INFINIRT_CALL_DEVICE_API(streamCreate, (stream_ptr));
+}
+
+__C infiniStatus_t infinirtStreamDestroy(infinirtStream_t stream) {
+    INFINIRT_CALL_DEVICE_API(streamDestroy, (stream));
+}
+
+__C infiniStatus_t infinirtStreamSynchronize(infinirtStream_t stream) {
+    INFINIRT_CALL_DEVICE_API(streamSynchronize, (stream));
+}
+
+__C infiniStatus_t infinirtStreamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
+    INFINIRT_CALL_DEVICE_API(streamWaitEvent, (stream, event));
+}
+
+__C infiniStatus_t infinirtEventCreate(infinirtEvent_t *event_ptr) {
+    INFINIRT_CALL_DEVICE_API(eventCreate, (event_ptr));
+}
+
+__C infiniStatus_t infinirtEventRecord(infinirtEvent_t event, infinirtStream_t stream) {
+    INFINIRT_CALL_DEVICE_API(eventRecord, (event, stream));
+}
+
+__C infiniStatus_t infinirtEventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr) {
+    INFINIRT_CALL_DEVICE_API(eventQuery, (event, status_ptr));
+}
+
+__C infiniStatus_t infinirtEventSynchronize(infinirtEvent_t event) {
+    INFINIRT_CALL_DEVICE_API(eventSynchronize, (event));
+}
+
+__C infiniStatus_t infinirtEventDestroy(infinirtEvent_t event) {
+    INFINIRT_CALL_DEVICE_API(eventDestroy, (event));
+}
+
+__C infiniStatus_t infinirtMalloc(void **p_ptr, size_t size) {
+    INFINIRT_CALL_DEVICE_API(mallocDevice, (p_ptr, size));
+}
+
+__C infiniStatus_t infinirtMallocHost(void **p_ptr, size_t size) {
+    INFINIRT_CALL_DEVICE_API(mallocHost, (p_ptr, size));
+}
+
+__C infiniStatus_t infinirtFree(void *ptr) {
+    INFINIRT_CALL_DEVICE_API(freeDevice, (ptr));
+}
+
+__C infiniStatus_t infinirtFreeHost(void *ptr) {
+    INFINIRT_CALL_DEVICE_API(freeHost, (ptr));
+}
+
+__C infiniStatus_t infinirtMemcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind) {
+    INFINIRT_CALL_DEVICE_API(memcpy, (dst, src, size, kind));
+}
+
+__C infiniStatus_t infinirtMemcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream) {
+    INFINIRT_CALL_DEVICE_API(memcpyAsync, (dst, src, size, kind, stream));
+}
+
+__C infiniStatus_t infinirtMallocAsync(void **p_ptr, size_t size, infinirtStream_t stream) {
+    INFINIRT_CALL_DEVICE_API(mallocAsync, (p_ptr, size, stream));
+}
+
+__C infiniStatus_t infinirtFreeAsync(void *ptr, infinirtStream_t stream) {
+    INFINIRT_CALL_DEVICE_API(freeAsync, (ptr, stream));
+}

src/infinirt/infinirt_api.h

+#ifndef __INFINIRT_API_H__
+#define __INFINIRT_API_H__
+#include "infinirt.h"
+
+#define INFINIRT_DEVICE_API(IMPL, COUNT)                                                                                          \
+    infiniStatus_t getDeviceCount(int *count) COUNT;                                                                              \
+    infiniStatus_t setDevice(int device_id) IMPL;                                                                                 \
+    infiniStatus_t deviceSynchronize() IMPL;                                                                                      \
+                                                                                                                                  \
+    infiniStatus_t streamCreate(infinirtStream_t *stream_ptr) IMPL;                                                               \
+    infiniStatus_t streamDestroy(infinirtStream_t stream) IMPL;                                                                   \
+    infiniStatus_t streamSynchronize(infinirtStream_t stream) IMPL;                                                               \
+    infiniStatus_t streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) IMPL;                                          \
+                                                                                                                                  \
+    infiniStatus_t eventCreate(infinirtEvent_t *event_ptr) IMPL;                                                                  \
+    infiniStatus_t eventRecord(infinirtEvent_t event, infinirtStream_t stream) IMPL;                                              \
+    infiniStatus_t eventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr) IMPL;                                     \
+    infiniStatus_t eventSynchronize(infinirtEvent_t event) IMPL;                                                                  \
+    infiniStatus_t eventDestroy(infinirtEvent_t event) IMPL;                                                                      \
+                                                                                                                                  \
+    infiniStatus_t mallocDevice(void **p_ptr, size_t size) IMPL;                                                                  \
+    infiniStatus_t mallocHost(void **p_ptr, size_t size) IMPL;                                                                    \
+    infiniStatus_t freeDevice(void *ptr) IMPL;                                                                                    \
+    infiniStatus_t freeHost(void *ptr) IMPL;                                                                                      \
+                                                                                                                                  \
+    infiniStatus_t memcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind) IMPL;                               \
+    infiniStatus_t memcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream) IMPL; \
+                                                                                                                                  \
+    infiniStatus_t mallocAsync(void **p_ptr, size_t size, infinirtStream_t stream) IMPL;                                          \
+    infiniStatus_t freeAsync(void *ptr, infinirtStream_t stream) IMPL;
+
+#define INFINIRT_DEVICE_API_IMPL INFINIRT_DEVICE_API(, )
+#define INFINIRT_DEVICE_API_NOOP INFINIRT_DEVICE_API({ return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; }, \
+                                                     {*count = 0; return INFINI_STATUS_SUCCESS; })
+
+#endif // __INFINIRT_API_H__

xmake.lua

@@ -152,7 +152,6 @@ target("infiniop")
     add_files("src/infiniop/devices/handle.cc")
     add_files("src/infiniop/ops/*/operator.cc")
     add_files("src/infiniop/*.cc")
-    after_build(function (target) print(YELLOW .. "You can install the libraries with \"xmake install\"" .. NC) end)
 
     set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
     add_installfiles("include/infiniop/(**/*.h)", {prefixdir = "include/infiniop"})
@@ -160,3 +159,23 @@ target("infiniop")
     add_installfiles("include/infiniop.h", {prefixdir = "include"})
     add_installfiles("include/infinicore.h", {prefixdir = "include"})
 target_end()
+
+target("infinirt")
+    set_kind("shared")
+    if has_config("cpu") then
+        add_deps("infinirt-cpu")
+    end
+    if has_config("nv-gpu") then
+        add_deps("infinirt-cuda")
+    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")
+    add_installfiles("include/infinirt.h")
+target_end()
+
+target("all")
+    set_kind("phony")
+    add_deps("infiniop", "infinirt")
+    after_build(function (target) print(YELLOW .. "[Congratulations!] Now you can install the libraries with \"xmake install\"" .. NC) end)
+target_end()

xmake/cpu.lua

@@ -15,3 +15,17 @@ target("infiniop-cpu")
         add_ldflags("-fopenmp")
     end
 target_end()
+
+target("infinirt-cpu")
+on_install(function (target) end)
+    set_kind("static")
+
+    set_warnings("all", "error")
+
+    if not is_plat("windows") then
+        add_cxflags("-fPIC")
+    end
+
+    set_languages("cxx17")
+    add_files("../src/infinirt/cpu/*.cc")
+target_end()

xmake/cuda.lua

@@ -34,3 +34,23 @@ target("infiniop-cuda")
     set_languages("cxx17")
     add_files("../src/infiniop/devices/cuda/*.cu", "../src/infiniop/ops/*/cuda/*.cu")
 target_end()
+
+target("infinirt-cuda")
+    set_kind("static")
+    on_install(function (target) end)
+    set_policy("build.cuda.devlink", true)
+
+    if is_plat("windows") then
+        add_cuflags("-Xcompiler=/utf-8", "--expt-relaxed-constexpr", "--allow-unsupported-compiler")
+    else
+        add_cuflags("-Xcompiler=-fPIC")
+        add_culdflags("-Xcompiler=-fPIC")
+        add_cxflags("-fPIC")
+    end
+    
+    set_toolchains("cuda")
+    add_links("cudart")
+
+    set_languages("cxx17")
+    add_files("../src/infinirt/cuda/*.cu")
+target_end()