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)