-
Notifications
You must be signed in to change notification settings - Fork 71
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #331 from eric-wieser/layout-in-type
numba: Change layout to be stored in type itself
- Loading branch information
Showing
5 changed files
with
97 additions
and
62 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,58 +1,75 @@ | ||
import numba | ||
import numba.extending | ||
from numba.extending import NativeValue | ||
import llvmlite.ir | ||
try: | ||
# module locations as of numba 0.49.0 | ||
from numba.core import cgutils, types | ||
from numba.core import 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 import types | ||
from numba.targets.imputils import lower_constant | ||
|
||
from .._layout import Layout | ||
from .._layout import Layout, _cached_property | ||
from .._layout_helpers import layout_short_name | ||
from .._multivector import MultiVector | ||
|
||
|
||
opaque_layout = types.Opaque('Opaque(Layout)') | ||
# In future we want to store some of the layout in the type (the `order` etc), | ||
# but store the `names` in the layout instances, so that we can reuse jitted | ||
# functions across different basis vector names. | ||
|
||
|
||
class LayoutType(types.Type): | ||
def __init__(self): | ||
super().__init__("LayoutType") | ||
class LayoutType(types.Dummy): | ||
def __init__(self, layout): | ||
self.obj = layout | ||
# cache of multivector types for this layout | ||
self._cache = {} | ||
layout_name = layout_short_name(layout) | ||
if layout_name is not None: | ||
name = "LayoutType({})".format(layout_name) | ||
else: | ||
name = "LayoutType({!r})".format(layout) | ||
super().__init__(name) | ||
|
||
|
||
@numba.extending.register_model(LayoutType) | ||
class LayoutModel(numba.extending.models.StructModel): | ||
def __init__(self, dmm, fe_typ): | ||
members = [ | ||
('obj', opaque_layout), | ||
] | ||
super().__init__(dmm, fe_typ, members) | ||
class LayoutModel(numba.extending.models.OpaqueModel): | ||
pass | ||
|
||
# The docs say we should use register a function to determine the numba type | ||
# with `@numba.extending.typeof_impl.register(LayoutType)`, but this is way | ||
# too slow (https://github.com/numba/numba/issues/5839). Instead, we use the | ||
# undocumented `_numba_type_` attribute, and use our own cache. | ||
|
||
@numba.extending.typeof_impl.register(Layout) | ||
def _typeof_Layout(val: Layout, c) -> LayoutType: | ||
return LayoutType() | ||
@_cached_property | ||
def _numba_type_(self): | ||
return LayoutType(self) | ||
|
||
Layout._numba_type_ = _numba_type_ | ||
|
||
# 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() | ||
def lower_constant_Layout(context, builder, typ: LayoutType, pyval: Layout) -> llvmlite.ir.Value: | ||
return context.get_dummy_value() | ||
|
||
|
||
@numba.extending.unbox(LayoutType) | ||
def unbox_Layout(typ, obj, context): | ||
layout = cgutils.create_struct_proxy(typ)(context.context, context.builder) | ||
layout.obj = obj | ||
return numba.extending.NativeValue(layout._getvalue()) | ||
def unbox_Layout(typ: LayoutType, obj: Layout, c) -> NativeValue: | ||
return NativeValue(c.context.get_dummy_value()) | ||
|
||
# Derived from the `Dispatcher` boxing | ||
|
||
@numba.extending.box(LayoutType) | ||
def box_Layout(typ, val, context): | ||
val = cgutils.create_struct_proxy(typ)(context.context, context.builder, value=val) | ||
obj = val.obj | ||
context.pyapi.incref(obj) | ||
def box_Layout(typ: LayoutType, val: llvmlite.ir.Value, c) -> Layout: | ||
obj = c.context.add_dynamic_addr(c.builder, id(typ.obj), info=typ.name) | ||
c.pyapi.incref(obj) | ||
return obj | ||
|
||
# methods | ||
|
||
@numba.extending.overload_method(LayoutType, 'MultiVector') | ||
def Layout_MultiVector(self, value): | ||
def impl(self, value): | ||
return MultiVector(self, value) | ||
return impl |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters