Skip to content
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.

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

[stdlib] Add vectorized Span.apply() #3877

Open
wants to merge 13 commits into
base: main
Choose a base branch
from
71 changes: 71 additions & 0 deletions mojo/stdlib/src/memory/span.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -473,3 +473,74 @@ struct Span[
address_space=address_space,
alignment=alignment,
](ptr=self._data, length=self._len)

fn apply[
D: DType,
O: MutableOrigin, //,
func: fn[w: Int] (SIMD[D, w]) -> SIMD[D, w],
](mut self: Span[Scalar[D], O]):
"""Apply the function to the `Span` inplace.

Parameters:
D: The DType.
O: The origin of the `Span`.
func: The function to evaluate.
"""

alias widths = (256, 128, 64, 32, 16, 8, 4)
var ptr = self.unsafe_ptr()
var length = len(self)
var processed = 0

@parameter
for i in range(len(widths)):
alias w = widths[i]

@parameter
if simdwidthof[D]() >= w:
for _ in range((length - processed) // w):
var p_curr = ptr + processed
p_curr.store(func(p_curr.load[width=w]()))
processed += w

for i in range(length - processed):
(ptr + processed + i).init_pointee_move(func(ptr[processed + i]))

fn apply[
D: DType,
O: MutableOrigin, //,
func: fn[w: Int] (SIMD[D, w]) -> SIMD[D, w],
*,
where: fn[w: Int] (SIMD[D, w]) -> SIMD[DType.bool, w],
](mut self: Span[Scalar[D], O]):
"""Apply the function to the `Span` inplace where the condition is
`True`.

Parameters:
D: The DType.
O: The origin of the `Span`.
func: The function to evaluate.
where: The condition to apply the function.
"""

alias widths = (256, 128, 64, 32, 16, 8, 4)
var ptr = self.unsafe_ptr()
var length = len(self)
var processed = 0

@parameter
for i in range(len(widths)):
alias w = widths[i]

@parameter
if simdwidthof[D]() >= w:
for _ in range((length - processed) // w):
var p_curr = ptr + processed
var vec = p_curr.load[width=w]()
p_curr.store(where(vec).select(func(vec), vec))
processed += w

for i in range(length - processed):
var vec = ptr[processed + i]
if where(vec):
(ptr + processed + i).init_pointee_move(func(vec))
41 changes: 41 additions & 0 deletions mojo/stdlib/test/memory/test_span.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,46 @@ def test_span_coerce():
takes_span(a)


def test_apply():
fn _twice[D: DType, w: Int](x: SIMD[D, w]) -> SIMD[D, w]:
return x * 2

fn _where[D: DType, w: Int](x: SIMD[D, w]) -> SIMD[DType.bool, w]:
return x % 2 == 0

def _test[D: DType]():
items = List[Scalar[D]](
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
)
twice = items
span = Span(twice)
span.apply[func = _twice[D]]()
for i in range(len(items)):
assert_true(span[i] == items[i] * 2)

# twice only even numbers
twice = items
span = Span(twice)
span.apply[func = _twice[D], where = _where[D]]()
for i in range(len(items)):
if items[i] % 2 == 0:
assert_true(span[i] == items[i] * 2)
else:
assert_true(span[i] == items[i])

_test[DType.uint8]()
_test[DType.uint16]()
_test[DType.uint32]()
_test[DType.uint64]()
_test[DType.int8]()
_test[DType.int16]()
_test[DType.int32]()
_test[DType.int64]()
_test[DType.float16]()
_test[DType.float32]()
_test[DType.float64]()


def main():
test_span_list_int()
test_span_list_str()
Expand All @@ -234,3 +274,4 @@ def main():
test_fill()
test_ref()
test_reversed()
test_apply()