@@ -45,14 +45,14 @@ def ping_pong(
45
45
token_second_done = i + 2
46
46
if curr_rank == first_rank :
47
47
iris .store (data + offsets , i , curr_rank , peer_rank , heap_bases , mask = data_mask )
48
- iris .store (flag + offsets , token_first_done , curr_rank , peer_rank , heap_bases , mask = flag_mask )
48
+ iris .atomic_xchg (flag + offsets , token_first_done , curr_rank , peer_rank , heap_bases , mask = flag_mask )
49
49
while tl .load (flag , cache_modifier = ".cv" , volatile = True ) != token_second_done :
50
50
pass
51
51
else :
52
52
while tl .load (flag , cache_modifier = ".cv" , volatile = True ) != token_first_done :
53
53
pass
54
54
iris .store (data + offsets , i , curr_rank , peer_rank , heap_bases , mask = data_mask )
55
- iris .store (flag + offsets , token_second_done , curr_rank , peer_rank , heap_bases , mask = flag_mask )
55
+ iris .atomic_xchg (flag + offsets , token_second_done , curr_rank , peer_rank , heap_bases , mask = flag_mask )
56
56
57
57
stop = read_realtime ()
58
58
tl .store (mm_end_timestamp_ptr + peer_rank * BLOCK_SIZE + offsets , stop , time_stmp_mask )
0 commit comments