123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197 |
- #include <cnrt.h>
- #include <string.h>
- #include "cnstream_logging.hpp"
- #include "cnstream_common.hpp"
- #include "cnstream_syncmem.hpp"
- namespace cnstream {
- static void CNStreamMallocHost(void** ptr, size_t size) {
- void* __ptr = malloc(size);
- LOGF_IF(FRAME, nullptr == __ptr) << "Malloc memory on CPU failed, malloc size:" << size;
- *ptr = __ptr;
- }
- static void CNStreamFreeHost(void* ptr) {
- free(ptr);
- }
- CNSyncedMemory::CNSyncedMemory(size_t size) : size_(size) {}
- CNSyncedMemory::CNSyncedMemory(size_t size, int mlu_dev_id, int mlu_ddr_chn)
- : size_(size), dev_id_(mlu_dev_id), ddr_chn_(mlu_ddr_chn) {}
- CNSyncedMemory::~CNSyncedMemory() {
- std::lock_guard<std::mutex> lock(mutex_);
- if (0 == size_) return;
- if (cpu_ptr_ && own_cpu_data_) {
- free(cpu_ptr_);
- }
- if (mlu_ptr_ && own_mlu_data_) {
-
- CALL_CNRT_BY_CONTEXT(cnrtFree(mlu_ptr_), dev_id_, ddr_chn_);
- }
- }
- inline void CNSyncedMemory::ToCpu() {
- if (0 == size_) return;
- switch (head_) {
- case SyncedHead::UNINITIALIZED:
- CNStreamMallocHost(&cpu_ptr_, size_);
- memset(cpu_ptr_, 0, size_);
- head_ = SyncedHead::HEAD_AT_CPU;
- own_cpu_data_ = true;
- break;
- case SyncedHead::HEAD_AT_MLU:
- if (NULL == cpu_ptr_) {
- CNStreamMallocHost(&cpu_ptr_, size_);
- own_cpu_data_ = true;
- }
- CALL_CNRT_BY_CONTEXT(cnrtMemcpy(cpu_ptr_, mlu_ptr_, size_, CNRT_MEM_TRANS_DIR_DEV2HOST), dev_id_, ddr_chn_);
- head_ = SyncedHead::SYNCED;
- break;
- case SyncedHead::HEAD_AT_CPU:
- case SyncedHead::SYNCED:
- break;
- }
- }
- inline void CNSyncedMemory::ToMlu() {
- if (0 == size_) return;
- switch (head_) {
- case SyncedHead::UNINITIALIZED:
- CALL_CNRT_BY_CONTEXT(cnrtMalloc(&mlu_ptr_, size_), dev_id_, ddr_chn_);
- head_ = SyncedHead::HEAD_AT_MLU;
- own_mlu_data_ = true;
- break;
- case SyncedHead::HEAD_AT_CPU:
- if (NULL == mlu_ptr_) {
- CALL_CNRT_BY_CONTEXT(cnrtMalloc(&mlu_ptr_, size_), dev_id_, ddr_chn_);
- own_mlu_data_ = true;
- }
- CALL_CNRT_BY_CONTEXT(cnrtMemcpy(mlu_ptr_, cpu_ptr_, size_, CNRT_MEM_TRANS_DIR_HOST2DEV), dev_id_, ddr_chn_);
- head_ = SyncedHead::SYNCED;
- break;
- case SyncedHead::HEAD_AT_MLU:
- case SyncedHead::SYNCED:
- break;
- }
- }
- const void* CNSyncedMemory::GetCpuData() {
- std::lock_guard<std::mutex> lock(mutex_);
- ToCpu();
- return const_cast<const void*>(cpu_ptr_);
- }
- void CNSyncedMemory::SetCpuData(void* data) {
- std::lock_guard<std::mutex> lock(mutex_);
- if (0 == size_) return;
- LOGF_IF(FRAME, NULL == data) << "data is NULL.";
- if (own_cpu_data_) {
- CNStreamFreeHost(cpu_ptr_);
- }
- cpu_ptr_ = data;
- head_ = SyncedHead::HEAD_AT_CPU;
- own_cpu_data_ = false;
- }
- const void* CNSyncedMemory::GetMluData() {
- std::lock_guard<std::mutex> lock(mutex_);
- ToMlu();
- return const_cast<const void*>(mlu_ptr_);
- }
- void CNSyncedMemory::SetMluData(void* data) {
- std::lock_guard<std::mutex> lock(mutex_);
- if (0 == size_) return;
- LOGF_IF(FRAME, nullptr == data) << "data is NULL.";
- if (own_mlu_data_) {
- CALL_CNRT_BY_CONTEXT(cnrtFree(mlu_ptr_), dev_id_, ddr_chn_);
- }
- mlu_ptr_ = data;
- head_ = SyncedHead::HEAD_AT_MLU;
- own_mlu_data_ = false;
- }
- void CNSyncedMemory::SetMluDevContext(int dev_id, int ddr_chn) {
- std::lock_guard<std::mutex> lock(mutex_);
-
- cnrtDev_t dev;
- LOGF_IF(FRAME, CNRT_RET_SUCCESS != cnrtGetDeviceHandle(&dev, dev_id)) << "Can not find device by id: " << dev_id;
-
- dev_id_ = dev_id;
- ddr_chn_ = ddr_chn;
- }
- int CNSyncedMemory::GetMluDevId() const {
- std::lock_guard<std::mutex> lock(mutex_);
- return dev_id_;
- }
- int CNSyncedMemory::GetMluDdrChnId() const {
- std::lock_guard<std::mutex> lock(mutex_);
- return ddr_chn_;
- }
- void* CNSyncedMemory::GetMutableCpuData() {
- std::lock_guard<std::mutex> lock(mutex_);
- ToCpu();
- head_ = SyncedHead::HEAD_AT_CPU;
- return cpu_ptr_;
- }
- void* CNSyncedMemory::GetMutableMluData() {
- std::lock_guard<std::mutex> lock(mutex_);
- ToMlu();
- head_ = SyncedHead::HEAD_AT_MLU;
- return mlu_ptr_;
- }
- }
|