Skip to content

fix: Prevent OOB access from sentinel indices in atomic_rmw (closes #1335) #1353

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
botbikamordehai2-sketch wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: Prevent OOB access from sentinel indices in atomic_rmw (closes #1335) #1353

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions nki/test/test_nki_nl_atomic_rmw.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,11 @@ def atomic_rmw_indirect_indices(in_tensor, indices_tensor, value_tensor):
M = 512

# NKI_EXAMPLE_18_BEGIN
# Initialize sentinel mask for int32 max value
sentinel_mask = indices_tile != np.iinfo(np.int32).max
# Use safe indices (0 for sentinel lanes) to avoid OOB access
safe_indices = nl.where(sentinel_mask, indices_tile, nl.zeros(indices_tile.shape, dtype=indices_tile.dtype))

value: tensor[N, M] = nl.load(value_tensor)

# dynamic indices have to be in SBUF, with shape [N, 1]
Expand All @@ -41,12 +46,13 @@ def atomic_rmw_indirect_indices(in_tensor, indices_tensor, value_tensor):

########################################################################
# Atomic read-modify-write example:
# - read: values of rmw_tensor is indexed by values from indices_tile
# - modify: incremented by value
# - write: saved back into rmw_tensor
# resulting in rmw_tensor = rmw_tensor + value
# - read: values of rmw_tensor is indexed by safe_indices
# - modify: incremented by value (only for non-sentinel lanes)
# - write: saved back into rmw_tensor using predicate to skip sentinels
# resulting in rmw_tensor = rmw_tensor + value (only for valid indices)
########################################################################
nl.atomic_rmw(rmw_tensor[indices_tile, ix], value=value, op=np.add)
nl.atomic_rmw(rmw_tensor[safe_indices, ix], value=value, op=np.add,
predicate=sentinel_mask.squeeze(axis=-1))
# NKI_EXAMPLE_18_END
return rmw_tensor

Expand Down