From b3f0bc93f6b09d50c1cef1a3e3ab83ad922e1fba Mon Sep 17 00:00:00 2001
From: yinengy <yinengy@utexas.edu>
Date: Fri, 26 May 2023 17:44:04 +0000
Subject: [PATCH] fix: fix bugs on changes of owner during eviction

---
 src/python/parla/common/parray/coherence.py | 19 ++++++++++++-------
 testing/python/test_parray.py               | 16 +++++++++++++++-
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/src/python/parla/common/parray/coherence.py b/src/python/parla/common/parray/coherence.py
index b2d6350e..6ff3b37b 100644
--- a/src/python/parla/common/parray/coherence.py
+++ b/src/python/parla/common/parray/coherence.py
@@ -477,13 +477,8 @@ def evict(self, device_id: int, keep_one_copy: bool = True) -> List[MemoryOperat
                 # this device owns the last copy
                 if new_owner is None:
                     evict_last_copy = True
-            else:
-                # update states
-                self._local_states[device_id] = self.INVALID
-                self._versions[device_id] = -1
-                self._is_complete[device_id] = None
-
-                operations.append(MemoryOperation.evict(device_id))
+                else:
+                    self.owner = new_owner
         else:  # Modified, this device owns the last copy
             evict_last_copy = True
 
@@ -491,6 +486,9 @@ def evict(self, device_id: int, keep_one_copy: bool = True) -> List[MemoryOperat
             if keep_one_copy:  # write back to CPU
                 if device_id != CPU_INDEX:              
                     operations.extend(self._write_back_to(CPU_INDEX, self.MODIFIED, on_different_device=True, this_device_id=device_id))
+                    # special case, since `this_device_id` is set, _write_back will not evict this devic
+                    # need to do it manually
+                    operations.append(MemoryOperation.evict(device_id))
 
                     self.owner = CPU_INDEX
                     self._local_states[CPU_INDEX] = self.MODIFIED
@@ -500,5 +498,12 @@ def evict(self, device_id: int, keep_one_copy: bool = True) -> List[MemoryOperat
             else:
                 # do nothing and report error
                 return [MemoryOperation.error()]
+        else:
+            # update states for eviction
+            self._local_states[device_id] = self.INVALID
+            self._versions[device_id] = -1
+            self._is_complete[device_id] = None
+
+            operations.append(MemoryOperation.evict(device_id))
 
         return operations
diff --git a/testing/python/test_parray.py b/testing/python/test_parray.py
index b49d3334..2a41a6ef 100644
--- a/testing/python/test_parray.py
+++ b/testing/python/test_parray.py
@@ -87,7 +87,10 @@ def check_array_update():
                 assert a._coherence.owner == 1
 
             @spawn(ts[5], dependencies=[ts[4]], placement=gpu(1), inout=[(a,0)])
-            def check_array_update():
+            def check_array_evict():
+                a.print_overview()
+
+                print(a)
                 result = a.evict(1, False)
 
                 assert result == False
@@ -96,6 +99,17 @@ def check_array_update():
 
                 assert result == True
                 assert a._coherence.owner == -1
+                assert a._array._buffer[1] is None
+
+            @spawn(ts[6], dependencies=[ts[5]], placement=gpu(0), input=[(a,0)])
+            def check_array_evict2():
+                assert a._array._buffer[-1] is not None
+                result = a.evict(-1, False)
+
+                assert result == True
+
+                assert a._coherence.owner == 0
+                assert a._array._buffer[-1] is None
 
 if __name__=="__main__":
     test_parray_creation()