Merge pull request #80 from harshbaldwa/master

added atomic_dec

compyle/cython_generator.py

@@ -475,16 +475,28 @@ def _handle_cast_statement(self, name, call):
         ctype = call_args[1].strip()[1:-1]
         return '%s = <%s> (%s)' % (name, ctype, expr)
 
-    def _handle_atomic_statement(self, name, call, is_serial):
+    def _handle_atomic_statement_inc(self, name, call, is_serial):
         # FIXME: This won't handle casting to pointers
         # using something like 'intp'
         call_arg = call[11:-1].strip()
         if self._config.use_openmp and not is_serial:
-            return['openmp.omp_set_lock(&cy_lock)', '%s = %s' % (name, call_arg),
+            return['openmp.omp_set_lock(&cy_lock)',
+                   '%s = %s' % (name, call_arg),
                    '%s += 1' % call_arg, 'openmp.omp_unset_lock(&cy_lock)']
         else:
             return ['%s = %s' % (name, call_arg), '%s += 1' % call_arg]
 
+    def _handle_atomic_statement_dec(self, name, call, is_serial):
+        # FIXME: This won't handle casting to pointers
+        # using something like 'intp'
+        call_arg = call[11:-1].strip()
+        if self._config.use_openmp and not is_serial:
+            return['openmp.omp_set_lock(&cy_lock)',
+                   '%s = %s' % (name, call_arg),
+                   '%s -= 1' % call_arg, 'openmp.omp_unset_lock(&cy_lock)']
+        else:
+            return ['%s = %s' % (name, call_arg), '%s -= 1' % call_arg]
+
     def _parse_function(self, obj, declarations=None, is_serial=False):
         c_code, py_code = self._get_method_wrapper(obj, indent=' ' * 4,
                                                    declarations=declarations,
@@ -531,7 +543,19 @@ def _process_body_line(self, line, is_serial=False):
                 name = words[0]
                 call = words[1]
                 indent = line[:line.index(name)]
-                stmts = self._handle_atomic_statement(name, call, is_serial)
+                stmts = self._handle_atomic_statement_inc(
+                    name, call, is_serial)
+                result = ''
+                for stmt in stmts:
+                    result += indent + stmt + '\n'
+                return '', result + '\n'
+            elif words[1].startswith('atomic_dec') and \
+                    not line.strip().startswith('#'):
+                name = words[0]
+                call = words[1]
+                indent = line[:line.index(name)]
+                stmts = self._handle_atomic_statement_dec(
+                    name, call, is_serial)
                 result = ''
                 for stmt in stmts:
                     result += indent + stmt + '\n'

compyle/jit.py

@@ -237,6 +237,8 @@ def visit_Call(self, node):
             return self.visit_cast(node)
         if node.func.id == 'atomic_inc':
             return self.visit(node.args[0])
+        if node.func.id == 'atomic_dec':
+            return self.visit(node.args[0])
         if node.func.id == 'address':
             return self.visit_address(node)
         if node.func.id in self.external_funcs:

compyle/low_level.py

@@ -321,6 +321,14 @@ def __call__(self, *args, **kw):
         pass
 
 
+class _atomic_dec(Extern):
+    def code(self, backend):
+        return ''
+
+    def __call__(self, *args, **kw):
+        pass
+
+
 class _cast(Extern):
     def code(self, backend):
         return ''
@@ -334,6 +342,7 @@ def __call__(self, x, type_str):
 nogil = _nogil()
 address = _address()
 atomic_inc = _atomic_inc()
+atomic_dec = _atomic_dec()
 cast = _cast()
 
 

compyle/tests/test_parallel.py

@@ -8,7 +8,7 @@
 from ..array import wrap, zeros
 from ..types import annotate, declare
 from ..parallel import Elementwise, Reduction, Scan
-from ..low_level import atomic_inc
+from ..low_level import atomic_inc, atomic_dec
 from .test_jit import g
 
 MY_CONST = 42
@@ -191,6 +191,21 @@ def test_atomic_inc_cuda(self):
         importorskip('pycuda')
         self._test_atomic_inc(backend='cuda')
 
+    def test_atomic_dec_cython(self):
+        self._test_atomic_dec(backend='cython')
+
+    def test_atomic_dec_cython_parallel(self):
+        with use_config(use_openmp=True):
+            self._test_atomic_dec(backend='cython')
+
+    def test_atomic_dec_opencl(self):
+        importorskip('pyopencl')
+        self._test_atomic_dec(backend='opencl')
+
+    def test_atomic_dec_cuda(self):
+        importorskip('pycuda')
+        self._test_atomic_dec(backend='cuda')
+
     def test_repeated_scans_with_different_settings(self):
         importorskip('pyopencl')
         with use_config(use_double=False):
@@ -520,6 +535,25 @@ def reduce_knl(i, x, result):
         # Then
         self.assertTrue(result[0] == 50000)
 
+    def _test_atomic_dec(self, backend):
+        # Given
+        a = np.random.randint(0, 100, 50000, dtype=np.int32)
+        result = zeros(1, dtype=np.int32, backend=backend)
+
+        a = wrap(a, backend=backend)
+
+        @annotate(gintp='x, result', i='int')
+        def reduce_knl(i, x, result):
+            b = declare('int')
+            b = atomic_dec(result[0])
+
+        # When
+        knl = Elementwise(reduce_knl, backend=backend)
+        knl(a, result)
+
+        # Then
+        self.assertTrue(result[0] == -50000)
+
 
 class TestParallelUtilsJIT(ParallelUtilsBase, unittest.TestCase):
     def setUp(self):
@@ -835,3 +869,22 @@ def reduce_knl(i, x, result):
 
         # Then
         self.assertTrue(result[0] == 50000)
+
+    def _test_atomic_dec(self, backend):
+        # Given
+        a = np.random.randint(0, 100, 50000, dtype=np.int32)
+        result = zeros(1, dtype=np.int32, backend=backend)
+
+        a = wrap(a, backend=backend)
+
+        @annotate
+        def reduce_knl(i, x, result):
+            b = declare('int')
+            b = atomic_dec(result[0])
+
+        # When
+        knl = Elementwise(reduce_knl, backend=backend)
+        knl(a, result)
+
+        # Then
+        self.assertTrue(result[0] == -50000)

compyle/tests/test_translator.py

@@ -1335,7 +1335,7 @@ def f(x=1.0):
     assert code.strip() == expect.strip()
 
 
-def test_atomic_works():
+def test_atomic_inc_works():
     # Given
     def f(x=1.0):
         return atomic_inc(x)
@@ -1368,6 +1368,41 @@ def f(x=1.0):
 
     assert code.strip() == expect.strip()
 
+
+def test_atomic_dec_works():
+    # Given
+    def f(x=1.0):
+        return atomic_dec(x)
+
+    # When
+    t = OpenCLConverter()
+    code = t.parse_function(f)
+
+    # Then
+    expect = dedent('''
+    WITHIN_KERNEL double f(double x)
+    {
+        return atomic_dec(&x);
+    }
+    ''')
+
+    assert code.strip() == expect.strip()
+
+    # When
+    t = CUDAConverter()
+    code = t.parse_function(f)
+
+    # Then
+    expect = dedent('''
+    WITHIN_KERNEL double f(double x)
+    {
+        return atomicAdd(&x, -1);
+    }
+    ''')
+
+    assert code.strip() == expect.strip()
+
+
 def test_cast_works():
     # Given
     def f(x=1.0):

compyle/translator.py

@@ -781,8 +781,11 @@ def _get_self_type(self):
     def render_atomic(self, func, arg):
         if func == 'atomic_inc':
             return 'atomic_inc(&%s)' % self.visit(arg)
+        elif func == 'atomic_dec':
+            return 'atomic_dec(&%s)' % self.visit(arg)
         else:
-            raise NotImplementedError("Only atomic_inc supported right now")
+            raise NotImplementedError(
+                "Only atomic_inc, atomic_dec supported right now")
 
 
 class CUDAConverter(OpenCLConverter):
@@ -831,5 +834,8 @@ def _get_local_declarations(self):
     def render_atomic(self, func, arg):
         if func == 'atomic_inc':
             return 'atomicAdd(&%s, 1)' % self.visit(arg)
+        elif func == 'atomic_dec':
+            return 'atomicAdd(&%s, -1)' % self.visit(arg)
         else:
-            raise NotImplementedError("Only atomic_inc supported right now")
+            raise NotImplementedError(
+                "Only atomic_inc, atomic_dec supported right now")