Merge pull request #2860 from noamher/conditional-gilstatnode

Allow condition in GILStatNode

8 files changed, 456 insertions, 2 deletions

diff --git a/Cython/Compiler/Nodes.py b/Cython/Compiler/Nodes.py
index 42ef109b3..89ef8a192 100644
--- a/Cython/Compiler/Nodes.py
+++ b/Cython/Compiler/Nodes.py
@@ -7805,10 +7805,12 @@ class GILStatNode(NogilTryFinallyStatNode):
     #
     #   state   string   'gil' or 'nogil'
 
+    child_attrs = ["condition"] + NogilTryFinallyStatNode.child_attrs
     state_temp = None
 
-    def __init__(self, pos, state, body):
+    def __init__(self, pos, state, body, condition=None):
         self.state = state
+        self.condition = condition
         self.create_state_temp_if_needed(pos, state, body)
         TryFinallyStatNode.__init__(
             self, pos,
@@ -7835,11 +7837,18 @@ class GILStatNode(NogilTryFinallyStatNode):
         if self.state == 'gil':
             env.has_with_gil_block = True
 
+        if self.condition is not None:
+            self.condition.analyse_declarations(env)
+
         return super(GILStatNode, self).analyse_declarations(env)
 
     def analyse_expressions(self, env):
         env.use_utility_code(
             UtilityCode.load_cached("ForceInitThreads", "ModuleSetupCode.c"))
+
+        if self.condition is not None:
+            self.condition = self.condition.analyse_expressions(env)
+
         was_nogil = env.nogil
         env.nogil = self.state == 'nogil'
         node = TryFinallyStatNode.analyse_expressions(self, env)
diff --git a/Cython/Compiler/Optimize.py b/Cython/Compiler/Optimize.py
index 561c56152..d48f98155 100644
--- a/Cython/Compiler/Optimize.py
+++ b/Cython/Compiler/Optimize.py
@@ -4688,6 +4688,30 @@ class ConstantFolding(Visitor.VisitorTransform, SkipDeclarations):
             return None
         return node
 
+    def visit_GILStatNode(self, node):
+        self.visitchildren(node)
+        if node.condition is None:
+            return node
+
+        if node.condition.has_constant_result():
+            # Condition is True - Modify node to be a normal
+            # GILStatNode with condition=None
+            if node.condition.constant_result:
+                node.condition = None
+
+            # Condition is False - the body of the GILStatNode
+            # should run without changing the state of the gil
+            # return the body of the GILStatNode
+            else:
+                return node.body
+
+        # If condition is not constant we keep the GILStatNode as it is.
+        # Either it will later become constant (e.g. a `numeric is int`
+        # expression in a fused type function) and then when ConstantFolding
+        # runs again it will be handled or a later transform (i.e. GilCheck)
+        # will raise an error
+        return node
+
     # in the future, other nodes can have their own handler method here
     # that can replace them with a constant result node
 
diff --git a/Cython/Compiler/ParseTreeTransforms.py b/Cython/Compiler/ParseTreeTransforms.py
index c3e3ff8ad..189804ed0 100644
--- a/Cython/Compiler/ParseTreeTransforms.py
+++ b/Cython/Compiler/ParseTreeTransforms.py
@@ -2915,6 +2915,12 @@ class GilCheck(VisitorTransform):
         return node
 
     def visit_GILStatNode(self, node):
+        if node.condition is not None:
+            error(node.condition.pos,
+                  "Non-constant condition in a "
+                  "`with %s(<condition>)` statement" % node.state)
+            return node
+
         if self.nogil and node.nogil_check:
             node.nogil_check()
 
@@ -3254,6 +3260,13 @@ class ReplaceFusedTypeChecks(VisitorTransform):
         self.visitchildren(node)
         return self.transform(node)
 
+    def visit_GILStatNode(self, node):
+        """
+        Fold constant condition of GILStatNode.
+        """
+        self.visitchildren(node)
+        return self.transform(node)
+
     def visit_PrimaryCmpNode(self, node):
         with Errors.local_errors(ignore=True):
           type1 = node.operand1.analyse_as_type(self.local_scope)
diff --git a/Cython/Compiler/Parsing.py b/Cython/Compiler/Parsing.py
index 4e5d443eb..8ca9beebc 100644
--- a/Cython/Compiler/Parsing.py
+++ b/Cython/Compiler/Parsing.py
@@ -2051,12 +2051,20 @@ def p_with_items(s, is_async=False):
             s.error("with gil/nogil cannot be async")
         state = s.systring
         s.next()
+
+        # support conditional gil/nogil
+        condition = None
+        if s.sy == '(':
+            s.next()
+            condition = p_test(s)
+            s.expect(')')
+
         if s.sy == ',':
             s.next()
             body = p_with_items(s)
         else:
             body = p_suite(s)
-        return Nodes.GILStatNode(pos, state=state, body=body)
+        return Nodes.GILStatNode(pos, state=state, body=body, condition=condition)
     else:
         manager = p_test(s)
         target = None
diff --git a/docs/src/userguide/external_C_code.rst b/docs/src/userguide/external_C_code.rst
index b3e0546ea..9e55d0c22 100644
--- a/docs/src/userguide/external_C_code.rst
+++ b/docs/src/userguide/external_C_code.rst
@@ -580,6 +580,26 @@ The GIL may also be acquired through the ``with gil`` statement::
     with gil:
         <execute this block with the GIL acquired>
 
+.. _gil_conditional:
+
+Conditional Acquiring / Releasing the GIL
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Sometimes it is helpful to use a condition to decide whether to run a
+certain piece of code with or without the GIL. This code would run anyway,
+the difference is whether the GIL will be held or released.
+The condition must be constant (at compile time).
+
+This could be useful for profiling, debugging, performance testing, and
+for fused types (see :ref:`fused_gil_conditional`).::
+
+    DEF FREE_GIL = True
+
+    with nogil(FREE_GIL):
+        <code to be executed with the GIL released>
+
+        with gil(False):
+           <GIL is still released>
+
 Declaring a function as callable without the GIL
 --------------------------------------------------
 
diff --git a/docs/src/userguide/fusedtypes.rst b/docs/src/userguide/fusedtypes.rst
index b50bb0efd..d91b3ef91 100644
--- a/docs/src/userguide/fusedtypes.rst
+++ b/docs/src/userguide/fusedtypes.rst
@@ -249,6 +249,34 @@ to figure out whether a specialization is part of another set of types
         if bunch_of_types in string_t:
             print("s is a string!")
 
+.. _fused_gil_conditional:
+
+Conditional GIL Acquiring / Releasing
+=====================================
+
+Acquiring and releasing the GIL can be controlled by a condition
+which is known at compile time (see :ref:`_gil_conditional`).
+
+This is most useful when combined with fused types.
+A fused type function may have to handle both cython native types
+(e.g. cython.int or cython.double) and python types (e.g. object or bytes).
+Conditional Acquiring / Releasing the GIL provides a method for running
+the same piece of code either with the GIL released (for cython native types)
+and with the GIL held (for python types).::
+
+    cimport cython
+
+    ctypedef fused double_or_object:
+        cython.double
+        object
+
+    def increment(double_or_object x):
+        with nogil(double_or_object is cython.double):
+            # Same code handles both cython.double (GIL is released)
+            # and python object (GIL is not released).
+            x = x + 1
+        return x
+
 __signatures__
 ==============
 
diff --git a/tests/errors/nogil_conditional.pyx b/tests/errors/nogil_conditional.pyx
new file mode 100644
index 000000000..e800eb199
--- /dev/null
+++ b/tests/errors/nogil_conditional.pyx
@@ -0,0 +1,81 @@
+# cython: remove_unreachable=False
+# mode: error
+
+cdef int f_nogil(int x) nogil:
+    cdef int y
+    y = x + 10
+    return y
+
+
+def f_gil(x):
+    y = 0
+    y = x + 100
+    return y
+
+
+def illegal_gil_usage():
+    cdef int res = 0
+    with nogil(True):
+        res = f_gil(res)
+
+        with nogil(True):
+            res = f_gil(res)
+
+        with gil(False):
+            res = f_gil(res)
+
+    with nogil(False):
+        res = f_nogil(res)
+
+
+def foo(a):
+    return a < 10
+
+
+def non_constant_condition(int x) -> int:
+    cdef int res = x
+    with nogil(x < 10):
+        res = f_nogil(res)
+
+    with gil(foo(x)):
+         res = f_gil(res)
+
+
+ctypedef fused number_or_object:
+    float
+    object
+
+
+def fused_type(number_or_object x):
+    with nogil(number_or_object is object):
+        res = x + 1
+
+    # This should be fine
+    with nogil(number_or_object is float):
+        res = x + 1
+
+    return res
+
+
+_ERRORS = u"""
+19:14: Accessing Python global or builtin not allowed without gil
+19:19: Calling gil-requiring function not allowed without gil
+19:19: Coercion from Python not allowed without the GIL
+19:19: Constructing Python tuple not allowed without gil
+19:20: Converting to Python object not allowed without gil
+21:13: Trying to release the GIL while it was previously released.
+22:18: Accessing Python global or builtin not allowed without gil
+22:23: Calling gil-requiring function not allowed without gil
+22:23: Coercion from Python not allowed without the GIL
+22:23: Constructing Python tuple not allowed without gil
+22:24: Converting to Python object not allowed without gil
+25:18: Accessing Python global or builtin not allowed without gil
+25:23: Calling gil-requiring function not allowed without gil
+25:23: Coercion from Python not allowed without the GIL
+25:23: Constructing Python tuple not allowed without gil
+25:24: Converting to Python object not allowed without gil
+37:17: Non-constant condition in a `with nogil(<condition>)` statement
+40:16: Non-constant condition in a `with gil(<condition>)` statement
+51:8: Assignment of Python object not allowed without gil
+51:16: Calling gil-requiring function not allowed without gil
+"""
diff --git a/tests/run/nogil_conditional.pyx b/tests/run/nogil_conditional.pyx
new file mode 100644
index 000000000..eba22d5b2
--- /dev/null
+++ b/tests/run/nogil_conditional.pyx
@@ -0,0 +1,271 @@
+# mode: run
+
+try:
+    from StringIO import StringIO
+except ImportError:
+    from io import StringIO
+
+
+def test(int x):
+    """
+    >>> test(0)
+    110
+    """
+    with nogil(True):
+        x = f_nogil(x)
+        with gil(True):
+            x = f_gil(x)
+    return x
+
+
+cdef int f_nogil(int x) nogil:
+    cdef int y
+    y = x + 10
+    return y
+
+
+def f_gil(x):
+    y = 0
+    y = x + 100
+    return y
+
+
+cdef int with_gil_func() except? -1 with gil:
+    raise Exception("error!")
+
+
+cdef int nogil_func() nogil except? -1:
+    with_gil_func()
+
+
+def test_nogil_exception_propagation():
+    """
+    >>> test_nogil_exception_propagation()
+    Traceback (most recent call last):
+       ...
+    Exception: error!
+    """
+    with nogil:
+        with gil:
+            with nogil(True):
+                nogil_func()
+
+
+cdef int write_unraisable() nogil:
+    with gil:
+        raise ValueError()
+
+
+def test_unraisable():
+    """
+    >>> print(test_unraisable())  # doctest: +ELLIPSIS
+    ValueError
+    Exception...ignored...
+    """
+    import sys
+    old_stderr = sys.stderr
+    stderr = sys.stderr = StringIO()
+    try:
+        write_unraisable()
+    finally:
+        sys.stderr = old_stderr
+    return stderr.getvalue().strip()
+
+
+def test_nested():
+    """
+    >>> test_nested()
+    240
+    """
+    cdef int res = 0
+
+    with nogil(True):
+        res = f_nogil(res)
+        with gil(1 < 2):
+            res = f_gil(res)
+            with nogil:
+                res = f_nogil(res)
+
+        with gil:
+            res = f_gil(res)
+            with nogil(True):
+                res = f_nogil(res)
+            with nogil:
+                res = f_nogil(res)
+
+    return res
+
+
+DEF FREE_GIL = True
+DEF FREE_GIL_FALSE = False
+
+
+def test_nested_condition_false():
+    """
+    >>> test_nested_condition_false()
+    220
+    """
+    cdef int res = 0
+
+    with gil(FREE_GIL_FALSE):
+        res = f_gil(res)
+        with nogil(False):
+            res = f_gil(res)
+
+        with nogil(FREE_GIL):
+            res = f_nogil(res)
+            with gil(False):
+                res = f_nogil(res)
+
+    return res
+
+def test_try_finally():
+    """
+    >>> test_try_finally()
+    113
+    """
+    cdef int res = 0
+
+    try:
+        with nogil(True):
+            try:
+                res = f_nogil(res)
+                with gil(1 < 2):
+                    try:
+                        res = f_gil(res)
+                    finally:
+                        res += 1
+            finally:
+                res = res + 1
+    finally:
+        res += 1
+
+    return res
+
+
+ctypedef fused number_or_object:
+    int
+    float
+    object
+
+
+def test_fused(number_or_object x) -> number_or_object:
+    """
+    >>> test_fused[int](1)
+    2
+    >>> test_fused[float](1.0)
+    2.0
+    >>> test_fused[object](1)
+    2
+    >>> test_fused[object](1.0)
+    2.0
+    """
+    cdef number_or_object res = x
+
+    with nogil(number_or_object is not object):
+        res = res + 1
+
+    return res
+
+
+ctypedef fused int_or_object:
+    int
+    object
+
+
+def test_fused_object(int_or_object x):
+    """
+    >>> test_fused_object[object]("spam")
+    456
+    >>> test_fused_object[int](1000)
+    1000
+    """
+    cdef int res = 0
+
+    if int_or_object is object:
+        with nogil(False):
+            res += len(x)
+
+        try:
+            with nogil(int_or_object is object):
+                try:
+                    with gil(int_or_object is object):
+                        res = f_gil(res)
+                    with gil:
+                        res = f_gil(res)
+                    with gil(False):
+                        res = f_nogil(res)
+
+                    with gil(int_or_object is not object):
+                        res = f_nogil(res)
+                    with nogil(False):
+                        res = f_nogil(res)
+
+                    res = f_nogil(res)
+                finally:
+                    res = res + 1
+
+            with nogil(int_or_object is not object):
+                res = f_gil(res)
+
+            with gil(int_or_object is not object):
+                res = f_gil(res)
+
+                with nogil(int_or_object is object):
+                    res = f_nogil(res)
+
+        finally:
+            res += 1
+    else:
+        res = x
+
+    return res
+
+
+def test_fused_int(int_or_object x):
+    """
+    >>> test_fused_int[object]("spam")
+    4
+    >>> test_fused_int[int](1000)
+    1452
+    """
+    cdef int res = 0
+
+    if int_or_object is int:
+        res += x
+
+        try:
+            with nogil(int_or_object is int):
+                try:
+                    with gil(int_or_object is int):
+                        res = f_gil(res)
+                    with gil:
+                        res = f_gil(res)
+                    with gil(False):
+                        res = f_nogil(res)
+
+                    with gil(int_or_object is not int):
+                        res = f_nogil(res)
+                    with nogil(False):
+                        res = f_nogil(res)
+
+                    res = f_nogil(res)
+                finally:
+                    res = res + 1
+
+            with nogil(int_or_object is not int):
+                res = f_gil(res)
+
+            with gil(int_or_object is not int):
+                res = f_gil(res)
+
+                with nogil(int_or_object is int):
+                    res = f_nogil(res)
+
+        finally:
+            res += 1
+    else:
+        with nogil(False):
+            res = len(x)
+
+    return res