Experimental AST rewriter and JIT decorator

clifford/_ast_transformer.py

@@ -0,0 +1,44 @@
+
+import ast
+
+
+class GATransformer(ast.NodeTransformer):
+    """
+    This is an AST transformer that converts operations into
+    JITable counterparts that work on MultiVector value arrays.
+    We crawl the AST and convert BinOps and UnaryOps into numba
+    overloaded functions.
+    """
+    def visit_BinOp(self, node):
+        ops = {
+            ast.Mult: 'ga_mul',
+            ast.BitXor: 'ga_xor',
+            ast.BitOr: 'ga_or',
+            ast.Add: 'ga_add',
+            ast.Sub: 'ga_sub',
+        }
+        try:
+            func_name = ops[type(node.op)]
+        except KeyError:
+            return node
+        else:
+            return ast.Call(
+                func=ast.Name(id=func_name, ctx=ast.Load()),
+                args=[self.visit(node.left), self.visit(node.right)],
+                keywords=[]
+            )
+
+    def visit_UnaryOp(self, node):
+        ops = {
+            ast.Invert: 'ga_rev'
+        }
+        try:
+            func_name = ops[type(node.op)]
+        except KeyError:
+            return node
+        else:
+            return ast.Call(
+                func=ast.Name(id=func_name, ctx=ast.Load()),
+                args=[self.visit(node.operand)],
+                keywords=[]
+            )

clifford/_layout.py

@@ -5,6 +5,10 @@
 import numpy as np
 import sparse
 
+from numba.extending import overload
+from numba import types
+
+
 # TODO: move some of these functions to this file if they're not useful anywhere
 # else
 import clifford as cf
