diff --git a/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.cc b/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.cc index 6859243b2c2..0b5fc310664 100644 --- a/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.cc +++ b/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.cc @@ -213,7 +213,11 @@ void *OpenCLAllocator::_Malloc(MemType mem_type, void *data, size_t size, const total_size_ += size; const uint64_t max_size = ocl_runtime_->GetGlobalMemSize() * 0.8; - UNLOCK_AND_RETURN_NULL(total_size_ >= max_size, nullptr); + if (total_size_ >= max_size) { + is_oversize_ = true; + } else { + is_oversize_ = false; + } cl::Buffer *buffer = nullptr; cl::Image2D *image = nullptr; diff --git a/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.h b/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.h index 5cbe6fcab21..435e7c8cf5f 100644 --- a/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.h +++ b/mindspore/lite/src/litert/kernel/gpu/opencl/opencl_allocator.h @@ -83,6 +83,7 @@ class OpenCLAllocator : public mindspore::Allocator { } return ptr; } + bool IsOverSize() { return is_oversize_; } private: void Lock(); @@ -117,6 +118,7 @@ class OpenCLAllocator : public mindspore::Allocator { // 6 is empirical value int shift_factor_ = 6; bool lock_flag_ = true; + bool is_oversize_{false}; }; } // namespace mindspore::lite::opencl diff --git a/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.cc b/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.cc index 8e652811056..74932af6435 100644 --- a/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.cc +++ b/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.cc @@ -277,6 +277,7 @@ int OpenCLKernel::PreProcess() { } output->ResetRefCount(); } + is_oversize_kernel_ = ocl_runtime_->GetAllocator()->IsOverSize(); return RET_OK; } diff --git a/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.h b/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.h index 4eca3342151..23c57ba427b 100644 --- a/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.h +++ b/mindspore/lite/src/litert/kernel/opencl/opencl_kernel.h @@ -284,7 +284,12 @@ class OpenCLKernel : public LiteKernel { int PreProcess() override; int ReSize() override; int Run() override { return RET_ERROR; } - int PostProcess() override { return RET_OK; } + int PostProcess() override { + if (is_oversize_kernel_) { + return FreeInWorkTensor(); + } + return RET_OK; + } bool MallocDataDone(); std::string OpenCLKernelHeader(); @@ -336,6 +341,7 @@ class OpenCLKernel : public LiteKernel { cl::Event event_; void *restore_quant_data_{nullptr}; bool dequant_flag_{false}; + bool is_oversize_kernel_{false}; private: lite::opencl::OpenCLRuntimeInnerWrapper ocl_runtime_wrap_;