Experimental AST rewriter and JIT decorator #326

hugohadfield · 2020-06-05T16:23:38Z

No description provided.

clifford/jit_func.py

clifford/_ast_transformer.py

clifford/_layout.py

clifford/test/test_jit_func.py

eric-wieser · 2020-06-05T18:06:54Z

clifford/test/test_jit_func.py

+
+        np.testing.assert_allclose(test_func(e1, e2, einf).value, slow_test_func(e1, e2, einf).value)
+
+    def test_benchmark(self):


Let's make this an actual benchmark:

Suggested change

def test_benchmark(self):

@pytest.mark.parametrize('use_jit', [False, True])

def test_benchmark(self, use_jit, benchmark):

clifford/test/test_jit_func.py

lgtm-com · 2020-06-05T18:07:55Z

This pull request introduces 1 alert when merging 8094a61 into c07db7d - view on LGTM.com

new alerts:

1 for Except block handles 'BaseException'

clifford/jit_func.py

eric-wieser · 2020-06-05T18:10:47Z

clifford/_layout.py

+            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


Tempting to make these shorter:

Suggested change

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 lambda a, b: omt_func(a, b)

elif isinstance(a, types.Array) and isinstance(b, types.abstract.Number):

return lambda a, b: omt_func(a, as_ga(b))

elif isinstance(a, types.abstract.Number) and isinstance(b, types.Array):

return lambda a, b: omt_func(as_ga(a), b)

else:

return lambda a, b: return a^b

clifford/jit_func.py

clifford/_ast_transformer.py

eric-wieser

Happy to merge this into the experimental branch, if you think you're stabilized now

lgtm-com · 2020-06-06T08:40:54Z

This pull request introduces 2 alerts when merging e0263f8 into c07db7d - view on LGTM.com

new alerts:

1 for Unused local variable
1 for Except block handles 'BaseException'

lgtm-com · 2020-06-06T09:08:22Z

This pull request introduces 2 alerts when merging e878dbe into c07db7d - view on LGTM.com

new alerts:

1 for Unused local variable
1 for Except block handles 'BaseException'

eric-wieser · 2020-06-06T09:23:12Z

clifford/jit_func.py

+        fname = func.__name__
+        source = inspect.getsource(func)
+        # Remove the decorator first line.
+        source = '\n'.join(source.splitlines()[1:])


It would be better to remove this at the ast level

lgtm-com · 2020-06-06T16:38:30Z

This pull request introduces 2 alerts when merging 482b091 into c07db7d - view on LGTM.com

new alerts:

1 for Unused local variable
1 for Except block handles 'BaseException'

hugohadfield · 2020-06-08T08:39:13Z

clifford/test/test_jit_func.py

+
+        test_compound_func = jit_func(self.layout,
+                                      mv_constants={'test_func_1': test_func_1_jit,
+                                                    'test_func_2': test_func_2_jit})(compound_func)


This is a way to nest functions by using(abusing?) the mv_constants argument, it is clearly not a great solution but I'm not sure how to pull it off otherwise. Perhaps the decorator has knowledge of what functions are available?

The decorator can probably use inspect.currentframe().f_back.f_locals

I suspect you should be using f_globals instead of globals() too.

Lots of changes since I lasted looked.

clifford/_ast_transformer.py

eric-wieser · 2020-06-08T09:25:18Z

clifford/test/test_jit_func.py

+
+        test_compound_func = jit_func(self.layout,
+                                      mv_constants={'test_func_1': test_func_1_jit,
+                                                    'test_func_2': test_func_2_jit})(compound_func)


The decorator can probably use inspect.currentframe().f_back.f_locals

eric-wieser · 2020-06-08T09:27:30Z

clifford/_layout.py

+    @_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 project_to_grade_func(self):
+        return get_project_to_grade_func(self)
+
+    @_cached_property
+    def overload_call_func(self):
+        return get_overload_call(self)


I don't think these really belong here. You've put them here because it makes them easy to cache, but I think it would be better to cache them manually in a weakref.WeakKeyDictionary.

So perhaps something like

def weak_cache(f): _cache = weakref.WeakKeyDictionary() @functools.wraps(f) def wrapped(*args, **kwargs): a, *args = args try: return _cache[a] except KeyError: ret =_cache[a] = f(a, *args, **kwargs) return ret wrapped._cache = _cache return wrapped @weak_cache def _get_jit_impls(layout): return { 'as_ga': get_as_ga_value_vector_func(layout), 'ga_add': get_overload_add_func(layout), 'ga_sub': get_overload_sub_func(layout), 'ga_mul': get_overload_mul_func(layout), 'ga_xor': get_overload_xor_func(layout), 'ga_or': get_overload_or_func(layout), 'ga_rev': get_overload_reverse_func(layout), 'ga_call': get_overload_call_func(layout), }

eric-wieser · 2020-06-08T09:27:49Z

clifford/jit_func.py

@@ -0,0 +1,92 @@
+


Perhaps clifford/experimental/jit_func

eric-wieser · 2020-06-08T09:34:42Z

clifford/jit_func.py

+        # Remove the decorators from the function
+        tree = DecoratorRemover().visit(tree)


Strictly you want to leave behind all the non-jitfunc decorators, but that's fine as a todo.

clifford/jit_func.py

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

eric-wieser · 2020-06-09T11:12:09Z

clifford/__init__.py

