Add support for constant multivectors in jitted functions.

This changes the approach for storing `Layout` objects to be using globals via `add_dynamic_addr`.
This will disable numba caching, but we don't use that anyway.

Author: eric-wieser

clifford/numba/_layout.py

@@ -3,15 +3,16 @@
 try:
     # module locations as of numba 0.49.0
     from numba.core import cgutils, types
+    from numba.core.imputils import lower_constant
 except ImportError:
     # module locations prior to numba 0.49.0
     from numba import cgutils, types
+    from numba.targets.imputils import lower_constant
 
 from .._layout import Layout
 
 
-# Taken from numba_passthru
-opaque_pyobject = types.Opaque('Opaque(PyObject)')
+opaque_layout = types.Opaque('Opaque(Layout)')
 
 
 class LayoutType(types.Type):
@@ -23,7 +24,7 @@ def __init__(self):
 class LayoutModel(numba.extending.models.StructModel):
     def __init__(self, dmm, fe_typ):
         members = [
-            ('meminfo', types.MemInfoPointer(opaque_pyobject)),
+            ('obj', opaque_layout),
         ]
         super().__init__(dmm, fe_typ, members)
 
@@ -33,18 +34,25 @@ def _typeof_Layout(val: Layout, c) -> LayoutType:
     return LayoutType()
 
 
-# Derived from numba_passthru
+# Derived from the `Dispatcher` boxing
+
+@lower_constant(LayoutType)
+def lower_constant_dispatcher(context, builder, typ, pyval):
+    layout = cgutils.create_struct_proxy(typ)(context, builder)
+    layout.obj = context.add_dynamic_addr(builder, id(pyval), info=type(pyval).__name__)
+    return layout._getvalue()
+
 
 @numba.extending.unbox(LayoutType)
 def unbox_Layout(typ, obj, context):
     layout = cgutils.create_struct_proxy(typ)(context.context, context.builder)
-    layout.meminfo = context.pyapi.nrt_meminfo_new_from_pyobject(obj, obj)
+    layout.obj = obj
     return numba.extending.NativeValue(layout._getvalue())
 
 
 @numba.extending.box(LayoutType)
 def box_Layout(typ, val, context):
     val = cgutils.create_struct_proxy(typ)(context.context, context.builder, value=val)
-    obj = context.context.nrt.meminfo_data(context.builder, val.meminfo)
+    obj = val.obj
     context.pyapi.incref(obj)
     return obj

clifford/numba/_multivector.py

@@ -10,12 +10,12 @@
 try:
     # module locations as of numba 0.49.0
     import numba.np.numpy_support as _numpy_support
-    from numba.core.imputils import impl_ret_borrowed
+    from numba.core.imputils import impl_ret_borrowed, lower_constant
     from numba.core import cgutils, types
 except ImportError:
     # module locations prior to numba 0.49.0
     import numba.numpy_support as _numpy_support
-    from numba.targets.imputils import impl_ret_borrowed
+    from numba.targets.imputils import impl_ret_borrowed, lower_constant
     from numba import cgutils, types
 
 from .._multivector import MultiVector
@@ -46,7 +46,8 @@ def layout_type(self):
 
 @numba.extending.typeof_impl.register(MultiVector)
 def _typeof_MultiVector(val: MultiVector, c) -> MultiVectorType:
-    return MultiVectorType(dtype=val.value.dtype)
+    val._numba_type_ = MultiVectorType(dtype=val.value.dtype)
+    return val._numba_type_
 
 
 @numba.extending.register_model(MultiVectorType)
@@ -77,6 +78,18 @@ def impl_MultiVector(context, builder, sig, args):
     return impl_ret_borrowed(context, builder, sig.return_type, mv._getvalue())
 
 
+@lower_constant(MultiVectorType)
+def lower_constant_MultiVector(context, builder, typ: MultiVectorType, pyval: MultiVector):
+    value = context.get_constant_generic(builder, typ.value_type, pyval.value)
+    layout = context.get_constant_generic(builder, typ.layout_type, pyval.layout)
+    return impl_ret_borrowed(
+        context,
+        builder,
+        typ,
+        cgutils.pack_struct(builder, (layout, value)),
+    )
+
+
 @numba.extending.unbox(MultiVectorType)
 def unbox_MultiVector(typ: MultiVectorType, obj: MultiVector, c) -> NativeValue:
     value = c.pyapi.object_getattr_string(obj, "value")

clifford/test/test_numba_extensions.py

@@ -41,3 +41,10 @@ def add(a, b):
         ab = add(e1, e2)
         assert ab == e1 + e2
         assert ab.layout is e1.layout
+
+    def test_constant_multivector(self):
+        @numba.njit
+        def add_e1(a):
+            return cf.MultiVector(a.layout, a.value + e1.value)
+
+        assert add_e1(e2) == e1 + e2