imops 0.8.2__cp39-cp39-win32.whl → 0.8.3__cp39-cp39-win32.whl

imops/src/_fast_numeric.pyx

@@ -8,10 +8,19 @@
 import numpy as np
 
 cimport numpy as np
+from libc.stdint cimport uint16_t
 
 from cython.parallel import prange
 
 
+# https://stackoverflow.com/questions/47421443/using-half-precision-numpy-floats-in-cython
+cdef extern from "numpy/halffloat.h":
+    ctypedef uint16_t npy_half
+
+    float npy_half_to_float(npy_half h) nogil
+    npy_half npy_float_to_half(float f) nogil
+
+
 ctypedef fused NUM:
     short
     int
@@ -20,45 +29,214 @@ ctypedef fused NUM:
     double
 
 
-def _parallel_sum(NUM[:] nums, Py_ssize_t num_threads) -> NUM:
-    cdef NUM res = 0
-    cdef Py_ssize_t i, len_nums = len(nums)
+ctypedef fused NUM_AND_NPY_HALF:
+    NUM
+    npy_half
 
-    for i in prange(len_nums, num_threads=num_threads, nogil=True):
-        res += nums[i]
 
-    return res
-
-
-def _parallel_pointwise_mul(
+# TODO: Generalize code below to n-d
+def _pointwise_add_array_3d(
     NUM[:, :, :] nums1,
     NUM[:, :, :] nums2,
-    Py_ssize_t[:] res_shape,
-    Py_ssize_t num_threads
+    NUM[:, :, :] out,
+    Py_ssize_t num_threads,
 ) -> np.ndarray:
-    cdef NUM[:, :, ::1] contiguous_nums1 = np.ascontiguousarray(nums1), contiguous_nums2 = np.ascontiguousarray(nums2)
-    cdef Py_ssize_t rows = res_shape[0], cols = res_shape[1], dims = res_shape[2]
+    cdef Py_ssize_t rows = out.shape[0], cols = out.shape[1], dims = out.shape[2]
+    cdef Py_ssize_t i, j, k
+
+    for i in prange(rows, nogil=True, num_threads=num_threads):
+        for j in prange(cols):
+            for k in prange(dims):
+                out[i, j, k] = nums1[i, j, k] + nums2[i, j, k]
+
+    return np.asarray(out)
 
-    cdef char[:] broadcast_mask1 = np.array([x == y for x, y in zip(res_shape, nums1.shape)], dtype=np.int8)
-    cdef char[:] broadcast_mask2 = np.array([x == y for x, y in zip(res_shape, nums2.shape)], dtype=np.int8)
 
-    cdef NUM[:, :, ::1] mul = np.empty_like(nums1, shape=res_shape)
+def _pointwise_add_array_4d(
+    NUM[:, :, :, :] nums1,
+    NUM[:, :, :, :] nums2,
+    NUM[:, :, :, :] out,
+    Py_ssize_t num_threads,
+) -> np.ndarray:
+    cdef Py_ssize_t dim1 = out.shape[0], dim2 = out.shape[1], dim3 = out.shape[2], dim4 = out.shape[3]
+    cdef Py_ssize_t i1, i2, i3, i4
+
+    for i1 in prange(dim1, nogil=True, num_threads=num_threads):
+        for i2 in prange(dim2):
+            for i3 in prange(dim3):
+                for i4 in prange(dim4):
+                    out[i1, i2, i3, i4] = nums1[i1, i2, i3, i4] + nums2[i1, i2, i3, i4]
+
+    return np.asarray(out)
+
+
+def _pointwise_add_value_3d(
+    NUM[:, :, :] nums,
+    NUM value,
+    NUM[:, :, :] out,
+    Py_ssize_t num_threads,
+) -> np.ndarray:
+    cdef Py_ssize_t rows = out.shape[0], cols = out.shape[1], dims = out.shape[2]
     cdef Py_ssize_t i, j, k
 
     for i in prange(rows, nogil=True, num_threads=num_threads):
         for j in prange(cols):
             for k in prange(dims):
-                mul[i, j, k] = (
-                    contiguous_nums1[
-                        i * broadcast_mask1[0],
-                        j * broadcast_mask1[1],
-                        k * broadcast_mask1[2]
-                    ] *
-                    contiguous_nums2[
-                        i * broadcast_mask2[0],
-                        j * broadcast_mask2[1],
-                        k * broadcast_mask2[2]
-                    ]
-                )
-
-    return np.asarray(mul)
+                out[i, j, k] = nums[i, j, k] + value
+
+    return np.asarray(out)
+
+
+def _pointwise_add_value_4d(
+    NUM[:, :, :, :] nums,
+    NUM value,
+    NUM[:, :, :, :] out,
+    Py_ssize_t num_threads,
+) -> np.ndarray:
+    cdef Py_ssize_t dim1 = out.shape[0], dim2 = out.shape[1], dim3 = out.shape[2], dim4 = out.shape[3]
+    cdef Py_ssize_t i1, i2, i3, i4
+
+    for i1 in prange(dim1, nogil=True, num_threads=num_threads):
+        for i2 in prange(dim2):
+            for i3 in prange(dim3):
+                for i4 in prange(dim4):
+                    out[i1, i2, i3, i4] = nums[i1, i2, i3, i4] + value
+
+    return np.asarray(out)
+
+
+def _pointwise_add_array_3d_fp16(
+    npy_half[:, :, :] nums1,
+    npy_half[:, :, :] nums2,
+    npy_half[:, :, :] out,
+    Py_ssize_t num_threads,
+) -> np.ndarray:
+    cdef Py_ssize_t rows = out.shape[0], cols = out.shape[1], dims = out.shape[2]
+    cdef Py_ssize_t i, j, k
+
+    for i in prange(rows, nogil=True, num_threads=num_threads):
+        for j in prange(cols):
+            for k in prange(dims):
+                out[i, j, k] = (npy_float_to_half(npy_half_to_float(nums1[i, j, k]) +
+                                npy_half_to_float(nums2[i, j, k])))
+
+    return np.asarray(out)
+
+
+def _pointwise_add_array_4d_fp16(
+    npy_half[:, :, :, :] nums1,
+    npy_half[:, :, :, :] nums2,
+    npy_half[:, :, :, :] out,
+    Py_ssize_t num_threads,
+) -> np.ndarray:
+    cdef Py_ssize_t dim1 = out.shape[0], dim2 = out.shape[1], dim3 = out.shape[2], dim4 = out.shape[3]
+    cdef Py_ssize_t i1, i2, i3, i4
+
+    for i1 in prange(dim1, nogil=True, num_threads=num_threads):
+        for i2 in prange(dim2):
+            for i3 in prange(dim3):
+                for i4 in prange(dim4):
+                    out[i1, i2, i3, i4] = (npy_float_to_half(npy_half_to_float(nums1[i1, i2, i3, i4]) +
+                                           npy_half_to_float(nums2[i1, i2, i3, i4])))
+
+    return np.asarray(out)
+
+
+def _pointwise_add_value_3d_fp16(
+    npy_half[:, :, :] nums,
+    npy_half value,
+    npy_half[:, :, :] out,
+    Py_ssize_t num_threads,
+) -> np.ndarray:
+    cdef Py_ssize_t rows = out.shape[0], cols = out.shape[1], dims = out.shape[2]
+    cdef Py_ssize_t i, j, k
+
+    for i in prange(rows, nogil=True, num_threads=num_threads):
+        for j in prange(cols):
+            for k in prange(dims):
+                out[i, j, k] = npy_float_to_half(npy_half_to_float(nums[i, j, k]) + npy_half_to_float(value))
+
+    return np.asarray(out)
+
+
+def _pointwise_add_value_4d_fp16(
+    npy_half[:, :, :, :] nums,
+    npy_half value,
+    npy_half[:, :, :, :] out,
+    Py_ssize_t num_threads,
+) -> np.ndarray:
+    cdef Py_ssize_t dim1 = out.shape[0], dim2 = out.shape[1], dim3 = out.shape[2], dim4 = out.shape[3]
+    cdef Py_ssize_t i1, i2, i3, i4
+
+    for i1 in prange(dim1, nogil=True, num_threads=num_threads):
+        for i2 in prange(dim2):
+            for i3 in prange(dim3):
+                for i4 in prange(dim4):
+                    out[i1, i2, i3, i4] = (npy_float_to_half(npy_half_to_float(nums[i1, i2, i3, i4]) +
+                                           npy_half_to_float(value)))
+
+    return np.asarray(out)
+
+
+def _fill_3d(NUM_AND_NPY_HALF[:, :, :] nums, NUM_AND_NPY_HALF value, Py_ssize_t num_threads) -> None:
+    cdef Py_ssize_t rows = nums.shape[0], cols = nums.shape[1], dims = nums.shape[2]
+    cdef Py_ssize_t i, j, k
+
+    for i in prange(rows, nogil=True, num_threads=num_threads):
+        for j in prange(cols):
+            for k in prange(dims):
+                nums[i, j, k] = value
+
+
+def _fill_4d(NUM_AND_NPY_HALF[:, :, :, :] nums, NUM_AND_NPY_HALF value, Py_ssize_t num_threads) -> None:
+    cdef Py_ssize_t dim1 = nums.shape[0], dim2 = nums.shape[1], dim3 = nums.shape[2], dim4 = nums.shape[3]
+    cdef Py_ssize_t i1, i2, i3, i4
+
+    for i1 in prange(dim1, nogil=True, num_threads=num_threads):
+        for i2 in prange(dim2):
+            for i3 in prange(dim3):
+                for i4 in prange(dim4):
+                    nums[i1, i2, i3, i4] = value
+
+
+# FIXME: somehow `const NUM_AND_NPY_HALF` is not working
+cpdef void _copy_3d(const NUM[:, :, :] nums1, NUM[:, :, :] nums2, Py_ssize_t num_threads):
+    cdef Py_ssize_t rows = nums1.shape[0], cols = nums1.shape[1], dims = nums1.shape[2]
+    cdef Py_ssize_t i, j, k
+
+    for i in prange(rows, nogil=True, num_threads=num_threads):
+        for j in prange(cols):
+            for k in prange(dims):
+                nums2[i, j, k] = nums1[i, j, k]
+
+
+cpdef void _copy_4d(const NUM[:, :, :, :] nums1, NUM[:, :, :, :] nums2, Py_ssize_t num_threads):
+    cdef Py_ssize_t dim1 = nums1.shape[0], dim2 = nums1.shape[1], dim3 = nums1.shape[2], dim4 = nums1.shape[3]
+    cdef Py_ssize_t i1, i2, i3, i4
+
+    for i1 in prange(dim1, nogil=True, num_threads=num_threads):
+        for i2 in prange(dim2):
+            for i3 in prange(dim3):
+                for i4 in prange(dim4):
+                    nums2[i1, i2, i3, i4] = nums1[i1, i2, i3, i4]
+
+
+cpdef void _copy_3d_fp16(const npy_half[:, :, :] nums1, npy_half[:, :, :] nums2, Py_ssize_t num_threads):
+    cdef Py_ssize_t rows = nums1.shape[0], cols = nums1.shape[1], dims = nums1.shape[2]
+    cdef Py_ssize_t i, j, k
+
+    for i in prange(rows, nogil=True, num_threads=num_threads):
+        for j in prange(cols):
+            for k in prange(dims):
+                nums2[i, j, k] = nums1[i, j, k]
+
+
+cpdef void _copy_4d_fp16(const npy_half[:, :, :, :] nums1, npy_half[:, :, :, :] nums2, Py_ssize_t num_threads):
+    cdef Py_ssize_t dim1 = nums1.shape[0], dim2 = nums1.shape[1], dim3 = nums1.shape[2], dim4 = nums1.shape[3]
+    cdef Py_ssize_t i1, i2, i3, i4
+
+    for i1 in prange(dim1, nogil=True, num_threads=num_threads):
+        for i2 in prange(dim2):
+            for i3 in prange(dim3):
+                for i4 in prange(dim4):
+                    nums2[i1, i2, i3, i4] = nums1[i1, i2, i3, i4]