From bdfae5bbcaa916b97ad011df79cdcf0e5557cd11 Mon Sep 17 00:00:00 2001
From: chenzupeng <chenzupeng@huawei.com>
Date: Sat, 26 Dec 2020 15:39:44 +0800
Subject: [PATCH] fix bug: memory alloc over size

---
 mindspore/lite/src/runtime/opencl/opencl_allocator.cc | 4 ++--
 mindspore/lite/src/runtime/opencl/opencl_runtime.cc   | 4 ++--
 mindspore/lite/src/runtime/opencl/opencl_runtime.h    | 2 ++
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/mindspore/lite/src/runtime/opencl/opencl_allocator.cc b/mindspore/lite/src/runtime/opencl/opencl_allocator.cc
index b19d7d19a0..200a599406 100644
--- a/mindspore/lite/src/runtime/opencl/opencl_allocator.cc
+++ b/mindspore/lite/src/runtime/opencl/opencl_allocator.cc
@@ -137,7 +137,7 @@ void *OpenCLAllocator::Malloc(size_t size, const std::vector<size_t> &img_size,
     uint32_t image_alignment = ocl_runtime_->GetImagePitchAlignment();
     size = UP_ROUND(img_size[0], image_alignment) * img_size[1] * dtype_size;
   }
-  if (size > MAX_MALLOC_SIZE) {
+  if (size > ocl_runtime_->GetMaxAllocSize()) {
     MS_LOG(ERROR) << "MallocData out of max_size, size: " << size;
     return nullptr;
   }
@@ -148,7 +148,7 @@ void *OpenCLAllocator::Malloc(size_t size, const std::vector<size_t> &img_size,
     return host_ptr;
   }
   total_size_ += size;
-  const uint64_t max_size = ocl_runtime_->GetGlobalMemSize();
+  const uint64_t max_size = ocl_runtime_->GetGlobalMemSize() * 0.8;
   if (total_size_ >= max_size) {
     UnLock();
     MS_LOG(ERROR) << "Mem pool out of max_size, total size: " << total_size_ << ", max size: " << max_size;
diff --git a/mindspore/lite/src/runtime/opencl/opencl_runtime.cc b/mindspore/lite/src/runtime/opencl/opencl_runtime.cc
index 1e44fc0cc0..2d32db5d6b 100644
--- a/mindspore/lite/src/runtime/opencl/opencl_runtime.cc
+++ b/mindspore/lite/src/runtime/opencl/opencl_runtime.cc
@@ -226,11 +226,11 @@ int OpenCLRuntime::Init() {
     }
   }
   global_memery_size_ = device_->getInfo<CL_DEVICE_GLOBAL_MEM_SIZE>();
-
+  max_alloc_size_ = device_->getInfo<CL_DEVICE_MAX_MEM_ALLOC_SIZE>();
   MS_LOG(INFO) << "Address space bits: " << device_->getInfo<CL_DEVICE_ADDRESS_BITS>();
   MS_LOG(INFO) << "Global Mem Size: " << global_memery_size_;
   MS_LOG(INFO) << "Global Mem Cache Size: " << global_memery_cachesize_;
-  MS_LOG(INFO) << "Max Alloc Size: " << device_->getInfo<CL_DEVICE_MAX_MEM_ALLOC_SIZE>();
+  MS_LOG(INFO) << "Max Alloc Size: " << max_alloc_size_;
   MS_LOG(INFO) << "Compute Unit: " << compute_units_;
   MS_LOG(INFO) << "Clock Frequency: " << max_freq_ << " MHz";
 
diff --git a/mindspore/lite/src/runtime/opencl/opencl_runtime.h b/mindspore/lite/src/runtime/opencl/opencl_runtime.h
index 92e1e8de63..d8704d09a7 100644
--- a/mindspore/lite/src/runtime/opencl/opencl_runtime.h
+++ b/mindspore/lite/src/runtime/opencl/opencl_runtime.h
@@ -62,6 +62,7 @@ class OpenCLRuntime {
   uint64_t GetMaxWorkGroupSize(const cl::Kernel &kernel);
   uint32_t GetSubGroupSize(const cl::Kernel &kernel, const cl::NDRange &range = cl::NullRange);
   uint64_t GetGlobalMemSize() { return global_memery_size_; }
+  uint64_t GetMaxAllocSize() { return max_alloc_size_; }
   GpuInfo GetGpuInfo();
   bool GetFp16Enable() const;
   bool SetFp16Enable(bool enable);
@@ -170,6 +171,7 @@ class OpenCLRuntime {
   cl::Program binary_program_{0};
   uint64_t global_memery_cachesize_{0};
   uint64_t global_memery_size_{0};
+  uint64_t max_alloc_size_{0};
   int max_work_group_size_{1};
   uint32_t compute_units_{0};
   uint32_t max_freq_{0};