From 940c7ff4a0203486bdd2cb8b6dd0ec534ea3d06b Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@protonmail.com>
Date: Thu, 26 Nov 2020 16:17:34 +0100
Subject: [PATCH] Fix memory allocation limit

Signed-off-by: trivernis <trivernis@protonmail.com>
---
 src/kernel_controller/mod.rs | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/kernel_controller/mod.rs b/src/kernel_controller/mod.rs
index f38f6bb..1fa7f51 100644
--- a/src/kernel_controller/mod.rs
+++ b/src/kernel_controller/mod.rs
@@ -4,8 +4,11 @@
  * See LICENSE for more information
  */
 
+use ocl::core::DeviceInfo;
+use ocl::enums::DeviceInfoResult;
 use ocl::ProQue;
 use parking_lot::Mutex;
+use std::mem::size_of;
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -19,9 +22,28 @@ impl KernelController {
             .src(include_str!("kernel.cl"))
             .dims(1 << 20)
             .build()?;
+        let device = pro_que.device();
+        println!("Using device {}", device.name()?);
+        println!("Vendor: {}", device.vendor()?);
+        println!(
+            "Global Mem Size: {} bytes",
+            device.info(DeviceInfo::GlobalMemSize)?
+        );
+        println!(
+            "Max Mem Alloc: {} bytes",
+            device.info(DeviceInfo::MaxMemAllocSize)?
+        );
+        println!();
         Ok(Self { pro_que })
     }
 
+    fn available_memory(&self) -> ocl::Result<u64> {
+        match self.pro_que.device().info(DeviceInfo::GlobalMemSize)? {
+            DeviceInfoResult::GlobalMemSize(size) => Ok(size),
+            _ => Ok(0),
+        }
+    }
+
     pub fn filter_primes(&self, input: Vec<i64>) -> ocl::Result<Vec<i64>> {
         lazy_static::lazy_static! {static ref PRIME_CACHE: Arc<Mutex<Vec<i64>>> = Arc::new(Mutex::new(get_lower_primes(2048)));}
 
@@ -73,7 +95,8 @@ impl KernelController {
         let start = Instant::now();
         let mut prime_cache = PRIME_CACHE.lock();
 
-        if prime_cache.len() < 1024 * 1024 * 1024 {
+        if (prime_cache.len() + primes.len()) * size_of::<i64>() < self.available_memory()? as usize
+        {
             prime_cache.append(&mut primes.clone());
             prime_cache.sort();
             prime_cache.dedup();