Skip to content

Commit

Permalink
fix halo exchange (#30)
Browse files Browse the repository at this point in the history
* fix halo exchange

* cleanup

---------

Co-authored-by: Connor Ward <[email protected]>
  • Loading branch information
ksagiyam and connorjward authored Apr 29, 2024
1 parent 15e870d commit 540985c
Showing 1 changed file with 15 additions and 10 deletions.
25 changes: 15 additions & 10 deletions pyop3/lang.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,7 @@ def _buffer_exchanges(buffer, intent, *, touches_ghost_points):
# reductions
assert intent in {INC, MIN_WRITE, MIN_RW, MAX_WRITE, MAX_RW}
# We don't need to update roots if performing the same reduction
# again. For example we can increment into an buffer as many times
# again. For example we can increment into a buffer as many times
# as we want. The reduction only needs to be done when the
# data is read.
if buffer._roots_valid or intent == buffer._pending_reduction:
Expand All @@ -361,6 +361,20 @@ def _buffer_exchanges(buffer, intent, *, touches_ghost_points):
initializers.append(buffer._reduce_leaves_to_roots_begin)
reductions.append(buffer._reduce_leaves_to_roots_end)

# set leaves to appropriate nil value
if intent == INC:
nil = 0
elif intent in {MIN_WRITE, MIN_RW}:
nil = dtype_limits(buffer.dtype).max
else:
assert intent in {MAX_WRITE, MAX_RW}
nil = dtype_limits(buffer.dtype).min

def _init_nil():
buffer._data[buffer.sf.ileaf] = nil

reductions.append(_init_nil)

# We are modifying owned values so the leaves must now be wrong
buffer._leaves_valid = False

Expand All @@ -370,15 +384,6 @@ def _buffer_exchanges(buffer, intent, *, touches_ghost_points):
else:
buffer._pending_reduction = intent

# set leaves to appropriate nil value
if intent == INC:
buffer._data[buffer.sf.ileaf] = 0
elif intent in {MIN_WRITE, MIN_RW}:
buffer._data[buffer.sf.ileaf] = dtype_limits(buffer.dtype).max
else:
assert intent in {MAX_WRITE, MAX_RW}
buffer._data[buffer.sf.ileaf] = dtype_limits(buffer.dtype).min

return tuple(initializers), tuple(reductions), tuple(broadcasts)


Expand Down

0 comments on commit 540985c

Please sign in to comment.