Merge pull request #12 from keithcolbert/feature/alias

Add alias and assign for Ctxt memory allocations
vernamlab · Feb 9, 2019 · 1af2f2b · 1af2f2b
2 parents c795c58 + c137440
commit 1af2f2b
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 16 deletions.
diff --git a/cufhe/include/cufhe.h b/cufhe/include/cufhe.h
@@ -106,6 +106,9 @@ struct PubKey {
 struct Ctxt {
   Ctxt(bool is_alias = false);
   ~Ctxt();
+  Ctxt(const Ctxt& that) = delete;
+  Ctxt& operator=(const Ctxt& that) = delete;
+  void assign(void* host_ptr, void* device_ptr);
   LWESample* lwe_sample_;
   MemoryDeleter lwe_sample_deleter_;
   LWESample* lwe_sample_device_;

diff --git a/cufhe/lib/cufhe.cc b/cufhe/lib/cufhe.cc
@@ -279,25 +279,45 @@ PubKey::~PubKey() {
 }
 
 Ctxt::~Ctxt() {
-  if(lwe_sample_ != nullptr)
-  {
-    lwe_sample_deleter_(lwe_sample_->data());
-    lwe_sample_deleter_ = nullptr;
+  if(lwe_sample_ != nullptr) {
+    if(lwe_sample_deleter_ != nullptr) {
+      lwe_sample_deleter_(lwe_sample_->data());
+      lwe_sample_deleter_ = nullptr;
+    }
+
     lwe_sample_->set_data(nullptr);
     delete lwe_sample_;
     lwe_sample_ = nullptr;
   }
 
-  if(lwe_sample_device_ != nullptr)
-  {
-    lwe_sample_device_deleter_(lwe_sample_device_->data());
-    lwe_sample_device_deleter_ = nullptr;
+  if(lwe_sample_device_ != nullptr) {
+    if(lwe_sample_device_deleter_ != nullptr) {
+      lwe_sample_device_deleter_(lwe_sample_device_->data());
+      lwe_sample_device_deleter_ = nullptr;
+    }
+
     lwe_sample_device_->set_data(nullptr);
     delete lwe_sample_device_;
     lwe_sample_device_ = nullptr;
   }
 }
 
+void Ctxt::assign(void* host_ptr, void* device_ptr) {
+  if(lwe_sample_deleter_ != nullptr) {
+    lwe_sample_deleter_(lwe_sample_->data());
+    lwe_sample_deleter_ = nullptr;
+  }
+
+  lwe_sample_->set_data((LWESample::PointerType)host_ptr);
+
+  if(lwe_sample_device_deleter_ != nullptr) {
+    lwe_sample_device_deleter_(lwe_sample_device_->data());
+    lwe_sample_device_deleter_ = nullptr;
+  }
+
+  lwe_sample_device_->set_data((LWESample::PointerType)device_ptr);
+}
+
 void SetSeed(uint32_t seed) {
   srand(seed);
   RandomGeneratorSetSeed(&seed, 1);

diff --git a/cufhe/lib/cufhe_gpu.cu b/cufhe/lib/cufhe_gpu.cu
@@ -32,15 +32,23 @@ Ctxt::Ctxt(bool is_alias) {
   Param* param = GetDefaultParam();
 
   lwe_sample_ = new LWESample(param->lwe_n_);
-  //pair = AllocatorBoth::New(lwe_sample_->SizeMalloc());
-  pair = AllocatorCPU::New(lwe_sample_->SizeMalloc());
-  lwe_sample_->set_data((LWESample::PointerType)pair.first);
-  lwe_sample_deleter_ = pair.second;
-
   lwe_sample_device_ = new LWESample(param->lwe_n_);
-  pair = AllocatorGPU::New(lwe_sample_device_->SizeMalloc());
-  lwe_sample_device_->set_data((LWESample::PointerType)pair.first);
-  lwe_sample_device_deleter_ = pair.second;
+
+  if (is_alias) {
+    lwe_sample_->set_data(nullptr);
+    lwe_sample_deleter_ = nullptr;
+    lwe_sample_device_->set_data(nullptr);
+    lwe_sample_device_deleter_ = nullptr;
+  } else {
+    //pair = AllocatorBoth::New(lwe_sample_->SizeMalloc());
+    pair = AllocatorCPU::New(lwe_sample_->SizeMalloc());
+    lwe_sample_->set_data((LWESample::PointerType)pair.first);
+    lwe_sample_deleter_ = pair.second;
+
+    pair = AllocatorGPU::New(lwe_sample_device_->SizeMalloc());
+    lwe_sample_device_->set_data((LWESample::PointerType)pair.first);
+    lwe_sample_device_deleter_ = pair.second;
+  }
 }
 
 } // namespace cufhe