Added rudimentary support for complex numbers (no arithmetic yet).

Added complex types defined as LLVM structures in numba.llvm_types. Added ctypes support for the resulting LLVM structure in numba.utils. Lots of tweaks to the type-specific codes in numba.translate. Added numba.translate.MethodReference class to support method calls. Added code generation for calculating complex conjugate in _LLVMModuleUtils.build_conj(). Modified ctypes function wrapper generation to add additional layer of wrapper for functions that return a complex result. Modified numba.translate.Translate.op_LOAD_ATTR() to handle complex data. Added unit tests for complex numbers in new file .../tests/test_complex.py.
jriehl · Aug 24, 2012 · Jun 26, 2012 · Jun 26, 2012 · Jun 26, 2012 · Jun 26, 2012
commit b0fcd32d3ef0b3d894c6287a2057763ab227388e
diff --git a/numba/llvm_types.py b/numba/llvm_types.py
@@ -22,6 +22,10 @@
 _intp_star = lc.Type.pointer(_intp)
 _void_star = lc.Type.pointer(lc.Type.int(8))
 _void_star_star = lc.Type.pointer(_void_star)
+_float = lc.Type.float()
+_double = lc.Type.double()
+_complex64 = lc.Type.struct([_float, _float])
+_complex128 = lc.Type.struct([_double, _double])
 
 _pyobject_head = [_intp, lc.Type.pointer(_int32)]
 if hasattr(sys, 'getobjects'):

diff --git a/numba/translate.py b/numba/translate.py
@@ -10,12 +10,13 @@
 import llvm.passes as lp
 import llvm.ee as le
 
-from utils import itercode, debugout
+from utils import itercode, debugout, Complex64, Complex128
 from ._ext import make_ufunc
 from .cfg import ControlFlowGraph
 from .llvm_types import _plat_bits, _int1, _int8, _int32, _intp, _intp_star, \
-    _void_star, _pyobject_head, _trace_refs_, _head_len, _numpy_struct, \
-    _numpy_array, _numpy_array_field_ofs
+    _void_star, _float, _double, _complex64, _complex128, _pyobject_head, \
+    _trace_refs_, _head_len, _numpy_struct,  _numpy_array, \
+    _numpy_array_field_ofs
 from .multiarray_api import MultiarrayAPI
 
 if __debug__:
@@ -54,6 +55,10 @@ def llvmtype_to_strtype(typ):
                               # otherwise it won't ever be used.
         # FIXME: Need to find a way to stash dtype somewhere in here.
         return 'arr[]'
+    elif typ == _complex64:
+        return 'c64'
+    elif typ == _complex128:
+        return 'c128'
     elif typ.kind == lc.TYPE_POINTER:
         if typ.pointee.kind == lc.TYPE_FUNCTION:
             return ['func'] + typ.pointee.args
@@ -72,6 +77,8 @@ def pythontype_to_strtype(typ):
         return 'f64'
     elif issubclass(typ, int):
         return 'i%d' % _plat_bits
+    elif issubclass(typ, complex):
+        return 'c128'
     elif issubclass(typ, (types.BuiltinFunctionType, types.FunctionType)):
         return ["func"]
     elif issubclass(typ, tuple):
@@ -106,6 +113,11 @@ def get_inc(self):
     def get_stop(self):
         return self.args[0 if (len(self.args) == 1) else 1]
 
+class MethodReference(object):
+    def __init__(self, object_var, py_method):
+        self.object_var = object_var
+        self.py_method = py_method
+
 class _LLVMCaster(object):
     # NOTE: Using a class to lower namespace polution here.  The
     # following would be class methods, but we'd have to index them
@@ -218,14 +230,19 @@ def llvm(self, typ=None, mod=None, builder = None):
             return ret_val
         else:
             if typ is None:
-                typ = 'f64'
+                typ = 'f64' if self.typ is None else self.typ
             if typ == 'f64':
                 res = lc.Constant.real(lc.Type.double(), float(self.val))
             elif typ == 'f32':
                 res = lc.Constant.real(lc.Type.float(), float(self.val))
             elif typ[0] == 'i':
