[External] [stdlib] SIMD.from_bytes() and SIMD.as_bytes() (#54932)

[External] [stdlib] `SIMD.from_bytes()` and `SIMD.as_bytes()`

Similar to the Python's
[int.from_bytes()](https://docs.python.org/3/library/stdtypes.html#int.from_bytes)
and
[int.to_bytes()](https://docs.python.org/3/library/stdtypes.html#int.to_bytes)
one.

---------

Signed-off-by: Lukas Hermann <[email protected]>
Co-authored-by: Manuel Saelices <[email protected]>
Co-authored-by: Lukas Hermann <[email protected]>
Co-authored-by: abdul dakkak <[email protected]>
Closes #3966
MODULAR_ORIG_COMMIT_REV_ID: 6f00b8254c5e5bf87b10818da7021d0f682a8949

docs/changelog.md

@@ -107,6 +107,10 @@ what we publish.
 
 - Add a new `validate` parameter to the `b64decode()` function.
 
+- New `SIMD.from_bytes()` and `SIMD.as_bytes()` functions to convert a list of bytes
+  to a list of scalars and vice versa, accepting the endianess as an argument. Similar
+  to Python `int.from_bytes()` and `int.to_bytes()` functions.
+
 - The free floating functions for constructing different types have been
   deprecated for actual constructors:
 

stdlib/src/builtin/simd.mojo

@@ -34,6 +34,7 @@ from sys import (
     bitwidthof,
     has_neon,
     is_amd_gpu,
+    is_big_endian,
     is_gpu,
     is_nvidia_gpu,
     is_x86,
@@ -45,13 +46,13 @@ from sys import (
 from sys._assembly import inlined_assembly
 from sys.info import _current_arch, _is_sm_9x
 
-from bit import pop_count
+from bit import pop_count, byte_swap
 from builtin._format_float import _write_float
 from builtin.dtype import _uint_type_of_width
 from builtin.format_int import _try_write_int
 from builtin.io import _snprintf
 from documentation import doc_private
-from memory import UnsafePointer, bitcast, Span
+from memory import UnsafePointer, bitcast, memcpy, Span
 
 from utils import IndexList, StaticTuple
 from utils._visualizers import lldb_formatter_wrapping_type
@@ -268,6 +269,7 @@ struct SIMD[type: DType, size: Int](
     alias _Mask = SIMD[DType.bool, size]
 
     alias element_type = type
+
     var value: __mlir_type[`!pop.simd<`, size.value, `, `, type.value, `>`]
     """The underlying storage for the vector."""
 
@@ -1956,6 +1958,56 @@ struct SIMD[type: DType, size: Int](
 
         return bitcast[_integral_type_of[type](), size](self).cast[int_dtype]()
 
+    @staticmethod
+    fn from_bytes[
+        big_endian: Bool = is_big_endian()
+    ](bytes: InlineArray[Byte, type.sizeof()]) -> Scalar[type]:
+        """Converts a byte array to an scalar integer.
+
+        Args:
+            bytes: The byte array to convert.
+
+        Parameters:
+            big_endian: Whether the byte array is big-endian.
+
+        Returns:
+            The integer value.
+        """
+        var ptr: UnsafePointer[Scalar[type]] = bytes.unsafe_ptr().bitcast[
+            Scalar[type]
+        ]()
+        var value = ptr[]
+
+        @parameter
+        if is_big_endian() != big_endian:
+            return byte_swap(value)
+
+        return value
+
+    fn as_bytes[
+        big_endian: Bool = is_big_endian()
+    ](self) -> InlineArray[Byte, type.sizeof()]:
+        """Convert the scalar integer to a byte array.
+
+        Parameters:
+            big_endian: Whether the byte array should be big-endian.
+
+        Returns:
+            The byte array.
+        """
+        var value = self
+
+        @parameter
+        if is_big_endian() != big_endian:
+            value = byte_swap(value)
+
+        var ptr = UnsafePointer.address_of(value)
+        var array = InlineArray[Byte, type.sizeof()](fill=0)
+
+        memcpy(array.unsafe_ptr(), ptr.bitcast[Byte](), type.sizeof())
+
+        return array^
+
     fn _floor_ceil_trunc_impl[intrinsic: StringLiteral](self) -> Self:
         constrained[
             intrinsic == "llvm.floor"

stdlib/test/builtin/test_simd.mojo

@@ -1855,6 +1855,50 @@ def test_float_conversion():
     assert_almost_equal(Float64(UInt64(36)), 36.0)
 
 
+def test_from_bytes_as_bytes():
+    alias TwoBytes = InlineArray[Byte, DType.int16.sizeof()]
+    alias TwoUBytes = InlineArray[Byte, DType.uint16.sizeof()]
+    alias FourBytes = InlineArray[Byte, DType.int32.sizeof()]
+
+    assert_equal(Int16.from_bytes[big_endian=True](TwoBytes(0, 16)), 16)
+    assert_equal(Int16.from_bytes[big_endian=False](TwoBytes(0, 16)), 4096)
+    assert_equal(Int16.from_bytes[big_endian=True](TwoBytes(252, 0)), -1024)
+    assert_equal(UInt16.from_bytes[big_endian=True](TwoUBytes(252, 0)), 64512)
+    assert_equal(Int16.from_bytes[big_endian=False](TwoBytes(252, 0)), 252)
+    assert_equal(Int32.from_bytes[big_endian=True](FourBytes(0, 0, 0, 1)), 1)
+    assert_equal(
+        Int32.from_bytes[big_endian=False](FourBytes(0, 0, 0, 1)),
+        16777216,
+    )
+    assert_equal(
+        Int32.from_bytes[big_endian=True](FourBytes(1, 0, 0, 0)),
+        16777216,
+    )
+    assert_equal(
+        Int32.from_bytes[big_endian=True](FourBytes(1, 0, 0, 1)),
+        16777217,
+    )
+    assert_equal(
+        Int32.from_bytes[big_endian=False](FourBytes(1, 0, 0, 1)),
+        16777217,
+    )
+    assert_equal(
+        Int32.from_bytes[big_endian=True](FourBytes(255, 0, 0, 0)),
+        -16777216,
+    )
+    for x_ref in List[Int16](10, 100, -12, 0, 1, -1, 1000, -1000):
+        x = x_ref[]
+
+        @parameter
+        for b in range(2):
+            assert_equal(
+                Int16.from_bytes[big_endian=b](
+                    Int16(x).as_bytes[big_endian=b]()
+                ),
+                x,
+            )
+
+
 def test_reversed():
     fn test[D: DType]() raises:
         assert_equal(SIMD[D, 4](1, 2, 3, 4).reversed(), SIMD[D, 4](4, 3, 2, 1))
@@ -1885,6 +1929,7 @@ def main():
     test_extract()
     test_floor()
     test_floordiv()
+    test_from_bytes_as_bytes()
     test_iadd()
     test_indexing()
     test_insert()