def lower_sort(context, builder, sig, args):
#
key_arr = make_array(sig.args[0])(context, builder, args[0])
key_data = builder.bitcast(key_arr.data, lir.IntType(8).as_pointer())
# XXX: assuming key arr is 1D
assert key_arr.shape.type.count == 1
arr_len = builder.extract_value(key_arr.shape, 0)
num_other_cols = len(args) - 1
# build array of other column arrays arg
other_arrs = cgutils.alloca_once(builder,
lir.IntType(8).as_pointer(),
num_other_cols)
for i in range(num_other_cols):
ptr = cgutils.gep_inbounds(builder, other_arrs, i)
arr = make_array(sig.args[i + 1])(context, builder, args[i + 1])
arr_data = builder.bitcast(arr.data, lir.IntType(8).as_pointer())
builder.store(arr_data, ptr)
call_args = [
key_data, arr_len, other_arrs,
lir.Constant(lir.IntType(64), num_other_cols)
]
fnty = lir.FunctionType(lir.VoidType(), [
lir.IntType(8).as_pointer(),
lir.IntType(64),
lir.IntType(8).as_pointer().as_pointer(),
lir.IntType(64)
])
fn = builder.module.get_or_insert_function(fnty, name="timsort")
builder.call(fn, call_args)
return lir.Constant(lir.IntType(32), 0)
def lower_shuffle_arr(context, builder, sig, args):
args[0] = make_array(sig.args[0])(context, builder, args[0]).data
args[1] = make_array(sig.args[1])(context, builder, args[1]).data
args[2] = make_array(sig.args[2])(context, builder, args[2]).data
args[3] = make_array(sig.args[3])(context, builder, args[3]).data
sig.args = (c_buffer_type, c_buffer_type, c_buffer_type, c_buffer_type, *sig.args[4:])
return lower_shuffle(context, builder, sig, 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 _h5_typ_table, "invalid allgather type"
typ_enum = _h5_typ_table[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 lower_dist_arr_reduce(context, builder, sig, args):
op_typ = args[1].type
# store an int to specify data type
typ_enum = _h5_typ_table[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 lower_dist_isend(context, builder, sig, args):
# store an int to specify data type
typ_enum = hpat.pio_lower._h5_typ_table[sig.args[0].dtype]
typ_arg = cgutils.alloca_once_value(
builder, lir.Constant(lir.IntType(32), typ_enum))
out = make_array(sig.args[0])(context, builder, args[0])
call_args = [
builder.bitcast(out.data,
lir.IntType(8).as_pointer()), args[1],
builder.load(typ_arg), args[2], args[3], args[4]
]
# 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(lir.IntType(32), arg_typs)
fn = builder.module.get_or_insert_function(fnty, name="hpat_dist_isend")
return builder.call(fn, call_args)
def set_df_datetime_date_lower(context, builder, sig, args):
#
col_name = sig.args[1].literal_value
data_arr = make_array(sig.args[2])(context, builder, args[2])
num_elems = builder.extract_value(data_arr.shape, 0)
pyapi = context.get_python_api(builder)
gil_state = pyapi.gil_ensure() # acquire GIL
dt_class = pyapi.unserialize(pyapi.serialize_object(datetime.date))
fnty = lir.FunctionType(lir.IntType(8).as_pointer(), [lir.IntType(64).as_pointer(), lir.IntType(64),
lir.IntType(8).as_pointer()])
fn = builder.module.get_or_insert_function(fnty, name="np_datetime_date_array_from_packed_ints")
py_arr = builder.call(fn, [data_arr.data, num_elems, dt_class])
# get column as string obj
cstr = context.insert_const_string(builder.module, col_name)
cstr_obj = pyapi.string_from_string(cstr)
# set column array
pyapi.object_setitem(args[0], cstr_obj, py_arr)
pyapi.decref(py_arr)
pyapi.decref(cstr_obj)
pyapi.gil_release(gil_state) # release GIL
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_array1d(context, builder, arrayty, ary, index,
wraparound=False)
result.yield_(value)
nindex = cgutils.increment_index(builder, index)
builder.store(nindex, iterobj.index)
def h5_read(context, builder, sig, args):
# extra last arg type for type enum
arg_typs = [h5file_lir_type, lir.IntType(32),
lir.IntType(64).as_pointer(), lir.IntType(64).as_pointer(),
lir.IntType(64), lir.IntType(8).as_pointer(), lir.IntType(32)]
fnty = lir.FunctionType(lir.IntType(32), arg_typs)
fn = builder.module.get_or_insert_function(fnty, name="hpat_h5_read")
out = make_array(sig.args[5])(context, builder, args[5])
# store size vars array struct to pointer
count_ptr = cgutils.alloca_once(builder, args[2].type)
builder.store(args[2], count_ptr)
size_ptr = cgutils.alloca_once(builder, args[3].type)
builder.store(args[3], size_ptr)
# store an int to specify data type
typ_enum = _numba_to_c_type_map[sig.args[5].dtype]
typ_arg = cgutils.alloca_once_value(
builder, lir.Constant(lir.IntType(32), typ_enum))
call_args = [args[0], args[1],
builder.bitcast(count_ptr, lir.IntType(64).as_pointer()),
builder.bitcast(size_ptr, lir.IntType(
64).as_pointer()), args[4],
builder.bitcast(out.data, lir.IntType(8).as_pointer()),
builder.load(typ_arg)]
return builder.call(fn, call_args)
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_quantile(context, builder, sig, args):
# store an int to specify data type
typ_enum = _h5_typ_table[sig.args[0].dtype]
typ_arg = cgutils.alloca_once_value(builder, lir.Constant(lir.IntType(32), typ_enum))
assert sig.args[0].ndim == 1
arr = make_array(sig.args[0])(context, builder, args[0])
local_size = builder.extract_value(arr.shape, 0)
if len(args) == 3:
total_size = args[2]
else:
# sequential case
total_size = local_size
call_args = [builder.bitcast(arr.data, lir.IntType(8).as_pointer()),
local_size, total_size, args[1], builder.load(typ_arg)]
# array, size, total_size, quantile, type enum
arg_typs = [lir.IntType(8).as_pointer(), lir.IntType(64), lir.IntType(64),
lir.DoubleType(), lir.IntType(32)]
fnty = lir.FunctionType(lir.DoubleType(), arg_typs)
fn = builder.module.get_or_insert_function(fnty, name="quantile_parallel")
return builder.call(fn, call_args)
def lower_get_sendrecv_counts(context, builder, sig, args):
# prepare buffer args
pointer_to_cbuffer_typ = lir.IntType(8).as_pointer().as_pointer()
send_counts = cgutils.alloca_once(builder, lir.IntType(8).as_pointer())
recv_counts = cgutils.alloca_once(builder, lir.IntType(8).as_pointer())
send_disp = cgutils.alloca_once(builder, lir.IntType(8).as_pointer())
recv_disp = cgutils.alloca_once(builder, lir.IntType(8).as_pointer())
# prepare key array args
key_arr = make_array(sig.args[0])(context, builder, args[0])
# XXX: assuming key arr is 1D
assert key_arr.shape.type.count == 1
arr_len = builder.extract_value(key_arr.shape, 0)
# TODO: extend to other key types
assert sig.args[0].dtype == types.intp
key_typ_enum = _h5_typ_table[sig.args[0].dtype]
key_typ_arg = builder.load(cgutils.alloca_once_value(builder,
lir.Constant(lir.IntType(32), key_typ_enum)))
key_arr_data = builder.bitcast(key_arr.data, lir.IntType(8).as_pointer())
call_args = [send_counts, recv_counts, send_disp, recv_disp, arr_len,
key_typ_arg, key_arr_data]
fnty = lir.FunctionType(lir.IntType(64), [pointer_to_cbuffer_typ] * 4
+ [lir.IntType(64), lir.IntType(32), lir.IntType(8).as_pointer()])
fn = builder.module.get_or_insert_function(fnty,
name="get_join_sendrecv_counts")
total_size = builder.call(fn, call_args)
items = [builder.load(send_counts), builder.load(recv_counts),
builder.load(send_disp), builder.load(recv_disp), total_size]
out_tuple_typ = types.Tuple([c_buffer_type, c_buffer_type, c_buffer_type,
c_buffer_type, types.intp])
return context.make_tuple(builder, out_tuple_typ, items)
def lower_read_images_inner(context, builder, sig, args):
bag = args[1]
out = make_array(sig.args[0])(context, builder, args[0])
fnty = lir.FunctionType(lir.IntType(32),
[lir.IntType(8).as_pointer(),
lir.IntType(8).as_pointer()])
fn = builder.module.get_or_insert_function(fnty, name="read_images")
return builder.call(fn, [builder.bitcast(out.data, lir.IntType(8).as_pointer()), bag])
def gen_alltoallv(context, builder, arr_typ, send_arg, recv_arg, send_counts,
recv_counts, send_disp, recv_disp):
#
typ_enum = _h5_typ_table[arr_typ.dtype]
typ_arg = builder.load(cgutils.alloca_once_value(builder,
lir.Constant(lir.IntType(32), typ_enum)))
send_data = make_array(arr_typ)(context, builder, send_arg).data
recv_data = make_array(arr_typ)(context, builder, recv_arg).data
send_data = builder.bitcast(send_data, lir.IntType(8).as_pointer())
recv_data = builder.bitcast(recv_data, lir.IntType(8).as_pointer())
call_args = [send_data, recv_data, send_counts,
recv_counts, send_disp, recv_disp, typ_arg]
fnty = lir.FunctionType(lir.VoidType(), [lir.IntType(
8).as_pointer()] * 6 + [lir.IntType(32)])
fn = builder.module.get_or_insert_function(fnty, name="c_alltoallv")
builder.call(fn, call_args)
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 make_array(self, typ):
return arrayobj.make_array(typ)
def codegen(context, builder, sig, args):
buff_arr, ind, str, len_str = args
buff_arr = make_array(sig.args[0])(context, builder, buff_arr)
ptr = builder.gep(buff_arr.data, [ind])
cgutils.raw_memcpy(builder, ptr, str, len_str, 1)
return context.get_dummy_value()