@@ -117,7 +117,8 @@ def linear_operator_as_matrix(func, input_blades, output_blades):
    ndimout = len(output_blades)
    mat = np.zeros((ndimout, ndimin))
    for i, b in enumerate(input_blades):
-        mat[:, i] = np.array([func(b)[j] for j in output_blades])


Probably makes sense to spin out a quick PR with just this fix against master.

yeah I was thinking that

eric-wieser · 2020-06-18T12:57:08Z

What do we want to do with this now that we have proper numba support?

hugohadfield · 2020-06-18T14:37:22Z

Shall we just close it? I don't think it adds anything on top of what we have, although it was an interesting and fun idea

hugohadfield added 4 commits June 5, 2020 16:35

Added numba overloaded functions to layout

9e3724c

Added a GA specific ast transformer

03826c5

Added a jit_func decorator to ast transform and numba jit

ef61257

Corrected jit_func, added a test

c13bc94

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/jit_func.py Outdated Show resolved Hide resolved

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/_ast_transformer.py Outdated Show resolved Hide resolved

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/_layout.py Outdated Show resolved Hide resolved

hugohadfield added 6 commits June 5, 2020 18:20

remove duplication in ast_transformer

51f8a42

convert to abstract numeric types in the numba jit overload

8022092

Improved handling globals, added a TODO

f14521b

Added ast_pretty warning if not installed

5fdbb86

removed unnescary print

d6c6e06

Added reversion to AST rewriter and JIT

8094a61

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/test/test_jit_func.py Outdated Show resolved Hide resolved

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/test/test_jit_func.py Outdated Show resolved Hide resolved

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/test/test_jit_func.py Outdated Show resolved Hide resolved

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/test/test_jit_func.py Outdated Show resolved Hide resolved

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/jit_func.py Outdated Show resolved Hide resolved

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/jit_func.py Outdated Show resolved Hide resolved

Added grade selection via the call syntax

1767342

eric-wieser reviewed Jun 5, 2020

View reviewed changes

clifford/_ast_transformer.py Outdated Show resolved Hide resolved

eric-wieser previously approved these changes Jun 5, 2020

View reviewed changes

hugohadfield added 4 commits June 6, 2020 09:17

Set up pytest benchmark

81601ce

Make node visitation recursive for Call

d905393

Add ImportError type for astpretty

750ec85

Improve warning whitespace

e0263f8

Make the Call rewrite exception an AttributeError

e878dbe

eric-wieser reviewed Jun 6, 2020

View reviewed changes

Moved the decorator removal to the AST level

482b091

hugohadfield added 5 commits June 7, 2020 15:10

Add scalar and multivector constants to decorator arguments

ff9648d

Fix nested function call transformer

5d27874

Improve speed of linear_operator_to_matrix

6c2cea6

Add testing for new jit decorator features

307874f

Added a nested jitted function test

c5be87a

hugohadfield commented Jun 8, 2020

View reviewed changes

Fixed flake8 complaints

8f02960

eric-wieser reviewed Jun 8, 2020

View reviewed changes

hugohadfield and others added 4 commits June 9, 2020 10:41

Apply suggestions from Eric code review

8e96d81

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

Fix up review comments

2315f3f

Moved jit_impls into jit_func

87a41b9

Moved jit_func into an experimental directory

ccf5551

eric-wieser reviewed Jun 9, 2020

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Experimental AST rewriter and JIT decorator #326

Experimental AST rewriter and JIT decorator #326

hugohadfield commented Jun 5, 2020

eric-wieser Jun 5, 2020 •

edited

lgtm-com bot commented Jun 5, 2020

eric-wieser Jun 5, 2020

eric-wieser left a comment

lgtm-com bot commented Jun 6, 2020

lgtm-com bot commented Jun 6, 2020

eric-wieser Jun 6, 2020

lgtm-com bot commented Jun 6, 2020

hugohadfield Jun 8, 2020

eric-wieser Jun 8, 2020

eric-wieser Jun 9, 2020

eric-wieser Jun 8, 2020

eric-wieser Jun 8, 2020

eric-wieser Jun 8, 2020

eric-wieser Jun 8, 2020

eric-wieser Jun 8, 2020

eric-wieser Jun 9, 2020

eric-wieser Jun 9, 2020

hugohadfield Jun 9, 2020

eric-wieser commented Jun 18, 2020

hugohadfield commented Jun 18, 2020


		np.testing.assert_allclose(test_func(e1, e2, einf).value, slow_test_func(e1, e2, einf).value)

		def test_benchmark(self):

	def test_benchmark(self):
	@pytest.mark.parametrize('use_jit', [False, True])
	def test_benchmark(self, use_jit, benchmark):

		# Remove the decorators from the function
		tree = DecoratorRemover().visit(tree)

Experimental AST rewriter and JIT decorator #326

Are you sure you want to change the base?

Experimental AST rewriter and JIT decorator #326

Conversation

hugohadfield commented Jun 5, 2020

eric-wieser Jun 5, 2020 • edited

Choose a reason for hiding this comment

lgtm-com bot commented Jun 5, 2020

Choose a reason for hiding this comment

eric-wieser left a comment

Choose a reason for hiding this comment

lgtm-com bot commented Jun 6, 2020

lgtm-com bot commented Jun 6, 2020

Choose a reason for hiding this comment

lgtm-com bot commented Jun 6, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eric-wieser commented Jun 18, 2020

hugohadfield commented Jun 18, 2020

eric-wieser Jun 5, 2020 •

edited