Add atomic ops

TODO

@@ -1,2 +0,0 @@
-* add atomic operation to llvm-py
-* add volatile load/store to llvm-py

llvm_cbuilder/builder.py

@@ -4,6 +4,7 @@
 
 import contextlib
 import llvm.core as lc
+import llvm.ee as le
 
 
 def _is_int(ty):
@@ -115,6 +116,7 @@ class CBuilder(object):
         self.declare_block = self.function.append_basic_block('decl')
         self.first_body_block = self.function.append_basic_block('body')
         self.builder = lc.Builder.new(self.first_body_block)
+        self.target_data = le.TargetData.new(self.function.module.data_layout)
 
         # prepare arguments
         self.args = []
@@ -210,6 +212,81 @@ class CBuilder(object):
         '''
         return _is_block_terminated(self.builder.basic_block)
 
+    def atomic_cmpxchg(self, ptr, old, val, ordering, crossthread=True):
+        res = self.builder.atomic_cmpxchg(ptr.value, old.value, val.value,
+                                          ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_xchg(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_xchg(ptr.value, val.value,
+                                       ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_add(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_add(ptr.value, val.value,
+                                      ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_sub(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_sub(ptr.value, val.value,
+                                      ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_and(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_and(ptr.value, val.value,
+                                      ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_nand(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_nand(ptr.value, val.value,
+                                       ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_or(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_or(ptr.value, val.value,
+                                     ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_xor(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_xor(ptr.value, val.value,
+                                      ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_max(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_max(ptr.value, val.value,
+                                      ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_min(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_min(ptr.value, val.value,
+                                      ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_umax(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_umax(ptr.value, val.value,
+                                       ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_umin(self, ptr, val, ordering, crossthread=True):
+        res = self.builder.atomic_umin(ptr.value, val.value,
+                                       ordering, crossthread)
+        return CTemp(self, res)
+
+    def atomic_load(self, ptr, ordering, align=1, crossthread=True):
+        res = self.builder.atomic_load(ptr.value, ordering, align, crossthread)
+        return CTemp(self, res)
+
+    def atomic_store(self, val, ptr, ordering, align=1, crossthread=True):
+        res = self.builder.atomic_store(val.value, ptr.value, ordering,
+                                        align, crossthread)
+        return CTemp(self, res)
+
+    def fence(self, ordering, crossthread=True):
+        res = self.builder.fence(ordering, crossthread)
+        return CTemp(self, res)
+
+    def alignment(self, ty):
+        return self.target_data.abi_alignment(ty)
 
 class CValue(object):
     '''
@@ -419,7 +496,7 @@ class CFunc(CValue):
         self.function = func
 
     def __call__(self, *args):
-        arg_value = map(lambda x: x.value, args)
+        arg_value = list(map(lambda x: x.value, args))
         res = self.parent.builder.call(self.function, arg_value)
         return CTemp(self.parent, res)
 
@@ -441,6 +518,7 @@ class CTemp(CValue):
         return self.value.type
 
 class CVar(CValue):
+
     def __init__(self, parent, ptr):
         super(CVar, self).__init__(parent)
         self.ptr = ptr
@@ -478,15 +556,39 @@ class CVar(CValue):
     def type(self):
         return self.ptr.type.pointee
 
-    def load(self):
+    def load(self, volatile=False):
         self._ensure_is_pointer()
-        loaded = self.parent.builder.load(self.value)
+        loaded = self.parent.builder.load(self.value, volatile=volatile)
         return CTemp(self.parent, loaded)
 
     def store(self, val):
         self._ensure_is_pointer()
         self.parent.builder.store(val.value, self.value)
 
+    def atomic_load(self, ordering, align=None, crossthread=True, volatile=False):
+        self._ensure_is_pointer()
+        if align is None:
+            align = self.parent.alignment(self.type.pointee)
+        inst = self.parent.builder.atomic_load(self.value, ordering, align,
+                                               crossthread=crossthread,
+                                               volatile=volatile)
+        return CTemp(self.parent, inst)
+
+    def atomic_store(self, value, ordering, align=None,  crossthread=True,
+                     volatile=False):
+        self._ensure_is_pointer()
+        if align is None:
+            align = self.parent.alignment(self.type.pointee)
+        self.parent.builder.atomic_store(value.value, self.value, ordering,
+                                         align=align, crossthread=crossthread)
+
+    def atomic_cmpxchg(self, old, new, ordering, crossthread=True, volatile=False):
+        self._ensure_is_pointer()
+        inst = self.parent.builder.atomic_cmpxchg(self.value, old.value,
+                                                  new.value, ordering,
+                                                  crossthread=crossthread)
+        return CTemp(self.parent, inst)
+
     def reference(self):
         return CTemp(self.parent, self.ptr)
 

llvm_cbuilder/executor.py

