Skip to content

Commit

Permalink
Add __rfloordiv__ to SIMD
Browse files Browse the repository at this point in the history
The __rfloordiv__() is already supported by Int and Object. Furthermore,
the built-in SIMD has been extended to also support this method.

Signed-off-by: Peyman Barazandeh <peymanb@gmail.com>
  • Loading branch information
Peyman Barazandeh authored and peymanbr committed Apr 26, 2024
1 parent 4cbb4c8 commit 3075a7d
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 0 deletions.
17 changes: 17 additions & 0 deletions stdlib/src/builtin/simd.mojo
Expand Up @@ -653,6 +653,23 @@ struct SIMD[type: DType, size: Int = simdwidthof[type]()](
var mask = ((rhs < 0) ^ (self < 0)) & (mod != 0)
return div - mask.cast[type]()

@always_inline("nodebug")
fn __rfloordiv__(self, value: Self) -> Self:
"""Returns the division of value and self rounded down to the nearest
integer.
# Constraints:
The element type of the SIMD vector must be numeric.
Args:
value: The value to divide on.
Returns:
`value // self`.
"""
constrained[type.is_numeric(), "the type must be numeric"]()
return value // self

@always_inline("nodebug")
fn __mod__(self, rhs: Self) -> Self:
"""Returns the remainder of self divided by rhs.
Expand Down
18 changes: 18 additions & 0 deletions stdlib/test/builtin/test_simd.mojo
Expand Up @@ -102,6 +102,24 @@ def test_floordiv():
assert_equal(Float32(2) // Float32(-2), -1)
assert_equal(Float32(99) // Float32(-2), -50)

assert_equal(
SIMD[DType.int32, 4](2, 4, 1, 5) // Int32(2),
SIMD[DType.int32, 4](1, 2, 0, 2),
)
assert_equal(
SIMD[DType.float32, 4](3, -4, 1, 5) // Float32(3),
SIMD[DType.float32, 4](1, -2, 0, 1),
)

assert_equal(
Int32(2) // SIMD[DType.int32, 4](2, 4, -2, -4),
SIMD[DType.int32, 4](1, 0, -1, -1),
)
assert_equal(
Float32(3) // SIMD[DType.float32, 4](3, -4, 1, 5),
SIMD[DType.float32, 4](1, -1, 3, 0),
)


def test_mod():
assert_equal(Int32(99) % Int32(1), 0)
Expand Down

0 comments on commit 3075a7d

Please sign in to comment.