[docs] Parallelization tutorial (cython#5184)

Parallelization tutorial to try to explain prange/parallel in a little more user-friendly way

Author: da-woods

docs/examples/tutorial/parallelization/manual_work.py

@@ -0,0 +1,24 @@
+# tag: openmp
+
+from cython.parallel import parallel
+from cython.cimports.openmp import omp_get_thread_num
+import cython
+
+@cython.cfunc
+@cython.nogil
+def long_running_task1() -> cython.void:
+    pass
+
+@cython.cfunc
+@cython.nogil
+def long_running_task2() -> cython.void:
+    pass
+
+def do_two_tasks():
+    thread_num: cython.int
+    with cython.nogil, parallel(num_threads=2):
+        thread_num = omp_get_thread_num()
+        if thread_num == 0:
+            long_running_task1()
+        elif thread_num == 1:
+            long_running_task2()

docs/examples/tutorial/parallelization/manual_work.pyx

@@ -0,0 +1,25 @@
+# tag: openmp
+
+from cython.parallel cimport parallel
+from openmp cimport omp_get_thread_num
+
+
+
+
+cdef void long_running_task1() nogil:
+    pass
+
+
+
+cdef void long_running_task2() nogil:
+    pass
+
+def do_two_tasks():
+    cdef int thread_num
+    with nogil, parallel(num_threads=2):
+        thread_num = omp_get_thread_num()
+        if thread_num == 0:
+            long_running_task1()
+        elif thread_num == 1:
+            long_running_task2()
+

docs/examples/tutorial/parallelization/median.py

@@ -0,0 +1,35 @@
+# distutils: language = c++
+
+from cython.parallel import parallel, prange
+from cython.cimports.libc.stdlib import malloc, free
+from cython.cimports.libcpp.algorithm import nth_element
+import cython
+from cython.operator import dereference
+
+import numpy as np
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def median_along_axis0(x: cython.double[:,:]):
+    out: cython.double[::1] = np.empty(x.shape[1])
+    i: cython.Py_ssize_t
+    j: cython.Py_ssize_t
+    scratch: cython.pointer(cython.double)
+    median_it: cython.pointer(cython.double)
+    with cython.nogil, parallel():
+        # allocate scratch space per loop
+        scratch = cython.cast(
+            cython.pointer(cython.double),
+            malloc(cython.sizeof(cython.double)*x.shape[0]))
+        try:
+            for i in prange(x.shape[1]):
+                # copy row into scratch space
+                for j in range(x.shape[0]):
+                    scratch[j] = x[j, i]
+                median_it = scratch + x.shape[0]//2
+                nth_element(scratch, median_it, scratch + x.shape[0])
+                # for the sake of a simple example, don't handle even lengths...
+                out[i] = dereference(median_it)
+        finally:
+            free(scratch)
+    return np.asarray(out)

docs/examples/tutorial/parallelization/median.pyx

@@ -0,0 +1,34 @@
+# distutils: language = c++
+
+from cython.parallel cimport parallel, prange
+from libcpp.vector cimport vector
+from libcpp.algorithm cimport nth_element
+cimport cython
+from cython.operator cimport dereference
+
+import numpy as np
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def median_along_axis0(const double[:,:] x):
+    cdef double[::1] out = np.empty(x.shape[1])
+    cdef Py_ssize_t i, j
+
+    cdef vector[double] *scratch
+    cdef vector[double].iterator median_it
+    with nogil, parallel():
+        # allocate scratch space per loop
+        scratch = new vector[double](x.shape[0])
+        try:
+            for i in prange(x.shape[1]):
+                # copy row into scratch space
+                for j in range(x.shape[0]):
+                    dereference(scratch)[j] = x[j, i]
+                median_it = scratch.begin() + scratch.size()//2
+                nth_element(scratch.begin(), median_it, scratch.end())
+                # for the sake of a simple example, don't handle even lengths...
+                out[i] = dereference(median_it)
+        finally:
+            del scratch
+    return np.asarray(out)
+

docs/examples/tutorial/parallelization/norm.py

@@ -0,0 +1,12 @@
+from cython.parallel import prange
+import cython
+from cython.cimports.libc.math import sqrt
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def l2norm(x: cython.double[:]):
+    total: cython.double = 0
+    i: cython.Py_ssize_t
+    for i in prange(x.shape[0], nogil=True):
+        total += x[i]*x[i]
+    return sqrt(total)

docs/examples/tutorial/parallelization/norm.pyx

@@ -0,0 +1,12 @@
+from cython.parallel cimport prange
+cimport cython
+from libc.math cimport sqrt
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def l2norm(double[:] x):
+    cdef double total = 0
+    cdef Py_ssize_t i
+    for i in prange(x.shape[0], nogil=True):
+        total += x[i]*x[i]
+    return sqrt(total)

docs/examples/tutorial/parallelization/normalize.py

@@ -0,0 +1,16 @@
+from cython.parallel import parallel, prange
+import cython
+from cython.cimports.libc.math import sqrt
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def normalize(x: cython.double[:]):
+    i: cython.Py_ssize_t
+    total: cython.double = 0
+    norm: cython.double
+    with cython.nogil, parallel():
+        for i in prange(x.shape[0]):
+            total += x[i]*x[i]
+        norm = sqrt(total)
+        for i in prange(x.shape[0]):
+            x[i] /= norm

docs/examples/tutorial/parallelization/normalize.pyx

@@ -0,0 +1,17 @@
+from cython.parallel cimport parallel, prange
+cimport cython
+from libc.math cimport sqrt
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def normalize(double[:] x):
+    cdef Py_ssize_t i
+    cdef double total = 0
+    cdef double norm
+    with nogil, parallel():
+        for i in prange(x.shape[0]):
+            total += x[i]*x[i]
+        norm = sqrt(total)
+        for i in prange(x.shape[0]):
+            x[i] /= norm
+

docs/examples/tutorial/parallelization/parallel_sin.py

@@ -0,0 +1,16 @@
+from cython.parallel import prange
+import cython
+from cython.cimports.libc.math import sin
+
+import numpy as np
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def do_sine(input: cython.double[:,:]):
+    output : cython.double[:,:] = np.empty_like(input)
+    i : cython.Py_ssize_t
+    j : cython.Py_ssize_t
+    for i in prange(input.shape[0], nogil=True):
+        for j in range(input.shape[1]):
+            output[i, j] = sin(input[i, j])
+    return np.asarray(output)

docs/examples/tutorial/parallelization/parallel_sin.pyx

@@ -0,0 +1,16 @@
+from cython.parallel cimport prange
+cimport cython
+from libc.math cimport sin
+
+import numpy as np
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def do_sine(double[:,:] input):
+    cdef double[:,:] output = np.empty_like(input)
+    cdef Py_ssize_t i, j
+
+    for i in prange(input.shape[0], nogil=True):
+        for j in range(input.shape[1]):
+            output[i, j] = sin(input[i, j])
+    return np.asarray(output)