-                res = lc.Constant.int(lc.Type.int(int(typ[1:])), 
+                res = lc.Constant.int(lc.Type.int(int(typ[1:])),
                                       int(self.val))
+            elif typ[0] == 'c':
+                ltyp = _float if typ[1:] == '64' else _double
+                res = lc.Constant.struct([
+                        lc.Constant.real(ltyp, self.val.real),
+                        lc.Constant.real(ltyp, self.val.imag)])
             elif typ[0] == 'func':
                 res = map_to_function(self.val, typ[1:], mod)
             return res
@@ -252,6 +269,11 @@ def str_to_llvmtype(str):
     elif str[0] == 'i':
         num = int(str[1:])
         ret_val = lc.Type.int(num)
+    elif str[0] == 'c':
+        if str[1:] == '64':
+            ret_val = _complex64
+        elif str[1:] == '128':
+            ret_val = _complex128
     elif str.startswith('arr['):
         ret_val = _numpy_array
     else:
@@ -302,7 +324,17 @@ def convert_to_ctypes(typ):
         if __debug__:
             print("convert_to_ctypes(): n_pointer = %d, typ' = %r" %
                   (n_pointer, typ))
-    ret_val = _typecodes[np.dtype(typ).str]
+    dtype_str = np.dtype(typ).str
+    if dtype_str[1] == 'c':
+        if dtype_str[2:] == '8':
+            ret_val = Complex64
+        elif dtype_str[2:] == '16':
+            ret_val = Complex128
+        else:
+            raise NotImplementedError("Numba does not currently support "
+                                      "ctypes wrapping for type %r." % (typ,))
+    else:
+        ret_val = _typecodes[np.dtype(typ).str]
     for _ in xrange(n_pointer):
         ret_val = ctypes.POINTER(ret_val)
     return ret_val
@@ -554,6 +586,7 @@ def build_zeros_like(cls, translator, args):
                 args[0]._llvm is not None and
                 args[0]._llvm.type == _numpy_array), (
             "Expected Numpy array argument to numpy.zeros_like().")
