Skip to content

Commit db28199

Browse files
Refactor carray/farray to use @overload
1 parent c619212 commit db28199

2 files changed

Lines changed: 70 additions & 52 deletions

File tree

‎numba/core/typing/npydecl.py‎

Lines changed: 0 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1139,50 +1139,3 @@ def generic(self, args, kws):
11391139
return None
11401140

11411141
return signature(retty, *args)
1142-
1143-
# -----------------------------------------------------------------------------
1144-
# Numba helpers
1145-
1146-
@glue_typing(carray)
1147-
class NumbaCArray(CallableTemplate):
1148-
layout = 'C'
1149-
1150-
def generic(self):
1151-
func_name = self.key.__name__
1152-
1153-
def typer(ptr, shape, dtype=None):
1154-
if ptr is types.voidptr:
1155-
ptr_dtype = None
1156-
elif isinstance(ptr, types.CPointer):
1157-
ptr_dtype = ptr.dtype
1158-
else:
1159-
raise NumbaTypeError("%s(): pointer argument expected, got '%s'"
1160-
% (func_name, ptr))
1161-
1162-
if dtype is None:
1163-
if ptr_dtype is None:
1164-
raise NumbaTypeError("%s(): explicit dtype required for void* argument"
1165-
% (func_name,))
1166-
dtype = ptr_dtype
1167-
elif isinstance(dtype, types.DTypeSpec):
1168-
dtype = dtype.dtype
1169-
if ptr_dtype is not None and dtype != ptr_dtype:
1170-
raise NumbaTypeError("%s(): mismatching dtype '%s' for pointer type '%s'"
1171-
% (func_name, dtype, ptr))
1172-
else:
1173-
raise NumbaTypeError("%s(): invalid dtype spec '%s'"
1174-
% (func_name, dtype))
1175-
1176-
ndim = parse_shape(shape)
1177-
if ndim is None:
1178-
raise NumbaTypeError("%s(): invalid shape '%s'"
1179-
% (func_name, shape))
1180-
1181-
return types.Array(dtype, ndim, self.layout)
1182-
1183-
return typer
1184-
1185-
1186-
@glue_typing(farray)
1187-
class NumbaFArray(NumbaCArray):
1188-
layout = 'F'

‎numba/np/arrayobj.py‎

Lines changed: 70 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3967,7 +3967,8 @@ def _empty_nd_impl(context, builder, arrtype, shapes):
39673967

39683968
@overload_classmethod(types.Array, "_allocate")
39693969
def _ol_array_allocate(cls, allocsize, align):
3970-
"""Implements a Numba-only default target (cpu) classmethod on the array type.
3970+
"""Implements a Numba-only default target (cpu) classmethod on the array
3971+
type.
39713972
"""
39723973
def impl(cls, allocsize, align):
39733974
return intrin_alloc(allocsize, align)
@@ -4805,10 +4806,74 @@ def np_frombuffer(context, builder, sig, args):
48054806
return impl_ret_borrowed(context, builder, sig.return_type, res)
48064807

48074808

4808-
@glue_lowering(carray, types.Any, types.Any)
4809-
@glue_lowering(carray, types.Any, types.Any, types.DTypeSpec)
4810-
@glue_lowering(farray, types.Any, types.Any)
4811-
@glue_lowering(farray, types.Any, types.Any, types.DTypeSpec)
4809+
@overload(carray)
4810+
def impl_carray(ptr, shape, dtype=None):
4811+
if is_nonelike(dtype):
4812+
intrinsic_cfarray = get_cfarray_intrinsic('C', None)
4813+
4814+
def impl(ptr, shape, dtype=None):
4815+
return intrinsic_cfarray(ptr, shape)
4816+
return impl
4817+
elif isinstance(dtype, types.DTypeSpec):
4818+
intrinsic_cfarray = get_cfarray_intrinsic('C', dtype)
4819+
4820+
def impl(ptr, shape, dtype=None):
4821+
return intrinsic_cfarray(ptr, shape)
4822+
return impl
4823+
4824+
4825+
@overload(farray)
4826+
def impl_farray(ptr, shape, dtype=None):
4827+
if is_nonelike(dtype):
4828+
intrinsic_cfarray = get_cfarray_intrinsic('F', None)
4829+
4830+
def impl(ptr, shape, dtype=None):
4831+
return intrinsic_cfarray(ptr, shape)
4832+
return impl
4833+
elif isinstance(dtype, types.DTypeSpec):
4834+
intrinsic_cfarray = get_cfarray_intrinsic('F', dtype)
4835+
4836+
def impl(ptr, shape, dtype=None):
4837+
return intrinsic_cfarray(ptr, shape)
4838+
return impl
4839+
4840+
4841+
def get_cfarray_intrinsic(layout, dtype_):
4842+
@intrinsic
4843+
def intrinsic_cfarray(typingctx, ptr, shape):
4844+
if ptr is types.voidptr:
4845+
ptr_dtype = None
4846+
elif isinstance(ptr, types.CPointer):
4847+
ptr_dtype = ptr.dtype
4848+
else:
4849+
msg = f"pointer argument expected, got '{ptr}'"
4850+
raise errors.NumbaTypeError(msg)
4851+
4852+
if dtype_ is None:
4853+
if ptr_dtype is None:
4854+
msg = "explicit dtype required for void* argument"
4855+
raise errors.NumbaTypeError(msg)
4856+
dtype = ptr_dtype
4857+
elif isinstance(dtype_, types.DTypeSpec):
4858+
dtype = dtype_.dtype
4859+
if ptr_dtype is not None and dtype != ptr_dtype:
4860+
msg = f"mismatching dtype '{dtype}' for pointer type '{ptr}'"
4861+
raise errors.NumbaTypeError(msg)
4862+
else:
4863+
msg = f"invalid dtype spec '{dtype_}'"
4864+
raise errors.NumbaTypeError(msg)
4865+
4866+
ndim = ty_parse_shape(shape)
4867+
if ndim is None:
4868+
msg = f"invalid shape '{shape}'"
4869+
raise errors.NumbaTypeError(msg)
4870+
4871+
retty = types.Array(dtype, ndim, layout)
4872+
sig = signature(retty, ptr, shape)
4873+
return sig, np_cfarray
4874+
return intrinsic_cfarray
4875+
4876+
48124877
def np_cfarray(context, builder, sig, args):
48134878
"""
48144879
numba.numpy_support.carray(...) and

0 commit comments

Comments
 (0)