@@ -31,7 +31,7 @@ MAP_CTYPES = {
 class CExecutor(object):
     def __init__(self, mod_or_engine):
         if isinstance(mod_or_engine, Module):
-            self.engine = le.ExecutionEngine.new(mod_or_engine)
+            self.engine = le.EngineBuilder.new(mod_or_engine).create()
         else:
             self.engine = mod_or_engine
 

tests/test_atomic_add.py

@@ -0,0 +1,115 @@
+'''
+Base on the test_pthread.py and extend to use atomic instructions
+'''
+
+from llvm.core import *
+from llvm.passes import *
+from llvm.ee import *
+from llvm_cbuilder import *
+import llvm_cbuilder.shortnames as C
+import unittest, logging
+
+logging.basicConfig(level=logging.DEBUG)
+
+NUM_OF_THREAD = 4
+REPEAT = 10000
+
+def gen_test_worker(mod):
+    cb = CBuilder.new_function(mod, 'worker', C.void, [C.pointer(C.int)])
+    pval = cb.args[0]
+    one = cb.constant(pval.type.pointee, 1)
+
+    ct = cb.var(C.int, 0)
+    limit = cb.constant(C.int, REPEAT)
+    with cb.loop() as loop:
+        with loop.condition() as setcond:
+            setcond( ct < limit )
+
+        with loop.body():
+            cb.atomic_add(pval, one, 'acq_rel')
+            ct += one
+
+    cb.ret()
+    cb.close()
+    return cb.function
+
+def gen_test_pthread(mod):
+    cb = CBuilder.new_function(mod, 'manager', C.int, [C.int])
+    arg = cb.args[0]
+
+    worker_func = cb.get_function_named('worker')
+    pthread_create = cb.get_function_named('pthread_create')
+    pthread_join = cb.get_function_named('pthread_join')
+
+
+    NULL = cb.constant_null(C.void_p)
+    cast_to_null = lambda x: x.cast(C.void_p)
+
+    threads = cb.array(C.void_p, NUM_OF_THREAD)
+
+    for tid in range(NUM_OF_THREAD):
+        pthread_create_args = [threads[tid].reference(),
+                               NULL,
+                               worker_func,
+                               arg.reference()]
+        pthread_create(*map(cast_to_null, pthread_create_args))
+
+    worker_func(arg.reference())
+
+    for tid in range(NUM_OF_THREAD):
+        pthread_join_args = threads[tid], NULL
+        pthread_join(*map(cast_to_null, pthread_join_args))
+
+
+    cb.ret(arg)
+    cb.close()
+    return cb.function
+
+class TestPThread(unittest.TestCase):
+    def test_pthread(self):
+        mod = Module.new(__name__)
+        # add pthread functions
+
+        mod.add_function(Type.function(C.int,
+                                       [C.void_p, C.void_p, C.void_p, C.void_p]),
+                         'pthread_create')
+
+        mod.add_function(Type.function(C.int,
+                                       [C.void_p, C.void_p]),
+                         'pthread_join')
+
+        lf_test_worker = gen_test_worker(mod)
+        lf_test_pthread = gen_test_pthread(mod)
+        logging.debug(mod)
+        mod.verify()
+
+        # optimize
+        fpm = FunctionPassManager.new(mod)
+        mpm = PassManager.new()
+        pmb = PassManagerBuilder.new()
+        pmb.vectorize = True
+        pmb.opt_level = 3
+        pmb.populate(fpm)
+        pmb.populate(mpm)
+
+        fpm.run(lf_test_worker)
+        fpm.run(lf_test_pthread)
+        mpm.run(mod)
+        logging.debug(mod)
+        mod.verify()
+
+        # run
+        exe = CExecutor(mod)
+        exe.engine.get_pointer_to_function(mod.get_function_named('worker'))
+        func = exe.get_ctype_function(lf_test_pthread, 'int, int')
+
+        inarg = 1234
+        gold = inarg + (NUM_OF_THREAD + 1) * REPEAT
+
+        for _ in range(1000): # run many many times to catch race condition
+            self.assertEqual(func(inarg), gold, "Unexpected race condition")
+
+
+if __name__ == '__main__':
+    unittest.main()
+

tests/test_atomic_ldst.py

@@ -0,0 +1,122 @@
+'''
+Base on the test_pthread.py and extend to use atomic instructions
+'''
+
+from llvm.core import *
+from llvm.passes import *
+from llvm.ee import *
+from llvm_cbuilder import *
+import llvm_cbuilder.shortnames as C
+import unittest, logging
+
+logging.basicConfig(level=logging.DEBUG)
+
+NUM_OF_THREAD = 4
+REPEAT = 10000
+
+def gen_test_worker(mod):
+    cb = CBuilder.new_function(mod, 'worker', C.void, [C.pointer(C.int)])
+    pval = cb.args[0]
+    one = cb.constant(pval.type.pointee, 1)
+
+    ct = cb.var(C.int, 0)
+    limit = cb.constant(C.int, REPEAT)
+    with cb.loop() as loop:
+        with loop.condition() as setcond:
+            setcond( ct < limit )
+
+        with loop.body():
+            oldval = pval.atomic_load('acquire')
+            updated = oldval + one
+            castmp = pval.atomic_cmpxchg(oldval, updated, 'release')
+
+            with cb.ifelse( castmp == oldval ) as ifelse:
+                with ifelse.then():
+                    ct += one
+
+    cb.ret()
+    cb.close()
+    return cb.function
+
+def gen_test_pthread(mod):
+    cb = CBuilder.new_function(mod, 'manager', C.int, [C.int])
+    arg = cb.args[0]
+
+    worker_func = cb.get_function_named('worker')
+    pthread_create = cb.get_function_named('pthread_create')
+    pthread_join = cb.get_function_named('pthread_join')
+
+
+    NULL = cb.constant_null(C.void_p)
+    cast_to_null = lambda x: x.cast(C.void_p)
+
+    threads = cb.array(C.void_p, NUM_OF_THREAD)
+
+    for tid in range(NUM_OF_THREAD):
+        pthread_create_args = [threads[tid].reference(),
+                               NULL,
+                               worker_func,
+                               arg.reference()]
+        pthread_create(*map(cast_to_null, pthread_create_args))
+
+    worker_func(arg.reference())
+
+    for tid in range(NUM_OF_THREAD):
+        pthread_join_args = threads[tid], NULL
+        pthread_join(*map(cast_to_null, pthread_join_args))
+
+
+    cb.ret(arg)
+    cb.close()
+    return cb.function
+
+class TestPThread(unittest.TestCase):
+    def test_pthread(self):
+        mod = Module.new(__name__)
+        # add pthread functions
+
+        mod.add_function(Type.function(C.int,
+                                       [C.void_p, C.void_p, C.void_p, C.void_p]),
+                         'pthread_create')
+
+        mod.add_function(Type.function(C.int,
+                                       [C.void_p, C.void_p]),
+                         'pthread_join')
+
+        lf_test_worker = gen_test_worker(mod)
+        lf_test_pthread = gen_test_pthread(mod)
+        logging.debug(mod)
+        mod.verify()
+
+        # optimize
+        fpm = FunctionPassManager.new(mod)
+        mpm = PassManager.new()
+        pmb = PassManagerBuilder.new()
+        pmb.vectorize = True
+        pmb.opt_level = 3
+        pmb.populate(fpm)
+        pmb.populate(mpm)
+
+        fpm.run(lf_test_worker)
+        fpm.run(lf_test_pthread)
+        mpm.run(mod)
+        logging.debug(mod)
+        mod.verify()
+
+        # run
+        exe = CExecutor(mod)
+        exe.engine.get_pointer_to_function(mod.get_function_named('worker'))
+        func = exe.get_ctype_function(lf_test_pthread, 'int, int')
+
+        inarg = 1234
+        gold = inarg + (NUM_OF_THREAD + 1) * REPEAT
+
+        for _ in range(1000): # run many many times to catch race condition
+            res = func(inarg)
+            self.assertEqual(res, gold,
+                             "Unexpected race condition: res = %d" % res)
+
+
+if __name__ == '__main__':
+    unittest.main()
+

tests/test_isprime.py

@@ -73,8 +73,12 @@ def gen_is_prime_fast(mod):
             cb.ret(false)
 
     idx = cb.var(C.int, 3, name='idx')
+
     sqrt = cb.get_intrinsic(INTR_SQRT, [C.float])
+
     looplimit = one + sqrt(arg.cast(C.float)).cast(C.int)
+
+
     with cb.loop() as loop:
         with loop.condition() as setcond:
             setcond( idx < looplimit )
@@ -86,6 +90,7 @@ def gen_is_prime_fast(mod):
             # increment
             idx += two
 
+
     cb.ret(true)
     cb.close()
     return func

tests/test_pthread.py

@@ -39,12 +39,12 @@ def gen_test_pthread(mod):
                                arg.reference()]
         pthread_create(*map(cast_to_null, pthread_create_args))
 
+    worker_func(arg.reference())
+
     for tid in range(NUM_OF_THREAD):
         pthread_join_args = threads[tid], NULL
         pthread_join(*map(cast_to_null, pthread_join_args))
 
-    worker_func(arg.reference())
-
     cb.ret(arg)
     cb.close()
     return cb.function
@@ -72,7 +72,18 @@ class TestPThread(unittest.TestCase):
         func = exe.get_ctype_function(lf_test_pthread, 'int, int')
 
         inarg = 1234
-        self.assertEqual(func(inarg), inarg+NUM_OF_THREAD+1)
+        gold = inarg + NUM_OF_THREAD + 1
+        self.assertLessEqual(func(inarg), gold)
+        # Cannot determine the exact return value due to untamed race condition
+
+        count_race = 0
+        for _ in range(2**12):
+            if func(inarg) != gold:
+                count_race += 1
+
+        if count_race > 0:
+            logging.info("Race condition occured %d times.", count_race)
+            logging.info("Race condition is expected.")
 
 if __name__ == '__main__':
     unittest.main()