Merge pull request #109 from ut-parla/fix/eviction-owner

fix: fix bugs on changes of owner during eviction

src/python/parla/common/parray/coherence.py

@@ -477,20 +477,18 @@ 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
 
         if evict_last_copy:
             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

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()