163 lines
4.1 KiB
C++
163 lines
4.1 KiB
C++
// Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
#include <any>
|
|
#include <mutex>
|
|
|
|
#include "core/device_impl.h"
|
|
#include "core/types.h"
|
|
#include "cuda.h"
|
|
#include "cuda_runtime.h"
|
|
|
|
namespace mmdeploy {
|
|
|
|
using CudaTask = std::function<void(cudaStream_t)>;
|
|
|
|
class CudaPlatformImpl : public PlatformImpl {
|
|
public:
|
|
CudaPlatformImpl();
|
|
|
|
const char* GetPlatformName() const noexcept override { return "cuda"; }
|
|
|
|
shared_ptr<BufferImpl> CreateBuffer(Device device) override;
|
|
|
|
shared_ptr<StreamImpl> CreateStream(Device device) override;
|
|
|
|
shared_ptr<EventImpl> CreateEvent(Device device) override;
|
|
|
|
Result<void> Copy(const void* host_ptr, Buffer dst, size_t size, size_t dst_offset,
|
|
Stream stream) override;
|
|
|
|
Result<void> Copy(Buffer src, void* host_ptr, size_t size, size_t src_offset,
|
|
Stream stream) override;
|
|
|
|
Result<void> Copy(Buffer src, Buffer dst, size_t size, size_t src_offset, size_t dst_offset,
|
|
Stream stream) override;
|
|
|
|
Result<Stream> GetDefaultStream(int32_t device_id) override;
|
|
|
|
Allocator GetDefaultAllocator(int32_t device_id);
|
|
|
|
Device GetDevice(int device_id) { return Device(platform_id_, device_id); }
|
|
|
|
private:
|
|
static bool CheckCopyParam(size_t src_size, size_t dst_size, size_t src_offset, size_t dst_offset,
|
|
size_t copy_size);
|
|
|
|
static bool CheckCopyDevice(const Device& src, const Device& dst, const Device& st);
|
|
|
|
static Result<void> CopyImpl(Device device, const void* src, void* dst, size_t src_size,
|
|
size_t dst_size, size_t src_offset, size_t dst_offset, size_t size,
|
|
Stream st);
|
|
|
|
class PerDeviceData {
|
|
public:
|
|
explicit PerDeviceData(int device_id) : device_id_(device_id) {}
|
|
void init();
|
|
Stream& default_stream() {
|
|
init();
|
|
return default_stream_;
|
|
}
|
|
Allocator& default_allocator() {
|
|
init();
|
|
return default_allocator_;
|
|
}
|
|
|
|
private:
|
|
int device_id_;
|
|
std::once_flag init_flag_;
|
|
Stream default_stream_;
|
|
Allocator default_allocator_;
|
|
};
|
|
|
|
std::vector<std::unique_ptr<PerDeviceData>> per_device_data_storage_;
|
|
std::vector<PerDeviceData*> per_device_data_;
|
|
};
|
|
|
|
CudaPlatformImpl& gCudaPlatform();
|
|
|
|
class CudaDeviceMemory;
|
|
|
|
class CudaBufferImpl : public BufferImpl {
|
|
public:
|
|
explicit CudaBufferImpl(Device device);
|
|
|
|
Result<void> Init(size_t size, Allocator allocator, size_t alignment, uint64_t flags) override;
|
|
|
|
Result<void> Init(size_t size, std::shared_ptr<void> native, uint64_t flags) override;
|
|
|
|
Result<BufferImplPtr> SubBuffer(size_t offset, size_t size, uint64_t flags) override;
|
|
|
|
void* GetNative(ErrorCode* ec) override;
|
|
|
|
Allocator GetAllocator() const override;
|
|
|
|
size_t GetSize(ErrorCode* ec) override;
|
|
|
|
private:
|
|
std::shared_ptr<CudaDeviceMemory> memory_;
|
|
size_t offset_{0};
|
|
size_t size_{0};
|
|
};
|
|
|
|
class CudaStreamImpl : public StreamImpl {
|
|
public:
|
|
explicit CudaStreamImpl(Device device);
|
|
|
|
~CudaStreamImpl() override;
|
|
|
|
Result<void> Init(uint64_t flags) override;
|
|
|
|
Result<void> Init(std::shared_ptr<void> native, uint64_t flags) override;
|
|
|
|
Result<void> DependsOn(Event& event) override;
|
|
|
|
Result<void> Query() override;
|
|
|
|
Result<void> Wait() override;
|
|
|
|
Result<void> Submit(Kernel& kernel) override;
|
|
|
|
void* GetNative(ErrorCode* ec) override;
|
|
|
|
private:
|
|
cudaStream_t stream_;
|
|
bool owned_stream_;
|
|
std::shared_ptr<void> external_;
|
|
};
|
|
|
|
class CudaEventImpl : public EventImpl {
|
|
public:
|
|
explicit CudaEventImpl(Device device);
|
|
|
|
~CudaEventImpl() override;
|
|
|
|
Result<void> Init(uint64_t flags) override;
|
|
|
|
Result<void> Init(std::shared_ptr<void> native, uint64_t flags) override;
|
|
|
|
Result<void> Query() override;
|
|
|
|
Result<void> Record(Stream& stream) override;
|
|
|
|
Result<void> Wait() override;
|
|
|
|
void* GetNative(ErrorCode* ec) override;
|
|
|
|
private:
|
|
cudaEvent_t event_;
|
|
bool owned_event_;
|
|
std::shared_ptr<void> external_;
|
|
};
|
|
|
|
class CudaKernelImpl : public KernelImpl {
|
|
public:
|
|
explicit CudaKernelImpl(Device device, CudaTask task);
|
|
|
|
void* GetNative(ErrorCode* ec) override;
|
|
|
|
private:
|
|
CudaTask task_;
|
|
};
|
|
|
|
} // namespace mmdeploy
|