fix: replace saturating add and sub with llvm intrinsics

Signed-off-by: Seth Stadick <[email protected]>

fix: add tests for saturating sub and add

Signed-off-by: Seth Stadick <[email protected]>

Author: sstadick

ishlib/matcher/alignment/striped_utils.mojo

@@ -1,4 +1,6 @@
 from memory import memset_zero
+from sys import llvm_intrinsic
+
 from ishlib.matcher.alignment import AlignedMemory
 
 
@@ -8,26 +10,34 @@ fn saturating_sub[
 ](lhs: SIMD[data, width], rhs: SIMD[data, width]) -> SIMD[data, width]:
     """Saturating SIMD subtraction.
 
-    https://stackoverflow.com/questions/33481295/saturating-subtract-add-for-unsigned-bytes
+    https://llvm.org/docs/LangRef.html#llvm-usub-sat-intrinsics
+    https://llvm.org/docs/LangRef.html#llvm-ssub-sat-intrinsics
     """
-    constrained[data.is_unsigned()]()
-    var resp = lhs - rhs
-    resp &= -(resp <= lhs).cast[data]()
-    return resp
+    constrained[data.is_integral()]()
+
+    @parameter
+    if data.is_unsigned():
+        return llvm_intrinsic["llvm.usub.sat", __type_of(lhs)](lhs, rhs)
+    else:
+        return llvm_intrinsic["llvm.ssub.sat", __type_of(lhs)](lhs, rhs)
 
 
 @always_inline
 fn saturating_add[
     data: DType, width: Int
 ](lhs: SIMD[data, width], rhs: SIMD[data, width]) -> SIMD[data, width]:
-    """Saturating SIMD subtraction.
+    """Saturating SIMD addition.
 
-    https://stackoverflow.com/questions/33481295/saturating-subtract-add-for-unsigned-bytes
+    https://llvm.org/docs/LangRef.html#llvm-uadd-sat-intrinsics
+    https://llvm.org/docs/LangRef.html#llvm-sadd-sat-intrinsics
     """
-    constrained[data.is_unsigned()]()
-    var resp = lhs + rhs
-    resp |= -(resp < lhs).cast[data]()
-    return resp
+    constrained[data.is_integral()]()
+
+    @parameter
+    if data.is_unsigned():
+        return llvm_intrinsic["llvm.uadd.sat", __type_of(lhs)](lhs, rhs)
+    else:
+        return llvm_intrinsic["llvm.sadd.sat", __type_of(lhs)](lhs, rhs)
 
 
 @value

tests/matcher/alignment/test_striped_utils.mojo

@@ -0,0 +1,66 @@
+from testing import assert_equal
+
+from ishlib.matcher.alignment.striped_utils import (
+    saturating_add,
+    saturating_sub,
+)
+
+
+def test_saturating_add():
+    alias dtypes = List[DType](
+        DType.uint8,
+        DType.int8,
+        DType.uint16,
+        DType.int16,
+        DType.uint32,
+        DType.int32,
+    )
+    alias widths = List[Int](4, 8, 16, 32, 64, 128)
+
+    @parameter
+    for i in range(0, len(dtypes)):
+
+        @parameter
+        for j in range(0, len(widths)):
+            alias dtype = dtypes[i]
+            alias width = widths[j]
+            alias MIN = Scalar[dtype].MIN
+            alias MAX = Scalar[dtype].MAX
+
+            var lhs = SIMD[dtype, width](MAX)
+            var rhs = SIMD[dtype, width](1)
+            var expected = SIMD[dtype, width](MAX)
+            assert_equal(saturating_add(lhs, rhs), expected)
+
+            expected = SIMD[dtype, width](2)
+            assert_equal(saturating_add(rhs, rhs), expected)
+
+
+def test_saturating_sub():
+    alias dtypes = List[DType](
+        DType.uint8,
+        DType.int8,
+        DType.uint16,
+        DType.int16,
+        DType.uint32,
+        DType.int32,
+    )
+    alias widths = List[Int](4, 8, 16, 32, 64, 128)
+
+    @parameter
+    for i in range(0, len(dtypes)):
+
+        @parameter
+        for j in range(0, len(widths)):
+            alias dtype = dtypes[i]
+            alias width = widths[j]
+            alias MIN = Scalar[dtype].MIN
+            alias MAX = Scalar[dtype].MAX
+
+            var lhs = SIMD[dtype, width](MIN)
+            var rhs = SIMD[dtype, width](1)
+            var expected = SIMD[dtype, width](MIN)
+            assert_equal(saturating_sub(lhs, rhs), expected)
+
+            expected = SIMD[dtype, width](0)
+            assert_equal(saturating_sub(rhs, rhs), expected)