Skip to content

Commit 26192df

Browse files
authored
Merge pull request #7061 from sklam/enh/arrayallocapi
Add ndarray allocator classmethod
2 parents 0e53a23 + f756fd0 commit 26192df

18 files changed

Lines changed: 408 additions & 256 deletions

‎.flake8‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,6 @@ exclude =
186186
numba/tests/test_buffer_protocol.py
187187
numba/tests/test_casting.py
188188
numba/tests/test_lists.py
189-
numba/tests/test_blackscholes.py
190189
numba/tests/test_array_analysis.py
191190
numba/tests/test_serialize.py
192191
numba/tests/test_iteration.py

‎numba/core/cpu.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,8 @@ def finalize(self, flags, options):
296296

297297
flags.inherit_if_not_set("fastmath")
298298

299+
flags.inherit_if_not_set("error_model", default="python")
300+
299301
# Add "target_backend" as a option that inherits from the caller
300302
flags.inherit_if_not_set("target_backend")
301303

‎numba/core/errors.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -651,6 +651,17 @@ def __init__(self, exception):
651651
self.old_exception = exception
652652

653653

654+
class InternalTargetMismatchError(InternalError):
655+
"""For signalling a target mismatch error occurred internally within the
656+
compiler.
657+
"""
658+
def __init__(self, kind, target_hw, hw_clazz):
659+
msg = (f"{kind.title()} being resolved on a target from which it does "
660+
f"not inherit. Local target is {target_hw}, declared "
661+
f"target class is {hw_clazz}.")
662+
super().__init__(msg)
663+
664+
654665
class RequireLiteralValue(TypingError):
655666
"""
656667
For signalling that a function's typing requires a constant value for

‎numba/core/extending.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -203,10 +203,12 @@ def _overload_method_common(typ, attr, **kwargs):
203203
from numba.core.typing.templates import make_overload_method_template
204204

205205
def decorate(overload_func):
206+
copied_kwargs = kwargs.copy() # avoid mutating parent dict
206207
template = make_overload_method_template(
207208
typ, attr, overload_func,
208-
inline=kwargs.get('inline', 'never'),
209-
prefer_literal=kwargs.get('prefer_literal', False)
209+
inline=copied_kwargs.pop('inline', 'never'),
210+
prefer_literal=copied_kwargs.pop('prefer_literal', False),
211+
**copied_kwargs,
210212
)
211213
infer_getattr(template)
212214
overload(overload_func, **kwargs)(overload_func)

‎numba/core/target_extension.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from abc import ABC, abstractmethod
22
from numba.core.registry import DelayedRegistry, CPUDispatcher
33
from numba.core.decorators import jit
4+
from numba.core.errors import InternalTargetMismatchError
45
from threading import local as tls
56

67

@@ -77,6 +78,35 @@ def resolve_dispatcher_from_str(target_str):
7778
return dispatcher_registry[target_hw]
7879

7980

81+
def _get_local_target_checked(tyctx, hwstr, reason):
82+
"""Returns the local target if it is compatible with the given target
83+
name during a type resolution; otherwise, raises an exception.
84+
85+
Parameters
86+
----------
87+
tyctx: typing context
88+
hwstr: str
89+
target name to check against
90+
reason: str
91+
Reason for the resolution. Expects a noun.
92+
Returns
93+
-------
94+
target_hw : Target
95+
96+
Raises
97+
------
98+
InternalTargetMismatchError
99+
"""
100+
# Get the class for the target declared by the function
101+
hw_clazz = resolve_target_str(hwstr)
102+
# get the local target
103+
target_hw = get_local_target(tyctx)
104+
# make sure the target_hw is in the MRO for hw_clazz else bail
105+
if not target_hw.inherits_from(hw_clazz):
106+
raise InternalTargetMismatchError(reason, target_hw, hw_clazz)
107+
return target_hw
108+
109+
80110
class JitDecorator(ABC):
81111

82112
@abstractmethod

‎numba/core/types/functions.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,10 @@ def copy(self, this):
391391

392392
@property
393393
def key(self):
394-
return self.typing_key, self.this
394+
# FIXME: With target-overload, the MethodTemplate can change depending
395+
# on the target.
396+
unique_impl = getattr(self.template, "_overload_func", None)
397+
return self.typing_key, self.this, unique_impl
395398

396399
def get_impl_key(self, sig):
397400
"""