@@ -175,6 +179,201 @@ def construct_graded_mt(
     return sparse.COO(coords=coords, data=mult_table_vals, shape=(dims, dims, dims))
 
 
+def get_as_ga_vector_func(layout):
+    """
+    Returns a function that converts a scalar into a GA value vector
+    for the given algebra
+    """
+    scalar_index = layout._basis_blade_order.bitmap_to_index[0]
+    ndims = layout.gaDims
+    @_numba_utils.njit
+    def as_ga_value_vector(x):
+        op = np.zeros(ndims)
+        op[scalar_index] = x
+        return op
+    return as_ga_value_vector
+
+
+def get_overload_add(layout):
+    """
+    Returns an overloaded JITed function that works on
+    MultiVector value arrays
+    """
+    scalar_index = layout._basis_blade_order.bitmap_to_index[0]
+
+    def ga_add(a, b):
+        # dummy function to overload
+        pass
+
+    @overload(ga_add, inline='always')
+    def ol_ga_add(a, b):
+        if isinstance(a, types.abstract.Number) and isinstance(b, types.Array):
+            def impl(a, b):
+                op = b.astype(np.float32)
+                op[scalar_index] += a
+                return op
+            return impl
+        elif isinstance(a, types.Array) and isinstance(b, types.abstract.Number):
+            def impl(a, b):
+                op = a.astype(np.float32)
+                op[scalar_index] += b
+                return op
+            return impl
+        else:
+            def impl(a, b):
+                return a + b
+            return impl
+
+    return ga_add
+
+
+def get_overload_sub(layout):
+    """
+    Returns an overloaded JITed function that works on
+    MultiVector value arrays
+    """
+    scalar_index = layout._basis_blade_order.bitmap_to_index[0]
+
+    def ga_sub(a, b):
+        # dummy function to overload
+        pass
+
+    @overload(ga_sub, inline='always')
+    def ol_ga_sub(a, b):
+        if isinstance(a, types.abstract.Number) and isinstance(b, types.Array):
+            def impl(a, b):
+                op = -b.astype(np.float32)
+                op[scalar_index] += a
+                return op
+            return impl
+        elif isinstance(a, types.Array) and isinstance(b, types.abstract.Number):
+            def impl(a, b):
+                op = a.astype(np.float32)
+                op[scalar_index] -= b
+                return op
+            return impl
+        else:
+            def impl(a, b):
+                return a - b
+            return impl
+
+    return ga_sub
+
+
+def get_overload_mul(layout):
+    """
+    Returns an overloaded JITed function that works on
+    MultiVector value arrays
+    """
+    def ga_mul(a, b):
+        # dummy function to overload
+        pass
+
+    gmt_func = layout.gmt_func
+    @overload(ga_mul, inline='always')
+    def ol_ga_mul(a, b):
+        if isinstance(a, types.Array) and isinstance(b, types.Array):
+            def impl(a, b):
+                return gmt_func(a, b)
+            return impl
+        else:
+            def impl(a, b):
+                return a*b
+            return impl
+
+    return ga_mul
+
+
+def get_overload_xor(layout):
+    """
+    Returns an overloaded JITed function that works on
+    MultiVector value arrays
+    """
+    def ga_xor(a, b):
+        # dummy function to overload
+        pass
+
+    as_ga = layout.as_ga_value_vector_func
+    omt_func = layout.omt_func
+    @overload(ga_xor, inline='always')
+    def ol_ga_xor(a, b):
+        if isinstance(a, types.Array) and isinstance(b, types.Array):
+            def impl(a, b):
+                return omt_func(a, b)
+            return impl
+        elif isinstance(a, types.Array) and isinstance(b, types.abstract.Number):
+            def impl(a, b):
+                return omt_func(a, as_ga(b))
+            return impl
+        elif isinstance(a, types.abstract.Number) and isinstance(b, types.Array):
+            def impl(a, b):
+                return omt_func(as_ga(a), b)
+            return impl
+        else:
+            def impl(a, b):
+                return a^b
+            return impl
+
+    return ga_xor
+
+
+def get_overload_or(layout):
+    """
+    Returns an overloaded JITed function that works on
+    MultiVector value arrays
+    """
+    def ga_or(a, b):
+        # dummy function to overload
+        pass
+
+    as_ga = layout.as_ga_value_vector_func
+    imt_func = layout.imt_func
+    @overload(ga_or, inline='always')
+    def ol_ga_or(a, b):
+        if isinstance(a, types.Array) and isinstance(b, types.Array):
+            def impl(a, b):
+                return imt_func(a, b)
+            return impl
+        elif isinstance(a, types.Array) and isinstance(b, types.abstract.Number):
+            def impl(a, b):
+                return imt_func(a, as_ga(b))
+            return impl
+        elif isinstance(a, types.abstract.Number) and isinstance(b, types.Array):
+            def impl(a, b):
+                return imt_func(as_ga(a), b)
+            return impl
+        else:
+            def impl(a, b):
+                return a|b
+            return impl
+
+    return ga_or
+
+
+def get_overload_reverse(layout):
+    """
+    Returns an overloaded JITed function that works on
+    MultiVector value arrays
+    """
+    def ga_rev(x):
+        # dummy function to overload
+        pass
+
+    adjoint_func = layout.adjoint_func
+    @overload(ga_rev, inline='always')
+    def ol_ga_rev(x):
+        if isinstance(x, types.Array):
+            def impl(x):
+                return adjoint_func(x)
+            return impl
+        else:
+            def impl(x):
+                return ~x
+            return impl
+
+    return ga_rev
+
+
 class Layout(object):
     r""" Layout stores information regarding the geometric algebra itself and the
     internal representation of multivectors.
@@ -372,6 +571,11 @@ def __init__(self, *args, **kw):
         self.dual_func
         self.vee_func
         self.inv_func
+        self.overload_mul_func
+        self.overload_xor_func
+        self.overload_or_func
+        self.overload_add_func
+        self.overload_sub_func
 
     @_cached_property
     def gmt(self):
@@ -572,6 +776,10 @@ def comp_func(Xval):
             return Yval
         return comp_func
 
+    @_cached_property
+    def as_ga_value_vector_func(self):
+        return get_as_ga_vector_func(self)
+
     @_cached_property
     def gmt_func(self):
         return get_mult_function(self.gmt, self.gradeList)
@@ -596,6 +804,30 @@ def left_complement_func(self):
     def right_complement_func(self):
         return self._gen_complement_func(omt=self.omt.T)
 
+    @_cached_property
+    def overload_mul_func(self):
+        return get_overload_mul(self)
+
+    @_cached_property
+    def overload_xor_func(self):
+        return get_overload_xor(self)
+
+    @_cached_property
+    def overload_or_func(self):
+        return get_overload_or(self)
+
+    @_cached_property
+    def overload_add_func(self):
+        return get_overload_add(self)
+
+    @_cached_property
+    def overload_sub_func(self):
+        return get_overload_sub(self)
+
+    @_cached_property
+    def overload_reverse_func(self):
+        return get_overload_reverse(self)
+
     @_cached_property
     def adjoint_func(self):
         '''

clifford/jit_func.py

@@ -0,0 +1,80 @@
+
+import ast
+try:
+    import astpretty
+    AST_PRETTY_AVAILABLE = True
+except:
+    AST_PRETTY_AVAILABLE = False
+import inspect
+import warnings
+from ._numba_utils import njit
+from ._ast_transformer import GATransformer
+
+
+class jit_func(object):
+    """
+    This is a JIT decorator that re-writes the AST and then numba JITs
+    the resulting function.
+    """
+    def __init__(self, layout, ast_debug=False):
+        self.layout = layout
+        if AST_PRETTY_AVAILABLE:
+            self.ast_debug = ast_debug
+        else:
+            if ast_debug:
+                warnings.warn('''
+The ast_debug flag is set to True, but the astpretty module is not importable.
+To see ast_debug output please pip install astpretty
+''')
+            self.ast_debug = False
+
+    def __call__(self, func):
+        # Get the function source
+        fname = func.__name__
+        source = inspect.getsource(func)
+        # Remove the decorator first line.
+        source = '\n'.join(source.splitlines()[1:])
+        # Remove the indentation
+        indentation = source.splitlines()[0].find('def')
+        source = '\n'.join([line[indentation:] for line in source.splitlines()])
+
+        # Re-write the ast
+        tree = ast.parse(source)
+        if self.ast_debug:
+            print('\n\n\n\n TRANSFORMING FROM \n\n\n\n')
+            astpretty.pprint(tree)
+
+        tree = GATransformer().visit(tree)
+        ast.fix_missing_locations(tree)
+
+        if self.ast_debug:
+            print('\n\n\n\n TRANSFORMING TO \n\n\n\n')
+            astpretty.pprint(tree)
+
+        # Set things up into locals and globals so that they JIT ok...
+        locals_dict = {'as_ga': self.layout.as_ga_value_vector_func,
+                       'ga_add': self.layout.overload_add_func,
+                       'ga_sub': self.layout.overload_sub_func,
+                       'ga_mul': self.layout.overload_mul_func,
+                       'ga_xor': self.layout.overload_xor_func,
+                       'ga_or': self.layout.overload_or_func,
+                       'ga_rev': self.layout.overload_reverse_func}
+        # TODO: Work out a better way to deal with changes to globals
+        globs = {}
+        for k, v in globals().items():
+            globs[k] = v
+        for k, v in locals_dict.items():
+            globs[k] = v
+
+        # Compile the function
+        co = compile(tree, '<ast>', "exec")
+        exec(co, globs, locals_dict)
+        new_func = locals_dict[fname]
+
+        # JIT the function
+        jitted_func = njit(new_func)
+
+        # Wrap the JITed function
+        def wrapper(*args, **kwargs):
+            return self.layout.MultiVector(value=jitted_func(*[a.value for a in args], **kwargs))
+        return wrapper