def lower_dist_isend(context, builder, sig, args):
# store an int to specify data type
typ_enum = _numba_to_c_type_map[sig.args[0].dtype]
typ_arg = context.get_constant(types.int32, typ_enum)
out = make_array(sig.args[0])(context, builder, args[0])
if len(args) == 4:
cond_arg = context.get_constant(types.boolean, True)
else:
cond_arg = args[4]
call_args = [
builder.bitcast(out.data,
lir.IntType(8).as_pointer()), args[1], typ_arg,
args[2], args[3], cond_arg
]
# array, size, extra arg type for type enum
# pe, tag, cond
arg_typs = [
lir.IntType(8).as_pointer(),
lir.IntType(32),
lir.IntType(32),
lir.IntType(32),
lir.IntType(32),
lir.IntType(1)
]
fnty = lir.FunctionType(mpi_req_llvm_type, arg_typs)
fn = builder.module.get_or_insert_function(fnty, name="hpat_dist_isend")
return builder.call(fn, call_args)
def lower_dist_allgather(context, builder, sig, args):
arr_typ = sig.args[0]
val_typ = sig.args[1]
assert val_typ == arr_typ.dtype
# type enum arg
assert val_typ in _numba_to_c_type_map, "invalid allgather type"
typ_enum = _numba_to_c_type_map[val_typ]
typ_arg = context.get_constant(types.int32, typ_enum)
# size arg is 1 for now
size_arg = context.get_constant(types.int32, 1)
val_ptr = cgutils.alloca_once_value(builder, args[1])
out = make_array(sig.args[0])(context, builder, args[0])
call_args = [
builder.bitcast(out.data,
lir.IntType(8).as_pointer()), size_arg, val_ptr,
typ_arg
]
fnty = lir.FunctionType(lir.VoidType(), [
lir.IntType(8).as_pointer(),
lir.IntType(32), val_ptr.type,
lir.IntType(32)
])
fn = builder.module.get_or_insert_function(fnty, name="allgather")
builder.call(fn, call_args)
return context.get_dummy_value()
def iternext_series_array(context, builder, sig, args, result):
"""
Implementation of iternext() for the ArrayIterator type
:param context: context descriptor
:param builder: llvmlite IR Builder
:param sig: iterator signature
:param args: tuple with iterator arguments, such as instruction, operands and types
:param result: iternext result
"""
[iterty] = sig.args
[iter] = args
arrayty = iterty.array_type
if arrayty.ndim != 1:
raise NotImplementedError("iterating over %dD array" % arrayty.ndim)
iterobj = context.make_helper(builder, iterty, value=iter)
ary = make_array(arrayty)(context, builder, value=iterobj.array)
nitems, = cgutils.unpack_tuple(builder, ary.shape, count=1)
index = builder.load(iterobj.index)
is_valid = builder.icmp(lc.ICMP_SLT, index, nitems)
result.set_valid(is_valid)
with builder.if_then(is_valid):
value = _getitem_array_single_int(context, builder, iterty.yield_type,
arrayty, ary, index)
result.yield_(value)
nindex = cgutils.increment_index(builder, index)
builder.store(nindex, iterobj.index)
def from_buffer(context, builder, sig, args):
assert len(sig.args) == 1
assert len(args) == 1
[fromty] = sig.args
[val] = args
# Type inference should have prevented passing a buffer from an
# array to a pointer of the wrong type
assert fromty.dtype == sig.return_type.dtype
ary = arrayobj.make_array(fromty)(context, builder, val)
return ary.data
def codegen(context, builder, sig, args):
(iterty, ) = sig.args
(value, ) = args
intp_t = context.get_value_type(types.intp)
iterobj = context.make_helper(builder, iterty, value=value)
arrayty = iterty.array_type
ary = make_array(arrayty)(context, builder, value=iterobj.array)
shape = cgutils.unpack_tuple(builder, ary.shape)
# array iterates along the outer dimension
return impl_ret_untracked(context, builder, intp_t, shape[0])
def impl_myarray(context, builder, sig, args):
from numba.np.arrayobj import make_array, populate_array
srcaryty = sig.args[-1]
shape, dtype, buf = args
srcary = make_array(srcaryty)(context, builder, value=buf)
# Copy source array and remove the parent field to avoid boxer re-using
# the original ndarray instance.
retary = make_array(sig.return_type)(context, builder)
populate_array(retary,
data=srcary.data,
shape=srcary.shape,
strides=srcary.strides,
itemsize=srcary.itemsize,
meminfo=srcary.meminfo)
ret = retary._getvalue()
context.nrt.incref(builder, sig.return_type, ret)
return ret
def pq_read_lower(context, builder, sig, args):
fnty = lir.FunctionType(lir.IntType(64),
[lir.IntType(8).as_pointer(), lir.IntType(64),
lir.IntType(8).as_pointer()], lir.IntType(32))
out_array = make_array(sig.args[2])(context, builder, args[2])
fn = builder.module.get_or_insert_function(fnty, name="pq_read")
return builder.call(fn, [args[0], args[1],
builder.bitcast(
out_array.data, lir.IntType(8).as_pointer()),
args[3]])
def func_impl(context, builder, sig, args):
"""
array[a] = scalar_or_array
array[a,..,b] = scalar_or_array
"""
aryty, idxty, valty = sig.args
ary, idx, val = args
if isinstance(idxty, types.BaseTuple):
index_types = idxty.types
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
else:
index_types = (idxty, )
indices = (idx, )
ary = make_array(aryty)(context, builder, ary)
# First try basic indexing to see if a single array location is denoted.
index_types, indices = normalize_indices(context, builder,
index_types, indices)
dataptr, shapes, _strides = basic_indexing(
context,
builder,
aryty,
ary,
index_types,
indices,
boundscheck=context.enable_boundscheck,
)
if shapes:
raise NotImplementedError("Complex shapes are not supported")
# Store source value the given location
val = context.cast(builder, val, valty, aryty.dtype)
operation = None
if isinstance(aryty.dtype, types.Integer) and aryty.dtype.signed:
operation = iop
elif isinstance(aryty.dtype,
types.Integer) and not aryty.dtype.signed:
operation = uop
elif isinstance(aryty.dtype, types.Float):
operation = fop
if operation is None:
raise TypeError("Atomic operation not supported on " +
str(aryty))
return _atomic_rmw(context, builder, operation, aryty, val,
dataptr)
def lower_dist_arr_reduce(context, builder, sig, args):
op_typ = args[1].type
# store an int to specify data type
typ_enum = _numba_to_c_type_map[sig.args[0].dtype]
typ_arg = cgutils.alloca_once_value(
builder, lir.Constant(lir.IntType(32), typ_enum))
ndims = sig.args[0].ndim
out = make_array(sig.args[0])(context, builder, args[0])
# store size vars array struct to pointer
size_ptr = cgutils.alloca_once(builder, out.shape.type)
builder.store(out.shape, size_ptr)
size_arg = builder.bitcast(size_ptr, lir.IntType(64).as_pointer())
ndim_arg = cgutils.alloca_once_value(
builder, lir.Constant(lir.IntType(32), sig.args[0].ndim))
call_args = [
builder.bitcast(out.data,
lir.IntType(8).as_pointer()), size_arg,
builder.load(ndim_arg), args[1],
builder.load(typ_arg)
]
# array, shape, ndim, extra last arg type for type enum
arg_typs = [
lir.IntType(8).as_pointer(),
lir.IntType(64).as_pointer(),
lir.IntType(32), op_typ,
lir.IntType(32)
]
fnty = lir.FunctionType(lir.IntType(32), arg_typs)
fn = builder.module.get_or_insert_function(fnty,
name="hpat_dist_arr_reduce")
builder.call(fn, call_args)
res = out._getvalue()
return impl_ret_borrowed(context, builder, sig.return_type, res)
def make_array(self, typ):
from numba.np import arrayobj
return arrayobj.make_array(typ)
def make_array(self, typ):
return arrayobj.make_array(typ)
