Skip to content
Permalink
master
Switch branches/tags
Go to file
 
 
Cannot retrieve contributors at this time
"""
Implementation of operations on Array objects and objects supporting
the buffer protocol.
"""
import functools
import math
import operator
from llvmlite import ir
import llvmlite.llvmpy.core as lc
from llvmlite.llvmpy.core import Constant
import numpy as np
from numba import pndindex, literal_unroll
from numba.core import types, utils, typing, errors, cgutils, extending
from numba.np.numpy_support import (as_dtype, carray, farray, is_contiguous,
is_fortran, check_is_integer)
from numba.np.numpy_support import type_can_asarray, is_nonelike
from numba.core.imputils import (lower_builtin, lower_getattr,
lower_getattr_generic,
lower_setattr_generic,
lower_cast, lower_constant,
iternext_impl, impl_ret_borrowed,
impl_ret_new_ref, impl_ret_untracked,
RefType)
from numba.core.typing import signature
from numba.core.extending import (register_jitable, overload, overload_method,
intrinsic)
from numba.misc import quicksort, mergesort
from numba.cpython import slicing
from numba.cpython.unsafe.tuple import tuple_setitem, build_full_slice_tuple
from numba.core.overload_glue import glue_lowering
from numba.core.extending import overload_classmethod
def set_range_metadata(builder, load, lower_bound, upper_bound):
"""
Set the "range" metadata on a load instruction.
Note the interval is in the form [lower_bound, upper_bound).
"""
range_operands = [Constant.int(load.type, lower_bound),
Constant.int(load.type, upper_bound)]
md = builder.module.add_metadata(range_operands)
load.set_metadata("range", md)
def mark_positive(builder, load):
"""
Mark the result of a load instruction as positive (or zero).
"""
upper_bound = (1 << (load.type.width - 1)) - 1
set_range_metadata(builder, load, 0, upper_bound)
def make_array(array_type):
"""
Return the Structure representation of the given *array_type*
(an instance of types.ArrayCompatible).
Note this does not call __array_wrap__ in case a new array structure
is being created (rather than populated).
"""
real_array_type = array_type.as_array
base = cgutils.create_struct_proxy(real_array_type)
ndim = real_array_type.ndim
class ArrayStruct(base):
def _make_refs(self, ref):
sig = signature(real_array_type, array_type)
try:
array_impl = self._context.get_function('__array__', sig)
except NotImplementedError:
return super(ArrayStruct, self)._make_refs(ref)
# Return a wrapped structure and its unwrapped reference
datamodel = self._context.data_model_manager[array_type]
be_type = self._get_be_type(datamodel)
if ref is None:
outer_ref = cgutils.alloca_once(self._builder, be_type,
zfill=True)
else:
outer_ref = ref
# NOTE: __array__ is called with a pointer and expects a pointer
# in return!
ref = array_impl(self._builder, (outer_ref,))
return outer_ref, ref
@property
def shape(self):
"""
Override .shape to inform LLVM that its elements are all positive.
"""
builder = self._builder
if ndim == 0:
return base.__getattr__(self, "shape")
# Unfortunately, we can't use llvm.assume as its presence can
# seriously pessimize performance,
# *and* the range metadata currently isn't improving anything here,
# see https://llvm.org/bugs/show_bug.cgi?id=23848 !
ptr = self._get_ptr_by_name("shape")
dims = []
for i in range(ndim):
dimptr = cgutils.gep_inbounds(builder, ptr, 0, i)
load = builder.load(dimptr)
dims.append(load)
mark_positive(builder, load)
return cgutils.pack_array(builder, dims)
return ArrayStruct
def get_itemsize(context, array_type):
"""
Return the item size for the given array or buffer type.
"""
llty = context.get_data_type(array_type.dtype)
return context.get_abi_sizeof(llty)
def load_item(context, builder, arrayty, ptr):
"""
Load the item at the given array pointer.
"""
align = None if arrayty.aligned else 1
return context.unpack_value(builder, arrayty.dtype, ptr,
align=align)
def store_item(context, builder, arrayty, val, ptr):
"""
Store the item at the given array pointer.
"""
align = None if arrayty.aligned else 1
return context.pack_value(builder, arrayty.dtype, val, ptr, align=align)
def fix_integer_index(context, builder, idxty, idx, size):
"""
Fix the integer index' type and value for the given dimension size.
"""
if idxty.signed:
ind = context.cast(builder, idx, idxty, types.intp)
ind = slicing.fix_index(builder, ind, size)
else:
ind = context.cast(builder, idx, idxty, types.uintp)
return ind
def normalize_index(context, builder, idxty, idx):
"""
Normalize the index type and value. 0-d arrays are converted to scalars.
"""
if isinstance(idxty, types.Array) and idxty.ndim == 0:
assert isinstance(idxty.dtype, types.Integer)
idxary = make_array(idxty)(context, builder, idx)
idxval = load_item(context, builder, idxty, idxary.data)
return idxty.dtype, idxval
else:
return idxty, idx
def normalize_indices(context, builder, index_types, indices):
"""
Same as normalize_index(), but operating on sequences of
index types and values.
"""
if len(indices):
index_types, indices = zip(*[normalize_index(context, builder, idxty,
idx)
for idxty, idx in zip(index_types, indices)
])
return index_types, indices
def populate_array(array, data, shape, strides, itemsize, meminfo,
parent=None):
"""
Helper function for populating array structures.
This avoids forgetting to set fields.
*shape* and *strides* can be Python tuples or LLVM arrays.
"""
context = array._context
builder = array._builder
datamodel = array._datamodel
required_fields = set(datamodel._fields)
if meminfo is None:
meminfo = Constant.null(context.get_value_type(
datamodel.get_type('meminfo')))
intp_t = context.get_value_type(types.intp)
if isinstance(shape, (tuple, list)):
shape = cgutils.pack_array(builder, shape, intp_t)
if isinstance(strides, (tuple, list)):
strides = cgutils.pack_array(builder, strides, intp_t)
if isinstance(itemsize, int):
itemsize = intp_t(itemsize)
attrs = dict(shape=shape,
strides=strides,
data=data,
itemsize=itemsize,
meminfo=meminfo,)
# Set `parent` attribute
if parent is None:
attrs['parent'] = Constant.null(context.get_value_type(
datamodel.get_type('parent')))
else:
attrs['parent'] = parent
# Calc num of items from shape
nitems = context.get_constant(types.intp, 1)
unpacked_shape = cgutils.unpack_tuple(builder, shape, shape.type.count)
# (note empty shape => 0d array therefore nitems = 1)
for axlen in unpacked_shape:
nitems = builder.mul(nitems, axlen, flags=['nsw'])
attrs['nitems'] = nitems
# Make sure that we have all the fields
got_fields = set(attrs.keys())
if got_fields != required_fields:
raise ValueError("missing {0}".format(required_fields - got_fields))
# Set field value
for k, v in attrs.items():
setattr(array, k, v)
return array
def update_array_info(aryty, array):
"""
Update some auxiliary information in *array* after some of its fields
were changed. `itemsize` and `nitems` are updated.
"""
context = array._context
builder = array._builder
# Calc num of items from shape
nitems = context.get_constant(types.intp, 1)
unpacked_shape = cgutils.unpack_tuple(builder, array.shape, aryty.ndim)
for axlen in unpacked_shape:
nitems = builder.mul(nitems, axlen, flags=['nsw'])
array.nitems = nitems
array.itemsize = context.get_constant(types.intp,
get_itemsize(context, aryty))
@lower_builtin('getiter', types.Buffer)
def getiter_array(context, builder, sig, args):
[arrayty] = sig.args
[array] = args
iterobj = context.make_helper(builder, sig.return_type)
zero = context.get_constant(types.intp, 0)
indexptr = cgutils.alloca_once_value(builder, zero)
iterobj.index = indexptr
iterobj.array = array
# Incref array
if context.enable_nrt:
context.nrt.incref(builder, arrayty, array)
res = iterobj._getvalue()
# Note: a decref on the iterator will dereference all internal MemInfo*
out = impl_ret_new_ref(context, builder, sig.return_type, res)
return out
def _getitem_array_single_int(context, builder, return_type, aryty, ary, idx):
""" Evaluate `ary[idx]`, where idx is a single int. """
# optimized form of _getitem_array_generic
shapes = cgutils.unpack_tuple(builder, ary.shape, count=aryty.ndim)
strides = cgutils.unpack_tuple(builder, ary.strides, count=aryty.ndim)
offset = builder.mul(strides[0], idx)
dataptr = cgutils.pointer_add(builder, ary.data, offset)
view_shapes = shapes[1:]
view_strides = strides[1:]
if isinstance(return_type, types.Buffer):
# Build array view
retary = make_view(context, builder, aryty, ary, return_type,
dataptr, view_shapes, view_strides)
return retary._getvalue()
else:
# Load scalar from 0-d result
assert not view_shapes
return load_item(context, builder, aryty, dataptr)
@lower_builtin('iternext', types.ArrayIterator)
@iternext_impl(RefType.BORROWED)
def iternext_array(context, builder, sig, args, result):
[iterty] = sig.args
[iter] = args
arrayty = iterty.array_type
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)
# ------------------------------------------------------------------------------
# Basic indexing (with integers and slices only)
def basic_indexing(context, builder, aryty, ary, index_types, indices,
boundscheck=None):
"""
Perform basic indexing on the given array.
A (data pointer, shapes, strides) tuple is returned describing
the corresponding view.
"""
zero = context.get_constant(types.intp, 0)
shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
strides = cgutils.unpack_tuple(builder, ary.strides, aryty.ndim)
output_indices = []
output_shapes = []
output_strides = []
ax = 0
for indexval, idxty in zip(indices, index_types):
if idxty is types.ellipsis:
# Fill up missing dimensions at the middle
n_missing = aryty.ndim - len(indices) + 1
for i in range(n_missing):
output_indices.append(zero)
output_shapes.append(shapes[ax])
output_strides.append(strides[ax])
ax += 1
continue
# Regular index value
if isinstance(idxty, types.SliceType):
slice = context.make_helper(builder, idxty, value=indexval)
slicing.guard_invalid_slice(context, builder, idxty, slice)
slicing.fix_slice(builder, slice, shapes[ax])
output_indices.append(slice.start)
sh = slicing.get_slice_length(builder, slice)
st = slicing.fix_stride(builder, slice, strides[ax])
output_shapes.append(sh)
output_strides.append(st)
elif isinstance(idxty, types.Integer):
ind = fix_integer_index(context, builder, idxty, indexval,
shapes[ax])
if boundscheck:
cgutils.do_boundscheck(context, builder, ind, shapes[ax], ax)
output_indices.append(ind)
else:
raise NotImplementedError("unexpected index type: %s" % (idxty,))
ax += 1
# Fill up missing dimensions at the end
assert ax <= aryty.ndim
while ax < aryty.ndim:
output_shapes.append(shapes[ax])
output_strides.append(strides[ax])
ax += 1
# No need to check wraparound, as negative indices were already
# fixed in the loop above.
dataptr = cgutils.get_item_pointer(context, builder, aryty, ary,
output_indices,
wraparound=False, boundscheck=False)
return (dataptr, output_shapes, output_strides)
def make_view(context, builder, aryty, ary, return_type,
data, shapes, strides):
"""
Build a view over the given array with the given parameters.
"""
retary = make_array(return_type)(context, builder)
populate_array(retary,
data=data,
shape=shapes,
strides=strides,
itemsize=ary.itemsize,
meminfo=ary.meminfo,
parent=ary.parent)
return retary
def _getitem_array_generic(context, builder, return_type, aryty, ary,
index_types, indices):
"""
Return the result of indexing *ary* with the given *indices*,
returning either a scalar or a view.
"""
dataptr, view_shapes, view_strides = \
basic_indexing(context, builder, aryty, ary, index_types, indices,
boundscheck=context.enable_boundscheck)
if isinstance(return_type, types.Buffer):
# Build array view
retary = make_view(context, builder, aryty, ary, return_type,
dataptr, view_shapes, view_strides)
return retary._getvalue()
else:
# Load scalar from 0-d result
assert not view_shapes
return load_item(context, builder, aryty, dataptr)
@lower_builtin(operator.getitem, types.Buffer, types.Integer)
@lower_builtin(operator.getitem, types.Buffer, types.SliceType)
def getitem_arraynd_intp(context, builder, sig, args):
"""
Basic indexing with an integer or a slice.
"""
aryty, idxty = sig.args
ary, idx = args
assert aryty.ndim >= 1
ary = make_array(aryty)(context, builder, ary)
res = _getitem_array_generic(context, builder, sig.return_type,
aryty, ary, (idxty,), (idx,))
return impl_ret_borrowed(context, builder, sig.return_type, res)
@lower_builtin(operator.getitem, types.Buffer, types.BaseTuple)
def getitem_array_tuple(context, builder, sig, args):
"""
Basic or advanced indexing with a tuple.
"""
aryty, tupty = sig.args
ary, tup = args
ary = make_array(aryty)(context, builder, ary)
index_types = tupty.types
indices = cgutils.unpack_tuple(builder, tup, count=len(tupty))
index_types, indices = normalize_indices(context, builder,
index_types, indices)
if any(isinstance(ty, types.Array) for ty in index_types):
# Advanced indexing
return fancy_getitem(context, builder, sig, args,
aryty, ary, index_types, indices)
res = _getitem_array_generic(context, builder, sig.return_type,
aryty, ary, index_types, indices)
return impl_ret_borrowed(context, builder, sig.return_type, res)
@lower_builtin(operator.setitem, types.Buffer, types.Any, types.Any)
def setitem_array(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)
try:
dataptr, shapes, strides = \
basic_indexing(context, builder, aryty, ary, index_types, indices,
boundscheck=context.enable_boundscheck)
except NotImplementedError:
use_fancy_indexing = True
else:
use_fancy_indexing = bool(shapes)
if use_fancy_indexing:
# Index describes a non-trivial view => use generic slice assignment
# (NOTE: this also handles scalar broadcasting)
return fancy_setslice(context, builder, sig, args,
index_types, indices)
# Store source value the given location
val = context.cast(builder, val, valty, aryty.dtype)
store_item(context, builder, aryty, val, dataptr)
@lower_builtin(len, types.Buffer)
def array_len(context, builder, sig, args):
(aryty,) = sig.args
(ary,) = args
arystty = make_array(aryty)
ary = arystty(context, builder, ary)
shapeary = ary.shape
res = builder.extract_value(shapeary, 0)
return impl_ret_untracked(context, builder, sig.return_type, res)
@lower_builtin("array.item", types.Array)
def array_item(context, builder, sig, args):
aryty, = sig.args
ary, = args
ary = make_array(aryty)(context, builder, ary)
nitems = ary.nitems
with builder.if_then(builder.icmp_signed('!=', nitems, nitems.type(1)),
likely=False):
msg = "item(): can only convert an array of size 1 to a Python scalar"
context.call_conv.return_user_exc(builder, ValueError, (msg,))
return load_item(context, builder, aryty, ary.data)
@lower_builtin("array.itemset", types.Array, types.Any)
def array_itemset(context, builder, sig, args):
aryty, valty = sig.args
ary, val = args
assert valty == aryty.dtype
ary = make_array(aryty)(context, builder, ary)
nitems = ary.nitems
with builder.if_then(builder.icmp_signed('!=', nitems, nitems.type(1)),
likely=False):
msg = "itemset(): can only write to an array of size 1"
context.call_conv.return_user_exc(builder, ValueError, (msg,))
store_item(context, builder, aryty, val, ary.data)
return context.get_dummy_value()
# ------------------------------------------------------------------------------
# Advanced / fancy indexing
class Indexer(object):
"""
Generic indexer interface, for generating indices over a fancy indexed
array on a single dimension.
"""
def prepare(self):
"""
Prepare the indexer by initializing any required variables, basic
blocks...
"""
raise NotImplementedError
def get_size(self):
"""
Return this dimension's size as an integer.
"""
raise NotImplementedError
def get_shape(self):
"""
Return this dimension's shape as a tuple.
"""
raise NotImplementedError
def get_index_bounds(self):
"""
Return a half-open [lower, upper) range of indices this dimension
is guaranteed not to step out of.
"""
raise NotImplementedError
def loop_head(self):
"""
Start indexation loop. Return a (index, count) tuple.
*index* is an integer LLVM value representing the index over this
dimension.
*count* is either an integer LLVM value representing the current
iteration count, or None if this dimension should be omitted from
the indexation result.
"""
raise NotImplementedError
def loop_tail(self):
"""
Finish indexation loop.
"""
raise NotImplementedError
class EntireIndexer(Indexer):
"""
Compute indices along an entire array dimension.
"""
def __init__(self, context, builder, aryty, ary, dim):
self.context = context
self.builder = builder
self.aryty = aryty
self.ary = ary
self.dim = dim
self.ll_intp = self.context.get_value_type(types.intp)
def prepare(self):
builder = self.builder
self.size = builder.extract_value(self.ary.shape, self.dim)
self.index = cgutils.alloca_once(builder, self.ll_intp)
self.bb_start = builder.append_basic_block()
self.bb_end = builder.append_basic_block()
def get_size(self):
return self.size
def get_shape(self):
return (self.size,)
def get_index_bounds(self):
# [0, size)
return (self.ll_intp(0), self.size)
def loop_head(self):
builder = self.builder
# Initialize loop variable
self.builder.store(Constant.int(self.ll_intp, 0), self.index)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_start)
cur_index = builder.load(self.index)
with builder.if_then(builder.icmp_signed('>=', cur_index, self.size),
likely=False):
builder.branch(self.bb_end)
return cur_index, cur_index
def loop_tail(self):
builder = self.builder
next_index = cgutils.increment_index(builder, builder.load(self.index))
builder.store(next_index, self.index)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_end)
class IntegerIndexer(Indexer):
"""
Compute indices from a single integer.
"""
def __init__(self, context, builder, idx):
self.context = context
self.builder = builder
self.idx = idx
self.ll_intp = self.context.get_value_type(types.intp)
def prepare(self):
pass
def get_size(self):
return Constant.int(self.ll_intp, 1)
def get_shape(self):
return ()
def get_index_bounds(self):
# [idx, idx+1)
return (self.idx, self.builder.add(self.idx, self.get_size()))
def loop_head(self):
return self.idx, None
def loop_tail(self):
pass
class IntegerArrayIndexer(Indexer):
"""
Compute indices from an array of integer indices.
"""
def __init__(self, context, builder, idxty, idxary, size):
self.context = context
self.builder = builder
self.idxty = idxty
self.idxary = idxary
self.size = size
assert idxty.ndim == 1
self.ll_intp = self.context.get_value_type(types.intp)
def prepare(self):
builder = self.builder
self.idx_size = cgutils.unpack_tuple(builder, self.idxary.shape)[0]
self.idx_index = cgutils.alloca_once(builder, self.ll_intp)
self.bb_start = builder.append_basic_block()
self.bb_end = builder.append_basic_block()
def get_size(self):
return self.idx_size
def get_shape(self):
return (self.idx_size,)
def get_index_bounds(self):
# Pessimal heuristic, as we don't want to scan for the min and max
return (self.ll_intp(0), self.size)
def loop_head(self):
builder = self.builder
# Initialize loop variable
self.builder.store(Constant.int(self.ll_intp, 0), self.idx_index)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_start)
cur_index = builder.load(self.idx_index)
with builder.if_then(
builder.icmp_signed('>=', cur_index, self.idx_size),
likely=False
):
builder.branch(self.bb_end)
# Load the actual index from the array of indices
index = _getitem_array_single_int(
self.context, builder, self.idxty.dtype, self.idxty, self.idxary,
cur_index
)
index = fix_integer_index(self.context, builder,
self.idxty.dtype, index, self.size)
return index, cur_index
def loop_tail(self):
builder = self.builder
next_index = cgutils.increment_index(builder,
builder.load(self.idx_index))
builder.store(next_index, self.idx_index)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_end)
class BooleanArrayIndexer(Indexer):
"""
Compute indices from an array of boolean predicates.
"""
def __init__(self, context, builder, idxty, idxary):
self.context = context
self.builder = builder
self.idxty = idxty
self.idxary = idxary
assert idxty.ndim == 1
self.ll_intp = self.context.get_value_type(types.intp)
self.zero = Constant.int(self.ll_intp, 0)
def prepare(self):
builder = self.builder
self.size = cgutils.unpack_tuple(builder, self.idxary.shape)[0]
self.idx_index = cgutils.alloca_once(builder, self.ll_intp)
self.count = cgutils.alloca_once(builder, self.ll_intp)
self.bb_start = builder.append_basic_block()
self.bb_tail = builder.append_basic_block()
self.bb_end = builder.append_basic_block()
def get_size(self):
builder = self.builder
count = cgutils.alloca_once_value(builder, self.zero)
# Sum all true values
with cgutils.for_range(builder, self.size) as loop:
c = builder.load(count)
pred = _getitem_array_single_int(
self.context, builder, self.idxty.dtype,
self.idxty, self.idxary, loop.index
)
c = builder.add(c, builder.zext(pred, c.type))
builder.store(c, count)
return builder.load(count)
def get_shape(self):
return (self.get_size(),)
def get_index_bounds(self):
# Pessimal heuristic, as we don't want to scan for the
# first and last true items
return (self.ll_intp(0), self.size)
def loop_head(self):
builder = self.builder
# Initialize loop variable
self.builder.store(self.zero, self.idx_index)
self.builder.store(self.zero, self.count)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_start)
cur_index = builder.load(self.idx_index)
cur_count = builder.load(self.count)
with builder.if_then(builder.icmp_signed('>=', cur_index, self.size),
likely=False):
builder.branch(self.bb_end)
# Load the predicate and branch if false
pred = _getitem_array_single_int(
self.context, builder, self.idxty.dtype, self.idxty, self.idxary,
cur_index
)
with builder.if_then(builder.not_(pred)):
builder.branch(self.bb_tail)
# Increment the count for next iteration
next_count = cgutils.increment_index(builder, cur_count)
builder.store(next_count, self.count)
return cur_index, cur_count
def loop_tail(self):
builder = self.builder
builder.branch(self.bb_tail)
builder.position_at_end(self.bb_tail)
next_index = cgutils.increment_index(builder,
builder.load(self.idx_index))
builder.store(next_index, self.idx_index)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_end)
class SliceIndexer(Indexer):
"""
Compute indices along a slice.
"""
def __init__(self, context, builder, aryty, ary, dim, idxty, slice):
self.context = context
self.builder = builder
self.aryty = aryty
self.ary = ary
self.dim = dim
self.idxty = idxty
self.slice = slice
self.ll_intp = self.context.get_value_type(types.intp)
self.zero = Constant.int(self.ll_intp, 0)
def prepare(self):
builder = self.builder
# Fix slice for the dimension's size
self.dim_size = builder.extract_value(self.ary.shape, self.dim)
slicing.guard_invalid_slice(self.context, builder, self.idxty,
self.slice)
slicing.fix_slice(builder, self.slice, self.dim_size)
self.is_step_negative = cgutils.is_neg_int(builder, self.slice.step)
# Create loop entities
self.index = cgutils.alloca_once(builder, self.ll_intp)
self.count = cgutils.alloca_once(builder, self.ll_intp)
self.bb_start = builder.append_basic_block()
self.bb_end = builder.append_basic_block()
def get_size(self):
return slicing.get_slice_length(self.builder, self.slice)
def get_shape(self):
return (self.get_size(),)
def get_index_bounds(self):
lower, upper = slicing.get_slice_bounds(self.builder, self.slice)
return lower, upper
def loop_head(self):
builder = self.builder
# Initialize loop variable
self.builder.store(self.slice.start, self.index)
self.builder.store(self.zero, self.count)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_start)
cur_index = builder.load(self.index)
cur_count = builder.load(self.count)
is_finished = builder.select(self.is_step_negative,
builder.icmp_signed('<=', cur_index,
self.slice.stop),
builder.icmp_signed('>=', cur_index,
self.slice.stop))
with builder.if_then(is_finished, likely=False):
builder.branch(self.bb_end)
return cur_index, cur_count
def loop_tail(self):
builder = self.builder
next_index = builder.add(builder.load(self.index), self.slice.step,
flags=['nsw'])
builder.store(next_index, self.index)
next_count = cgutils.increment_index(builder, builder.load(self.count))
builder.store(next_count, self.count)
builder.branch(self.bb_start)
builder.position_at_end(self.bb_end)
class FancyIndexer(object):
"""
Perform fancy indexing on the given array.
"""
def __init__(self, context, builder, aryty, ary, index_types, indices):
self.context = context
self.builder = builder
self.aryty = aryty
self.shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
self.strides = cgutils.unpack_tuple(builder, ary.strides, aryty.ndim)
self.ll_intp = self.context.get_value_type(types.intp)
indexers = []
ax = 0
for indexval, idxty in zip(indices, index_types):
if idxty is types.ellipsis:
# Fill up missing dimensions at the middle
n_missing = aryty.ndim - len(indices) + 1
for i in range(n_missing):
indexer = EntireIndexer(context, builder, aryty, ary, ax)
indexers.append(indexer)
ax += 1
continue
# Regular index value
if isinstance(idxty, types.SliceType):
slice = context.make_helper(builder, idxty, indexval)
indexer = SliceIndexer(context, builder, aryty, ary, ax,
idxty, slice)
indexers.append(indexer)
elif isinstance(idxty, types.Integer):
ind = fix_integer_index(context, builder, idxty, indexval,
self.shapes[ax])
indexer = IntegerIndexer(context, builder, ind)
indexers.append(indexer)
elif isinstance(idxty, types.Array):
idxary = make_array(idxty)(context, builder, indexval)
if isinstance(idxty.dtype, types.Integer):
indexer = IntegerArrayIndexer(context, builder,
idxty, idxary,
self.shapes[ax])
elif isinstance(idxty.dtype, types.Boolean):
indexer = BooleanArrayIndexer(context, builder,
idxty, idxary)
else:
assert 0
indexers.append(indexer)
else:
raise AssertionError("unexpected index type: %s" % (idxty,))
ax += 1
# Fill up missing dimensions at the end
assert ax <= aryty.ndim, (ax, aryty.ndim)
while ax < aryty.ndim:
indexer = EntireIndexer(context, builder, aryty, ary, ax)
indexers.append(indexer)
ax += 1
assert len(indexers) == aryty.ndim, (len(indexers), aryty.ndim)
self.indexers = indexers
def prepare(self):
for i in self.indexers:
i.prepare()
# Compute the resulting shape
self.indexers_shape = sum([i.get_shape() for i in self.indexers], ())
def get_shape(self):
"""
Get the resulting data shape as Python tuple.
"""
return self.indexers_shape
def get_offset_bounds(self, strides, itemsize):
"""
Get a half-open [lower, upper) range of byte offsets spanned by
the indexer with the given strides and itemsize. The indexer is
guaranteed to not go past those bounds.
"""
assert len(strides) == self.aryty.ndim
builder = self.builder
is_empty = cgutils.false_bit
zero = self.ll_intp(0)
one = self.ll_intp(1)
lower = zero
upper = zero
for indexer, shape, stride in zip(self.indexers, self.indexers_shape,
strides):
is_empty = builder.or_(is_empty,
builder.icmp_unsigned('==', shape, zero))
# Compute [lower, upper) indices on this dimension
lower_index, upper_index = indexer.get_index_bounds()
lower_offset = builder.mul(stride, lower_index)
upper_offset = builder.mul(stride, builder.sub(upper_index, one))
# Adjust total interval
is_downwards = builder.icmp_signed('<', stride, zero)
lower = builder.add(lower,
builder.select(is_downwards,
upper_offset,
lower_offset))
upper = builder.add(upper,
builder.select(is_downwards,
lower_offset,
upper_offset))
# Make interval half-open
upper = builder.add(upper, itemsize)
# Adjust for empty shape
lower = builder.select(is_empty, zero, lower)
upper = builder.select(is_empty, zero, upper)
return lower, upper
def begin_loops(self):
indices, counts = zip(*(i.loop_head() for i in self.indexers))
return indices, counts
def end_loops(self):
for i in reversed(self.indexers):
i.loop_tail()
def fancy_getitem(context, builder, sig, args,
aryty, ary, index_types, indices):
shapes = cgutils.unpack_tuple(builder, ary.shape)
strides = cgutils.unpack_tuple(builder, ary.strides)
data = ary.data
indexer = FancyIndexer(context, builder, aryty, ary,
index_types, indices)
indexer.prepare()
# Construct output array
out_ty = sig.return_type
out_shapes = indexer.get_shape()
out = _empty_nd_impl(context, builder, out_ty, out_shapes)
out_data = out.data
out_idx = cgutils.alloca_once_value(builder,
context.get_constant(types.intp, 0))
# Loop on source and copy to destination
indices, _ = indexer.begin_loops()
# No need to check for wraparound, as the indexers all ensure
# a positive index is returned.
ptr = cgutils.get_item_pointer2(context, builder, data, shapes, strides,
aryty.layout, indices, wraparound=False,
boundscheck=context.enable_boundscheck)
val = load_item(context, builder, aryty, ptr)
# Since the destination is C-contiguous, no need for multi-dimensional
# indexing.
cur = builder.load(out_idx)
ptr = builder.gep(out_data, [cur])
store_item(context, builder, out_ty, val, ptr)
next_idx = cgutils.increment_index(builder, cur)
builder.store(next_idx, out_idx)
indexer.end_loops()
return impl_ret_new_ref(context, builder, out_ty, out._getvalue())
@lower_builtin(operator.getitem, types.Buffer, types.Array)
def fancy_getitem_array(context, builder, sig, args):
"""
Advanced or basic indexing with an array.
"""
aryty, idxty = sig.args
ary, idx = args
ary = make_array(aryty)(context, builder, ary)
if idxty.ndim == 0:
# 0-d array index acts as a basic integer index
idxty, idx = normalize_index(context, builder, idxty, idx)
res = _getitem_array_generic(context, builder, sig.return_type,
aryty, ary, (idxty,), (idx,))
return impl_ret_borrowed(context, builder, sig.return_type, res)
else:
# Advanced indexing
return fancy_getitem(context, builder, sig, args,
aryty, ary, (idxty,), (idx,))
def offset_bounds_from_strides(context, builder, arrty, arr, shapes, strides):
"""
Compute a half-open range [lower, upper) of byte offsets from the
array's data pointer, that bound the in-memory extent of the array.
This mimicks offset_bounds_from_strides() from
numpy/core/src/private/mem_overlap.c
"""
itemsize = arr.itemsize
zero = itemsize.type(0)
one = zero.type(1)
if arrty.layout in 'CF':
# Array is contiguous: contents are laid out sequentially
# starting from arr.data and upwards
lower = zero
upper = builder.mul(itemsize, arr.nitems)
else:
# Non-contiguous array: need to examine strides
lower = zero
upper = zero
for i in range(arrty.ndim):
# Compute the largest byte offset on this dimension
# max_axis_offset = strides[i] * (shapes[i] - 1)
# (shapes[i] == 0 is catered for by the empty array case below)
max_axis_offset = builder.mul(strides[i],
builder.sub(shapes[i], one))
is_upwards = builder.icmp_signed('>=', max_axis_offset, zero)
# Expand either upwards or downwards depending on stride
upper = builder.select(is_upwards,
builder.add(upper, max_axis_offset), upper)
lower = builder.select(is_upwards,
lower, builder.add(lower, max_axis_offset))
# Return a half-open range
upper = builder.add(upper, itemsize)
# Adjust for empty arrays
is_empty = builder.icmp_signed('==', arr.nitems, zero)
upper = builder.select(is_empty, zero, upper)
lower = builder.select(is_empty, zero, lower)
return lower, upper
def compute_memory_extents(context, builder, lower, upper, data):
"""
Given [lower, upper) byte offsets and a base data pointer,
compute the memory pointer bounds as pointer-sized integers.
"""
data_ptr_as_int = builder.ptrtoint(data, lower.type)