Skip to content
Permalink
master
Go to file
 
 
Cannot retrieve contributors at this time
5238 lines (4184 sloc) 185 KB
"""
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)
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
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)
start = builder.add(data_ptr_as_int, lower)
end = builder.add(data_ptr_as_int, upper)
return start, end
def get_array_memory_extents(context, builder, arrty, arr, shapes, strides,
data):
"""
Compute a half-open range [start, end) of pointer-sized integers
which fully contain the array data.
"""
lower, upper = offset_bounds_from_strides(context, builder, arrty, arr,
shapes, strides)
return compute_memory_extents(context, builder, lower, upper, data)
def extents_may_overlap(context, builder, a_start, a_end, b_start, b_end):
"""
Whether two memory extents [a_start, a_end) and [b_start, b_end)
may overlap.
"""
# Comparisons are unsigned, since we are really comparing pointers
may_overlap = builder.and_(
builder.icmp_unsigned('<', a_start, b_end),
builder.icmp_unsigned('<', b_start, a_end),
)
return may_overlap
def maybe_copy_source(context, builder, use_copy,
srcty, src, src_shapes, src_strides, src_data):
ptrty = src_data.type
copy_layout = 'C'
copy_data = cgutils.alloca_once_value(builder, src_data)
copy_shapes = src_shapes
copy_strides = None # unneeded for contiguous arrays
with builder.if_then(use_copy, likely=False):
# Allocate temporary scratchpad
# XXX: should we use a stack-allocated array for very small
# data sizes?
allocsize = builder.mul(src.itemsize, src.nitems)
data = context.nrt.allocate(builder, allocsize)
voidptrty = data.type
data = builder.bitcast(data, ptrty)
builder.store(data, copy_data)
# Copy source data into scratchpad
intp_t = context.get_value_type(types.intp)
with cgutils.loop_nest(builder, src_shapes, intp_t) as indices:
src_ptr = cgutils.get_item_pointer2(context, builder, src_data,
src_shapes, src_strides,
srcty.layout, indices)
dest_ptr = cgutils.get_item_pointer2(context, builder, data,
copy_shapes, copy_strides,
copy_layout, indices)
builder.store(builder.load(src_ptr), dest_ptr)
def src_getitem(source_indices):
assert len(source_indices) == srcty.ndim
src_ptr = cgutils.alloca_once(builder, ptrty)
with builder.if_else(use_copy, likely=False) as (if_copy, otherwise):
with if_copy:
builder.store(
cgutils.get_item_pointer2(context, builder,
builder.load(copy_data),
copy_shapes, copy_strides,
copy_layout, source_indices,
wraparound=False),
src_ptr)
with otherwise:
builder.store(
cgutils.get_item_pointer2(context, builder, src_data,
src_shapes, src_strides,
srcty.layout, source_indices,
wraparound=False),
src_ptr)
return load_item(context, builder, srcty, builder.load(src_ptr))
def src_cleanup():
# Deallocate memory
with builder.if_then(use_copy, likely=False):
data = builder.load(copy_data)
data = builder.bitcast(data, voidptrty)
context.nrt.free(builder, data)
return src_getitem, src_cleanup
def _bc_adjust_dimension(context, builder, shapes, strides, target_shape):
"""
Preprocess dimension for broadcasting.
Returns (shapes, strides) such that the ndim match *target_shape*.
When expanding to higher ndim, the returning shapes and strides are
prepended with ones and zeros, respectively.
When truncating to lower ndim, the shapes are checked (in runtime).
All extra dimension must have size of 1.
"""
zero = context.get_constant(types.uintp, 0)
one = context.get_constant(types.uintp, 1)
# Adjust for broadcasting to higher dimension
if len(target_shape) > len(shapes):
nd_diff = len(target_shape) - len(shapes)
# Fill missing shapes with one, strides with zeros
shapes = [one] * nd_diff + shapes
strides = [zero] * nd_diff + strides
# Adjust for broadcasting to lower dimension
elif len(target_shape) < len(shapes):
# Accepted if all extra dims has shape 1
nd_diff = len(shapes) - len(target_shape)
dim_is_one = [builder.icmp_unsigned('==', sh, one)
for sh in shapes[:nd_diff]]
accepted = functools.reduce(builder.and_, dim_is_one,
cgutils.true_bit)
# Check error
with builder.if_then(builder.not_(accepted), likely=False):
msg = "cannot broadcast source array for assignment"
context.call_conv.return_user_exc(builder, ValueError, (msg,))
# Truncate extra shapes, strides
shapes = shapes[nd_diff:]
strides = strides[nd_diff:]
return shapes, strides
def _bc_adjust_shape_strides(context, builder, shapes, strides, target_shape):
"""
Broadcast shapes and strides to target_shape given that their ndim already
matches. For each location where the shape is 1 and does not match the
dim for target, it is set to the value at the target and the stride is
set to zero.
"""
bc_shapes = []
bc_strides = []
zero = context.get_constant(types.uintp, 0)
one = context.get_constant(types.uintp, 1)
# Adjust all mismatching ones in shape
mismatch = [builder.icmp_signed('!=', tar, old)
for tar, old in zip(target_shape, shapes)]
src_is_one = [builder.icmp_signed('==', old, one) for old in shapes]
preds = [builder.and_(x, y) for x, y in zip(mismatch, src_is_one)]
bc_shapes = [builder.select(p, tar, old)
for p, tar, old in zip(preds, target_shape, shapes)]
bc_strides = [builder.select(p, zero, old)
for p, old in zip(preds, strides)]
return bc_shapes, bc_strides
def _broadcast_to_shape(context, builder, arrtype, arr, target_shape):
"""
Broadcast the given array to the target_shape.
Returns (array_type, array)
"""
# Compute broadcasted shape and strides
shapes = cgutils.unpack_tuple(builder, arr.shape)
strides = cgutils.unpack_tuple(builder, arr.strides)
shapes, strides = _bc_adjust_dimension(context, builder, shapes, strides,
target_shape)
shapes, strides = _bc_adjust_shape_strides(context, builder, shapes,
strides, target_shape)
new_arrtype = arrtype.copy(ndim=len(target_shape), layout='A')
# Create new view
new_arr = make_array(new_arrtype)(context, builder)
repl = dict(shape=cgutils.pack_array(builder, shapes),
strides=cgutils.pack_array(builder, strides))
cgutils.copy_struct(new_arr, arr, repl)
return new_arrtype, new_arr
def fancy_setslice(context, builder, sig, args, index_types, indices):
"""
Implement slice assignment for arrays. This implementation works for
basic as well as fancy indexing, since there's no functional difference
between the two for indexed assignment.
"""
aryty, _, srcty = sig.args
ary, _, src = args
ary = make_array(aryty)(context, builder, ary)
dest_shapes = cgutils.unpack_tuple(builder, ary.shape)
dest_strides = cgutils.unpack_tuple(builder, ary.strides)
dest_data = ary.data
indexer = FancyIndexer(context, builder, aryty, ary,
index_types, indices)
indexer.prepare()
if isinstance(srcty, types.Buffer):
# Source is an array
src_dtype = srcty.dtype
index_shape = indexer.get_shape()
src = make_array(srcty)(context, builder, src)
# Broadcast source array to shape
srcty, src = _broadcast_to_shape(context, builder, srcty, src,
index_shape)
src_shapes = cgutils.unpack_tuple(builder, src.shape)
src_strides = cgutils.unpack_tuple(builder, src.strides)
src_data = src.data
# Check shapes are equal
shape_error = cgutils.false_bit
assert len(index_shape) == len(src_shapes)
for u, v in zip(src_shapes, index_shape):
shape_error = builder.or_(shape_error,
builder.icmp_signed('!=', u, v))
with builder.if_then(shape_error, likely=False):
msg = "cannot assign slice from input of different size"
context.call_conv.return_user_exc(builder, ValueError, (msg,))
# Check for array overlap
src_start, src_end = get_array_memory_extents(context, builder, srcty,
src, src_shapes,
src_strides, src_data)
dest_lower, dest_upper = indexer.get_offset_bounds(dest_strides,
ary.itemsize)
dest_start, dest_end = compute_memory_extents(context, builder,
dest_lower, dest_upper,
dest_data)
use_copy = extents_may_overlap(context, builder, src_start, src_end,
dest_start, dest_end)
src_getitem, src_cleanup = maybe_copy_source(context, builder, use_copy,
srcty, src, src_shapes,
src_strides, src_data)
elif isinstance(srcty, types.Sequence):
src_dtype = srcty.dtype
# Check shape is equal to sequence length
index_shape = indexer.get_shape()
assert len(index_shape) == 1
len_impl = context.get_function(len, signature(types.intp, srcty))
seq_len = len_impl(builder, (src,))
shape_error = builder.icmp_signed('!=', index_shape[0], seq_len)
with builder.if_then(shape_error, likely=False):
msg = "cannot assign slice from input of different size"
context.call_conv.return_user_exc(builder, ValueError, (msg,))
def src_getitem(source_indices):
idx, = source_indices
getitem_impl = context.get_function(
operator.getitem,
signature(src_dtype, srcty, types.intp),
)
return getitem_impl(builder, (src, idx))
def src_cleanup():
pass
else:
# Source is a scalar (broadcast or not, depending on destination
# shape).
src_dtype = srcty
def src_getitem(source_indices):
return src
def src_cleanup():
pass
# Loop on destination and copy from source to destination
dest_indices, counts = indexer.begin_loops()
# Source is iterated in natural order
source_indices = tuple(c for c in counts if c is not None)
val = src_getitem(source_indices)
# Cast to the destination dtype (cross-dtype slice assignement is allowed)
val = context.cast(builder, val, src_dtype, aryty.dtype)
# No need to check for wraparound, as the indexers all ensure
# a positive index is returned.
dest_ptr = cgutils.get_item_pointer2(context, builder, dest_data,
dest_shapes, dest_strides,
aryty.layout, dest_indices,
wraparound=False)
store_item(context, builder, aryty, val, dest_ptr)
indexer.end_loops()
src_cleanup()
return context.get_dummy_value()
# ------------------------------------------------------------------------------
# Shape / layout altering
def vararg_to_tuple(context, builder, sig, args):
aryty = sig.args[0]
dimtys = sig.args[1:]
# values
ary = args[0]
dims = args[1:]
# coerce all types to intp
dims = [context.cast(builder, val, ty, types.intp)
for ty, val in zip(dimtys, dims)]
# make a tuple
shape = cgutils.pack_array(builder, dims, dims[0].type)
shapety = types.UniTuple(dtype=types.intp, count=len(dims))
new_sig = typing.signature(sig.return_type, aryty, shapety)
new_args = ary, shape
return new_sig, new_args
@lower_builtin('array.transpose', types.Array)
def array_transpose(context, builder, sig, args):
return array_T(context, builder, sig.args[0], args[0])
def permute_arrays(axis, shape, strides):
if len(axis) != len(set(axis)):
raise ValueError("repeated axis in transpose")
dim = len(shape)
for x in axis:
if x >= dim or abs(x) > dim:
raise ValueError("axis is out of bounds for array of "
"given dimension")
shape[:] = shape[axis]
strides[:] = strides[axis]
# Transposing an array involves permuting the shape and strides of the array
# based on the given axes.
@lower_builtin('array.transpose', types.Array, types.BaseTuple)
def array_transpose_tuple(context, builder, sig, args):
aryty = sig.args[0]
ary = make_array(aryty)(context, builder, args[0])
axisty, axis = sig.args[1], args[1]
num_axis, dtype = axisty.count, axisty.dtype
ll_intp = context.get_value_type(types.intp)
ll_ary_size = lc.Type.array(ll_intp, num_axis)
# Allocate memory for axes, shapes, and strides arrays.
arys = [axis, ary.shape, ary.strides]
ll_arys = [cgutils.alloca_once(builder, ll_ary_size) for _ in arys]
# Store axes, shapes, and strides arrays to the allocated memory.
for src, dst in zip(arys, ll_arys):
builder.store(src, dst)
np_ary_ty = types.Array(dtype=dtype, ndim=1, layout='C')
np_itemsize = context.get_constant(types.intp,
context.get_abi_sizeof(ll_intp))
# Form NumPy arrays for axes, shapes, and strides arrays.
np_arys = [make_array(np_ary_ty)(context, builder) for _ in arys]
# Roughly, `np_ary = np.array(ll_ary)` for each of axes, shapes, and strides
for np_ary, ll_ary in zip(np_arys, ll_arys):
populate_array(np_ary,
data=builder.bitcast(ll_ary, ll_intp.as_pointer()),
shape=[context.get_constant(types.intp, num_axis)],
strides=[np_itemsize],
itemsize=np_itemsize,
meminfo=None)
# Pass NumPy arrays formed above to permute_arrays function that permutes
# shapes and strides based on axis contents.
context.compile_internal(builder, permute_arrays,
typing.signature(types.void,
np_ary_ty, np_ary_ty, np_ary_ty),
[a._getvalue() for a in np_arys])
# Make a new array based on permuted shape and strides and return it.
ret = make_array(sig.return_type)(context, builder)
populate_array(ret,
data=ary.data,
shape=builder.load(ll_arys[1]),
strides=builder.load(ll_arys[2]),
itemsize=ary.itemsize,
meminfo=ary.meminfo,
parent=ary.parent)
res = ret._getvalue()
return impl_ret_borrowed(context, builder, sig.return_type, res)
@lower_builtin('array.transpose', types.Array, types.VarArg(types.Any))
def array_transpose_vararg(context, builder, sig, args):
new_sig, new_args = vararg_to_tuple(context, builder, sig, args)
return array_transpose_tuple(context, builder, new_sig, new_args)
@overload(np.transpose)
def numpy_transpose(a, axes=None):
if isinstance(a, types.BaseTuple):
raise errors.UnsupportedError("np.transpose does not accept tuples")
if axes is None:
def np_transpose_impl(a, axes=None):
return a.transpose()
else:
def np_transpose_impl(a, axes=None):
return a.transpose(axes)
return np_transpose_impl
@lower_getattr(types.Array, 'T')
def array_T(context, builder, typ, value):
if typ.ndim <= 1:
res = value
else:
ary = make_array(typ)(context, builder, value)
ret = make_array(typ)(context, builder)
shapes = cgutils.unpack_tuple(builder, ary.shape, typ.ndim)
strides = cgutils.unpack_tuple(builder, ary.strides, typ.ndim)
populate_array(ret,
data=ary.data,
shape=cgutils.pack_array(builder, shapes[::-1]),
strides=cgutils.pack_array(builder, strides[::-1]),
itemsize=ary.itemsize,
meminfo=ary.meminfo,
parent=ary.parent)
res = ret._getvalue()
return impl_ret_borrowed(context, builder, typ, res)
def _attempt_nocopy_reshape(context, builder, aryty, ary,
newnd, newshape, newstrides):
"""
Call into Numba_attempt_nocopy_reshape() for the given array type
and instance, and the specified new shape.
Return value is non-zero if successful, and the array pointed to
by *newstrides* will be filled up with the computed results.
"""
ll_intp = context.get_value_type(types.intp)
ll_intp_star = ll_intp.as_pointer()
ll_intc = context.get_value_type(types.intc)
fnty = lc.Type.function(ll_intc, [
# nd, *dims, *strides
ll_intp, ll_intp_star, ll_intp_star,
# newnd, *newdims, *newstrides
ll_intp, ll_intp_star, ll_intp_star,
# itemsize, is_f_order
ll_intp, ll_intc])
fn = builder.module.get_or_insert_function(
fnty, name="numba_attempt_nocopy_reshape")
nd = ll_intp(aryty.ndim)
shape = cgutils.gep_inbounds(builder, ary._get_ptr_by_name('shape'), 0, 0)
strides = cgutils.gep_inbounds(builder, ary._get_ptr_by_name('strides'),
0, 0)
newnd = ll_intp(newnd)
newshape = cgutils.gep_inbounds(builder, newshape, 0, 0)
newstrides = cgutils.gep_inbounds(builder, newstrides, 0, 0)
is_f_order = ll_intc(0)
res = builder.call(fn, [nd, shape, strides,
newnd, newshape, newstrides,
ary.itemsize, is_f_order])
return res
def normalize_reshape_value(origsize, shape):
num_neg_value = 0
known_size = 1
for ax, s in enumerate(shape):
if s < 0:
num_neg_value += 1
neg_ax = ax
else:
known_size *= s
if num_neg_value == 0:
if origsize != known_size:
raise ValueError("total size of new array must be unchanged")
elif num_neg_value == 1:
# Infer negative dimension
if known_size == 0:
inferred = 0
ok = origsize == 0
else:
inferred = origsize // known_size
ok = origsize % known_size == 0
if not ok:
raise ValueError("total size of new array must be unchanged")
shape[neg_ax] = inferred
else:
raise ValueError("multiple negative shape values")
@lower_builtin('array.reshape', types.Array, types.BaseTuple)
def array_reshape(context, builder, sig, args):
aryty = sig.args[0]
retty = sig.return_type
shapety = sig.args[1]
shape = args[1]
ll_intp = context.get_value_type(types.intp)
ll_shape = lc.Type.array(ll_intp, shapety.count)
ary = make_array(aryty)(context, builder, args[0])
# We will change the target shape in this slot
# (see normalize_reshape_value() below)
newshape = cgutils.alloca_once(builder, ll_shape)
builder.store(shape, newshape)
# Create a shape array pointing to the value of newshape.
# (roughly, `shape_ary = np.array(ary.shape)`)
shape_ary_ty = types.Array(dtype=shapety.dtype, ndim=1, layout='C')
shape_ary = make_array(shape_ary_ty)(context, builder)
shape_itemsize = context.get_constant(types.intp,
context.get_abi_sizeof(ll_intp))
populate_array(shape_ary,
data=builder.bitcast(newshape, ll_intp.as_pointer()),
shape=[context.get_constant(types.intp, shapety.count)],
strides=[shape_itemsize],
itemsize=shape_itemsize,
meminfo=None)
# Compute the original array size
size = ary.nitems
# Call our normalizer which will fix the shape array in case of negative
# shape value
context.compile_internal(builder, normalize_reshape_value,
typing.signature(types.void,
types.uintp, shape_ary_ty),
[size, shape_ary._getvalue()])
# Perform reshape (nocopy)
newnd = shapety.count
newstrides = cgutils.alloca_once(builder, ll_shape)
ok = _attempt_nocopy_reshape(context, builder, aryty, ary, newnd,
newshape, newstrides)
fail = builder.icmp_unsigned('==', ok, ok.type(0))
with builder.if_then(fail):
msg = "incompatible shape for array"
context.call_conv.return_user_exc(builder, NotImplementedError, (msg,))
ret = make_array(retty)(context, builder)
populate_array(ret,