‎numba/core/typing/context.py‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -284,7 +284,33 @@ def core(typ):
284284
return attrty
285285

286286
def find_matching_getattr_template(self, typ, attr):
287-
for template in self._get_attribute_templates(typ):
287+
from numba.core.target_extension import (target_registry,
288+
get_local_target)
289+
290+
templates = list(self._get_attribute_templates(typ))
291+
292+
# get the current target target
293+
target_hw = get_local_target(self)
294+
295+
# fish out templates that are specific to the target if a target is
296+
# specified
297+
DEFAULT_TARGET = 'generic'
298+
usable = []
299+
for ix, temp_cls in enumerate(templates):
300+
md = getattr(temp_cls, "metadata", {})
301+
hw = md.get('target', DEFAULT_TARGET)
302+
if hw is not None:
303+
hw_clazz = target_registry[hw]
304+
if target_hw.inherits_from(hw_clazz):
305+
usable.append((temp_cls, hw_clazz, ix))
306+
307+
# sort templates based on target specificity
308+
def key(x):
309+
return target_hw.__mro__.index(x[1])
310+
311+
order = [x[0] for x in sorted(usable, key=key)]
312+
313+
for template in order:
288314
return_type = template.resolve(typ, attr)
289315
if return_type is not None:
290316
return {

‎numba/core/typing/templates.py‎

Lines changed: 56 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,11 @@
1313

1414
import numba
1515
from numba.core import types, utils
16-
from numba.core.errors import TypingError, InternalError
16+
from numba.core.errors import (
17+
TypingError,
18+
InternalError,
19+
InternalTargetMismatchError,
20+
)
1721
from numba.core.cpu_options import InlineOptions
1822

1923
# info store for inliner callback functions e.g. cost model
@@ -887,23 +891,13 @@ def generic(self, args, kws):
887891
"""
888892
Type the intrinsic by the arguments.
889893
"""
890-
from numba.core.target_extension import (get_local_target,
891-
resolve_target_str,
894+
from numba.core.target_extension import (_get_local_target_checked,
892895
dispatcher_registry)
893896
from numba.core.imputils import builtin_registry
894897

895898
cache_key = self.context, args, tuple(kws.items())
896899
hwstr = self.metadata.get('target', 'generic')
897-
# Get the class for the target declared by the function
898-
hw_clazz = resolve_target_str(hwstr)
899-
# get the local target
900-
target_hw = get_local_target(self.context)
901-
# make sure the target_hw is in the MRO for hw_clazz else bail
902-
if not target_hw.inherits_from(hw_clazz):
903-
msg = (f"Intrinsic being resolved on a target from which it does "
904-
f"not inherit. Local target is {target_hw}, declared "
905-
f"target class is {hw_clazz}.")
906-
raise InternalError(msg)
900+
target_hw = _get_local_target_checked(self.context, hwstr, "intrinsic")
907901
disp = dispatcher_registry[target_hw]
908902
tgtctx = disp.targetdescr.target_context
909903
# This is all workarounds...
@@ -990,28 +984,12 @@ def make_intrinsic_template(handle, defn, name, kwargs):
990984

991985

992986
class AttributeTemplate(object):
993-
_initialized = False
994-
995987
def __init__(self, context):
996-
self._lazy_class_init()
997988
self.context = context
998989

999990
def resolve(self, value, attr):
1000991
return self._resolve(value, attr)
1001992

1002-
@classmethod
1003-
def _lazy_class_init(cls):
1004-
if not cls._initialized:
1005-
cls.do_class_init()
1006-
cls._initialized = True
1007-
1008-
@classmethod
1009-
def do_class_init(cls):
1010-
"""
1011-
Class-wide initialization. Can be overridden by subclasses to
1012-
register permanent typing or target hooks.
1013-
"""
1014-
1015993
def _resolve(self, value, attr):
1016994
fn = getattr(self, "resolve_%s" % attr, None)
1017995
if fn is None:
@@ -1038,15 +1016,33 @@ class _OverloadAttributeTemplate(AttributeTemplate):
10381016
def __init__(self, context):
10391017
super(_OverloadAttributeTemplate, self).__init__(context)
10401018
self.context = context
1019+
self._init_once()
10411020

1042-
@classmethod
1043-
def do_class_init(cls):
1044-
"""
1045-
Register attribute implementation.
1021+
def _get_target_registry(self):
1022+
"""Returns the registry for the current target
10461023
"""
1047-
from numba.core.imputils import lower_getattr
1024+
from numba.core.target_extension import (_get_local_target_checked,
1025+
dispatcher_registry)
1026+
hwstr = self.metadata.get('target', 'generic')
1027+
target_hw = _get_local_target_checked(self.context, hwstr, "attribute")
1028+
# Get resgistry for the current hardware
1029+
disp = dispatcher_registry[target_hw]
1030+
tgtctx = disp.targetdescr.target_context
1031+
tgtctx.refresh()
1032+
if builtin_registry in tgtctx._registries:
1033+
reg = builtin_registry
1034+
else:
1035+
# Pick a registry in which to install intrinsics
1036+
registries = iter(tgtctx._registries)
1037+
reg = next(registries)
1038+
return reg
1039+
1040+
def _init_once(self):
1041+
cls = type(self)
10481042
attr = cls._attr
10491043

1044+
lower_getattr = self._get_target_registry().lower_getattr
1045+
10501046
@lower_getattr(cls.key, attr)
10511047
def getattr_impl(context, builder, typ, value):
10521048
typingctx = context.typing_context
@@ -1082,24 +1078,30 @@ class _OverloadMethodTemplate(_OverloadAttributeTemplate):
10821078
"""
10831079
is_method = True
10841080

1085-
@classmethod
1086-
def do_class_init(cls):
1081+
def _init_once(self):
10871082
"""
1088-
Register generic method implementation.
1083+
Overriding parent definition
10891084
"""
1090-
from numba.core.imputils import lower_builtin
1091-
attr = cls._attr
1085+
attr = self._attr
10921086

1093-
@lower_builtin((cls.key, attr), cls.key, types.VarArg(types.Any))
1094-
def method_impl(context, builder, sig, args):
1095-
typ = sig.args[0]
1096-
typing_context = context.typing_context
1097-
fnty = cls._get_function_type(typing_context, typ)
1098-
sig = cls._get_signature(typing_context, fnty, sig.args, {})
1099-
call = context.get_function(fnty, sig)
1100-
# Link dependent library
1101-
context.add_linking_libs(getattr(call, 'libs', ()))
1102-
return call(builder, args)
1087+
try:
1088+
registry = self._get_target_registry()
1089+
except InternalTargetMismatchError:
1090+
# Target mismatch. Do not register attribute lookup here.
1091+
pass
1092+
else:
1093+
lower_builtin = registry.lower
1094+
1095+
@lower_builtin((self.key, attr), self.key, types.VarArg(types.Any))
1096+
def method_impl(context, builder, sig, args):
1097+
typ = sig.args[0]
1098+
typing_context = context.typing_context
1099+
fnty = self._get_function_type(typing_context, typ)
1100+
sig = self._get_signature(typing_context, fnty, sig.args, {})
1101+
call = context.get_function(fnty, sig)
1102+
# Link dependent library
1103+
context.add_linking_libs(getattr(call, 'libs', ()))
1104+
return call(builder, args)
11031105

11041106
def _resolve(self, typ, attr):
11051107
if self._attr != attr:
@@ -1132,7 +1134,8 @@ def generic(_, args, kws):
11321134

11331135
def make_overload_attribute_template(typ, attr, overload_func, inline,
11341136
prefer_literal=False,
1135-
base=_OverloadAttributeTemplate):
1137+
base=_OverloadAttributeTemplate,
1138+
**kwargs):
11361139
"""
11371140
Make a template class for attribute *attr* of *typ* overloaded by
11381141
*overload_func*.
@@ -1145,20 +1148,22 @@ def make_overload_attribute_template(typ, attr, overload_func, inline,
11451148
_inline_overloads={},
11461149
_overload_func=staticmethod(overload_func),
11471150
prefer_literal=prefer_literal,
1151+
metadata=kwargs,
11481152
)
11491153
obj = type(base)(name, (base,), dct)
11501154
return obj
11511155

11521156

11531157
def make_overload_method_template(typ, attr, overload_func, inline,
1154-
prefer_literal=False):
1158+
prefer_literal=False, **kwargs):
11551159
"""
11561160
Make a template class for method *attr* of *typ* overloaded by
11571161
*overload_func*.
11581162
"""
11591163
return make_overload_attribute_template(
11601164
typ, attr, overload_func, inline=inline,
11611165
base=_OverloadMethodTemplate, prefer_literal=prefer_literal,
1166+
**kwargs,
11621167
)
11631168

11641169

‎numba/np/arrayobj.py‎

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
from numba.cpython import slicing
3333
from numba.cpython.unsafe.tuple import tuple_setitem, build_full_slice_tuple
3434
from numba.core.overload_glue import glue_lowering
35+
from numba.core.extending import overload_classmethod
3536

3637

3738
def set_range_metadata(builder, load, lower_bound, upper_bound):
@@ -3501,10 +3502,16 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
35013502
" the maximum possible size.",)
35023503
)
35033504

3504-
align = context.get_preferred_array_alignment(arrtype.dtype)
3505-
meminfo = context.nrt.meminfo_alloc_aligned(builder, size=allocsize,
3506-
align=align)
3505+
dtype = arrtype.dtype
3506+
align_val = context.get_preferred_array_alignment(dtype)
3507+
align = context.get_constant(types.uint32, align_val)
3508+
args = (context.get_dummy_value(), allocsize, align)
35073509

3510+
mip = types.MemInfoPointer(types.voidptr)
3511+
arytypeclass = types.TypeRef(type(arrtype))
3512+
argtypes = signature(mip, arytypeclass, types.intp, types.uint32)
3513+
3514+
meminfo = context.compile_internal(builder, _call_allocator, argtypes, args)
35083515
data = context.nrt.meminfo_data(builder, meminfo)
35093516

35103517
intp_t = context.get_value_type(types.intp)
@@ -3521,6 +3528,35 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
35213528
return ary
35223529

35233530

3531+
@overload_classmethod(types.Array, "_allocate")
3532+
def _ol_array_allocate(cls, allocsize, align):
3533+
"""Implements a Numba-only default target (cpu) classmethod on the array type.
3534+
"""
3535+
def impl(cls, allocsize, align):
3536+
return intrin_alloc(allocsize, align)
3537+
return impl
3538+
3539+
3540+
def _call_allocator(arrtype, size, align):
3541+
"""Trampoline to call the intrinsic used for allocation
3542+
"""
3543+
return arrtype._allocate(size, align)
3544+
3545+
3546+
@intrinsic
3547+
def intrin_alloc(typingctx, allocsize, align):
3548+
"""Intrinsic to call into the allocator for Array
3549+
"""
3550+
def codegen(context, builder, signature, args):
3551+
[allocsize, align] = args
3552+
meminfo = context.nrt.meminfo_alloc_aligned(builder, allocsize, align)
3553+
return meminfo
3554+
3555+
mip = types.MemInfoPointer(types.voidptr) # return untyped pointer
3556+
sig = signature(mip, allocsize, align)
3557+
return sig, codegen
3558+
3559+
35243560
def _zero_fill_array(context, builder, ary):
35253561
"""
35263562
Zero-fill an array. The array must be contiguous.

0 commit comments

Comments
 (0)