Leak fix. (tensorflow#73)

kinshu · Feb 17, 2020 · 2e8dc93 · 2e8dc93
1 parent 3e8cc3a
commit 2e8dc93
Show file tree

Hide file tree

Showing 6 changed files with 10 additions and 13 deletions.
diff --git a/tensorflow_quantum/core/qsim/mux_test.cc b/tensorflow_quantum/core/qsim/mux_test.cc
@@ -27,6 +27,7 @@ namespace {
 TEST(MuxTest, GetStateSpace) {
   auto simulator = GetStateSpace(1, 1);
   EXPECT_FALSE(simulator == nullptr);
+  delete simulator;
 }
 
 }  // namespace

diff --git a/tensorflow_quantum/core/qsim/state_space.cc b/tensorflow_quantum/core/qsim/state_space.cc
@@ -106,8 +106,6 @@ bool StateSpace::Valid() const {
 
 float* StateSpace::GetRawState() const { return state_; };
 
-void StateSpace::SetRawState(float* state_update) { state_ = state_update; }
-
 uint64_t StateSpace::GetNumEntries() const { return size_; }
 
 uint64_t StateSpace::GetDimension() const { return size_ / 2; }

diff --git a/tensorflow_quantum/core/qsim/state_space.h b/tensorflow_quantum/core/qsim/state_space.h
@@ -51,9 +51,6 @@ class StateSpace {
   // Pointer to the raw state managed by this StateSpace
   float* GetRawState() const;
 
-  // Replace the pointer in this object with a new one
-  void SetRawState(float* state_update);
-
   // Dimension of the complex Hilbert space represented by this StateSpace
   uint64_t GetDimension() const;
 
@@ -105,7 +102,7 @@ class StateSpace {
   // Set the amplitude at the given state index
   virtual void SetAmpl(const uint64_t i, const std::complex<float>& val) = 0;
 
- private:
+ protected:
   float* state_;
   uint64_t size_;
   uint64_t num_qubits_;

diff --git a/tensorflow_quantum/core/qsim/state_space_avx.cc b/tensorflow_quantum/core/qsim/state_space_avx.cc
@@ -37,10 +37,10 @@ StateSpaceAVX::~StateSpaceAVX() { DeleteState(); }
 StateSpaceType StateSpaceAVX::GetType() const { return StateSpaceType::AVX; }
 
 void StateSpaceAVX::CreateState() {
-  SetRawState((float*)qsim::_aligned_malloc(sizeof(float) * GetNumEntries()));
+  state_ = (float*)qsim::_aligned_malloc(sizeof(float) * GetNumEntries());
 }
 
-void StateSpaceAVX::DeleteState() { qsim::_aligned_free(GetRawState()); }
+void StateSpaceAVX::DeleteState() { qsim::_aligned_free(state_); }
 
 StateSpace* StateSpaceAVX::Clone() const {
   StateSpaceAVX* state_copy =

diff --git a/tensorflow_quantum/core/qsim/state_space_slow.cc b/tensorflow_quantum/core/qsim/state_space_slow.cc
@@ -34,13 +34,14 @@ StateSpaceSlow::~StateSpaceSlow() { DeleteState(); }
 StateSpaceType StateSpaceSlow::GetType() const { return StateSpaceType::SLOW; }
 
 void StateSpaceSlow::CreateState() {
-  SetRawState((float*)malloc(sizeof(float) * GetNumEntries()));
+  state_ = (float*)malloc(sizeof(float) * GetNumEntries());
 }
 
-void StateSpaceSlow::DeleteState() { free(GetRawState()); }
+void StateSpaceSlow::DeleteState() { free(state_); }
 
 StateSpace* StateSpaceSlow::Clone() const {
-  StateSpace* state_copy = new StateSpaceSlow(GetNumQubits(), GetNumThreads());
+  StateSpaceSlow* state_copy =
+      new StateSpaceSlow(GetNumQubits(), GetNumThreads());
   return state_copy;
 }
 

diff --git a/tensorflow_quantum/core/qsim/state_space_sse.cc b/tensorflow_quantum/core/qsim/state_space_sse.cc
@@ -38,10 +38,10 @@ StateSpaceSSE::~StateSpaceSSE() { DeleteState(); }
 StateSpaceType StateSpaceSSE::GetType() const { return StateSpaceType::SSE; }
 
 void StateSpaceSSE::CreateState() {
-  SetRawState((float*)qsim::_aligned_malloc(sizeof(float) * GetNumEntries()));
+  state_ = (float*)qsim::_aligned_malloc(sizeof(float) * GetNumEntries());
 }
 
-void StateSpaceSSE::DeleteState() { qsim::_aligned_free(GetRawState()); }
+void StateSpaceSSE::DeleteState() { qsim::_aligned_free(state_); }
 
 StateSpace* StateSpaceSSE::Clone() const {
   StateSpaceSSE* state_copy =