How to reverse data using SIMD in Mojo

Mojo
, ,
1 
from memory import Span
from sys.info import simdwidthof


fn reverse[D: DType, O: MutableOrigin, //](span: Span[Scalar[D], O]):
    """Reverse the elements of the Span inplace.

    Parameters:
        D: The DType of the scalars the Span stores.
        O: The origin of the Span.
    """

    alias widths = (256, 128, 64, 32, 16, 8, 4, 2)
    var ptr = span.unsafe_ptr()
    var length = len(span)
    middle, is_odd = length // 2, length % 2 != 0
    var processed = 0

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

        @parameter
        if simdwidthof[D]() >= w:
            for _ in range((middle - processed) // w):
                var lhs_ptr = ptr + processed
                var rhs_ptr = ptr + length - (processed + w)
                var lhs_v = lhs_ptr.load[width=w]().reversed()
                var rhs_v = rhs_ptr.load[width=w]().reversed()
                lhs_ptr.store(rhs_v)
                rhs_ptr.store(lhs_v)
                processed += w

    if is_odd:
        var value = ptr[middle + 1]
        (ptr + middle - 1).move_pointee_into(ptr + middle + 1)
        (ptr + middle - 1).init_pointee_move(value)


fn main():
    forward = List[Byte](1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)
    backward = List[Byte](11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1)
    span = Span(forward)
    reverse(span)
    i = 0
    for num in span:
        print(num[] == backward[i])
        i += 1
2 Likes