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()