+        translator.init_multiarray()
         larr = args[0]._llvm
         largs = [translator.builder.load(
                 translator.builder.gep(larr, [
@@ -606,10 +639,30 @@ def get_py_modulo(cls, module, typ):
             builder.ret(res)
         return ret_val
 
+    @classmethod
+    def build_conj(cls, translator, args):
+        print args
+        assert ((len(args) == 1) and (args[0]._llvm is not None) and
+                (args[0]._llvm.type in (_complex64, _complex128)))
+        larg = args[0]._llvm
+        elem_ltyp = _float if larg.type == _complex64 else _double
+        new_imag_lval = translator.builder.fmul(
+            lc.Constant.real(elem_ltyp, -1.),
+            translator.builder.extract_value(larg, 1))
+        assert hasattr(translator.builder, 'insert_value'), (
+            "llvm-py support for LLVMBuildInsertValue() required to build "
+            "code for complex conjugates.")
+        res = translator.builder.insert_value(larg, new_imag_lval, 1)
+        return res, None
+
 PY_CALL_TO_LLVM_CALL_MAP = {
     debugout : _LLVMModuleUtils.build_debugout,
     len : _LLVMModuleUtils.build_len,
     np.zeros_like : _LLVMModuleUtils.build_zeros_like,
+    np.complex64.conj : _LLVMModuleUtils.build_conj,
+    np.complex128.conj : _LLVMModuleUtils.build_conj,
+    np.complex64.conjugate : _LLVMModuleUtils.build_conj,
+    np.complex128.conjugate : _LLVMModuleUtils.build_conj,
 }
 
 class LLVMControlFlowGraph (ControlFlowGraph):
@@ -1016,9 +1069,14 @@ def check_locals(self, i):
     def get_ctypes_func(self, llvm=True):
         ee = self._get_ee()
         import ctypes
-        prototype = ctypes.CFUNCTYPE(
-            convert_to_ctypes(self.ret_type),
-            *[convert_to_ctypes(x) for x in self.arg_types])
+        restype = convert_to_ctypes(self.ret_type)
+        prototype = ctypes.CFUNCTYPE(restype,
+                                     *[convert_to_ctypes(x)
+                                       for x in self.arg_types])
+        if hasattr(restype, 'make_ctypes_prototype_wrapper'):
+            # See numba.utils.ComplexMixin for an example of
+            # make_ctypes_prototype_wrapper().
+            prototype = restype.make_ctypes_prototype_wrapper(prototype)
         if llvm:
             PY_CALL_TO_LLVM_CALL_MAP[self.func] = \
                 self.build_call_to_translated_function
@@ -1062,9 +1120,14 @@ def func_resolve_type(self, func, args):
         llvm_args = [arg.llvm(typ) for typ, arg in zip(typs, args)]
         return lfunc, llvm_args
 
-    def _init_ma_obj(self):
+    def init_multiarray(self):
         '''Builds the MultiarrayAPI object and adds a PyArray_API
-        variable to the current module under construction.'''
+        variable to the current module under construction.
+
+        Should be called by code generators (like
+        _LLVMModuleUtils.build_zeros_like()) that build Numpy-like
+        functionality to ensure the ma_obj member has been
+        initialized.'''
         if self.ma_obj is None:
             self.ma_obj = MultiarrayAPI()
             self.ma_obj.set_PyArray_API(self.mod)
@@ -1273,8 +1336,10 @@ def op_CALL_FUNCTION(self, i, op, arg):
         if __debug__:
             print("op_CALL_FUNCTION():", func)
         if func.val in PY_CALL_TO_LLVM_CALL_MAP:
-            self._init_ma_obj()
             res, ret_typ = PY_CALL_TO_LLVM_CALL_MAP[func.val](self, args)
+        elif isinstance(func.val, MethodReference):
+            res, ret_val = PY_CALL_TO_LLVM_CALL_MAP[func.val.py_method](
+                self, [func.val.object_var])
         else:
             func, args = self.func_resolve_type(func, args)
             if func is None: # A delayed-result (i.e. range or xrange)
@@ -1343,17 +1408,34 @@ def op_LOAD_ATTR(self, i, op, arg):
             # Make this a map on types in the future (thinking this is
             # what typemap was destined to do...)
             objarg_llvm_val = objarg.llvm()
+            res = None
             if __debug__:
                 print "op_LOAD_ATTR():", objarg_llvm_val.type
             if objarg_llvm_val.type == _numpy_array:
                 field_index = _numpy_array_field_ofs[self.names[arg]]
+                field_indices = [_int32_zero,
+                                 lc.Constant.int(_int32, field_index)]
+            elif objarg_llvm_val.type in (_complex64, _complex128):
+                # Re-wrap the instance, if we had to convert to an LLVM value.
+                objarg = objarg if objarg._llvm else Variable(objarg_llvm_val)
+                field_name = self.names[arg]
+                if field_name == 'real':
+                    res = self.builder.extract_value(objarg_llvm_val, 0)
+                elif field_name == 'imag':
+                    res = self.builder.extract_value(objarg_llvm_val, 1)
+                elif objarg_llvm_val.type == _complex64:
+                    res = MethodReference(objarg,
+                                          getattr(np.complex64, field_name))
+                else:
+                    res = MethodReference(objarg,
+                                          getattr(np.complex128, field_name))
             else:
-                raise NotImplementedError('LOAD_ATTR only supported for Numpy '
-                                          'arrays.')
-            res_addr = self.builder.gep(objarg_llvm_val, 
-                                        [_int32_zero,
-                                         lc.Constant.int(_int32, field_index)])
-            res = self.builder.load(res_addr)
+                raise NotImplementedError(
+                    'LOAD_ATTR does not presently support LLVM type %r.' %
+                    (str(objarg_llvm_val.type),))
+            if res is None:
+                res_addr = self.builder.gep(objarg_llvm_val, field_indices)
+                res = self.builder.load(res_addr)
         self.stack.append(Variable(res))
 
     def op_JUMP_ABSOLUTE(self, i, op, arg):

diff --git a/numba/utils.py b/numba/utils.py
@@ -1,5 +1,6 @@
 import opcode
-
+import ctypes
+from numpy import complex64, complex128
 
 def itercode(code):
     """Return a generator of byte-offset, opcode, and argument 
@@ -33,3 +34,50 @@ def debugout(*args):
     if __debug__:
         print("debugout (non-translated): %s" % (''.join((str(arg)
                                                           for arg in args)),))
+
+# NOTE: The following ctypes structures were inspired by Joseph
+# Heller's response to python-list question about ctypes complex
+# support.  In that response, he said these were only suitable for
+# Linux.  Might our milage vary?
+
+class ComplexMixin (object):
+    def _get (self):
+        # FIXME: Ensure there will not be a loss of precision here!
+        return self._numpy_ty_(self.real + (self.imag * 1j))
+
+    def _set (self, value):
+        self.real = value.real
+        self.imag = value.imag
+
+    value = property(_get, _set)
+
+    @classmethod
+    def from_param(cls, param):
+        ret_val = cls()
+        ret_val.value = param
+        return ret_val
+
+    @classmethod
+    def make_ctypes_prototype_wrapper(cls, ctypes_prototype):
+        '''This is a hack so that functions that return a complex type
+        will construct a new Python value from the result, making the
+        Numba compiled function a drop-in replacement for a Python
+        function.'''
+        # FIXME: See if there is some way of avoiding this additional
+        # wrapper layer.
+        def _make_complex_result_wrapper(in_func):
+            ctypes_function = ctypes_prototype(in_func)
+            def _complex_result_wrapper(*args, **kws):
+                # Return the value property, not the ComplexMixin
+                # instance built by ctypes.
+                return ctypes_function(*args, **kws).value
+            return _complex_result_wrapper
+        return _make_complex_result_wrapper
+
+class Complex64 (ctypes.Structure, ComplexMixin):
+    _fields_ = [('real', ctypes.c_float), ('imag', ctypes.c_float)]
+    _numpy_ty_ = complex64
+
+class Complex128 (ctypes.Structure, ComplexMixin):
+    _fields_ = [('real', ctypes.c_double), ('imag', ctypes.c_double)]
+    _numpy_ty_ = complex128
diff --git a/tests/test_all.py b/tests/test_all.py
@@ -9,6 +9,7 @@
 import unittest
 
 from test_cfg import TestCFG
+from test_complex import TestComplex
 from test_extern_call import TestExternCall
 from test_filter2d import TestFilter2d
 from test_forloop import TestForLoop

diff --git a/tests/test_complex.py b/tests/test_complex.py
@@ -0,0 +1,78 @@
+#! /usr/bin/env python
+# ______________________________________________________________________
+'''test_complex
+
+Test Numba's ability to generate code that supports complex numbers.
+'''
+# ______________________________________________________________________
+
+import unittest
+
+from numba.decorators import numba_compile
+from numba.utils import debugout
+
+import numpy
+
+# ______________________________________________________________________
+
+def get_real_fn (in_num):
+    return in_num.real
+
+def get_imag_fn (in_num):
+    ret_val = in_num.imag
+    debugout("get_imag_fn() in_num.real = ", in_num.real, " in_num.imag = ",
+             ret_val)
+    return ret_val
+
+def get_conj_fn (in_num):
+    return in_num.conjugate()
+
+def get_complex_constant_fn ():
+    return (3. + 4.j).conjugate()
+
+# ______________________________________________________________________
+
+class TestComplex (unittest.TestCase):
+    def test_get_real_fn (self):
+        num0 = 3 + 2j
+        num1 = numpy.complex128(num0)
+        compiled_get_real_fn = numba_compile(arg_types = ['D'])(get_real_fn)
+        self.assertEqual(compiled_get_real_fn(num0), 3.)
+        self.assertEqual(get_real_fn(num0), compiled_get_real_fn(num0))
+        self.assertEqual(compiled_get_real_fn(num1), 3.)
+        self.assertEqual(get_real_fn(num1), compiled_get_real_fn(num1))
+
+    def test_get_imag_fn (self):
+        num0 = 0 - 2j
+        print num0
+        num1 = numpy.complex128(num0)
+        print num1
+        compiled_get_imag_fn = numba_compile(arg_types = ['D'])(get_imag_fn)
+        self.assertEqual(compiled_get_imag_fn(num0), -2.)
+        self.assertEqual(get_imag_fn(num0), compiled_get_imag_fn(num0))
+        self.assertEqual(compiled_get_imag_fn(num1), -2.)
+        self.assertEqual(get_imag_fn(num1), compiled_get_imag_fn(num1))
+
+    def test_get_conj_fn (self):
+        num0 = 4 - 1.5j
+        num1 = numpy.complex128(num0)
+        compiled_get_conj_fn = numba_compile(arg_types = ['D'],
+                                             ret_type = 'D')(get_conj_fn)
+        self.assertEqual(compiled_get_conj_fn(num0), 4 + 1.5j)
+        self.assertEqual(get_conj_fn(num0), compiled_get_conj_fn(num0))
+        self.assertEqual(compiled_get_conj_fn(num1), 4 + 1.5j)
+        self.assertEqual(get_conj_fn(num1), compiled_get_conj_fn(num1))
+
+    def test_get_complex_constant_fn (self):
+        compiled_get_complex_constant_fn = numba_compile(
+            arg_types = [], ret_type = 'D')(get_complex_constant_fn)
+        self.assertEqual(get_complex_constant_fn(),
+                         compiled_get_complex_constant_fn())
+
+# ______________________________________________________________________
+
+if __name__ == "__main__":
+    unittest.main()
+
+# ______________________________________________________________________
+# End of test_complex.py