[pypy-svn] r46379 - in pypy/dist/pypy/rpython/numpy: . test
simonb at codespeak.net
simonb at codespeak.net
Fri Sep 7 01:18:51 CEST 2007
Author: simonb
Date: Fri Sep 7 01:18:49 2007
New Revision: 46379
Modified:
pypy/dist/pypy/rpython/numpy/aarray.py
pypy/dist/pypy/rpython/numpy/rarray.py
pypy/dist/pypy/rpython/numpy/test/test_array.py
Log:
broadcasting and slicing, many more cases work
Modified: pypy/dist/pypy/rpython/numpy/aarray.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/aarray.py (original)
+++ pypy/dist/pypy/rpython/numpy/aarray.py Fri Sep 7 01:18:49 2007
@@ -99,11 +99,10 @@
if len(s_index.items)>s_array.ndim:
raise AnnotatorError("invalid index")
if isinstance(s_value, SomeArray):
- if s_value.ndim != ndim:
- # XX allow broadcasting..
+ if s_value.ndim > ndim:
raise AnnotatorError("shape mismatch")
- elif ndim > 0:
- raise AnnotatorError("need to set from array")
+ #elif ndim > 0:
+ # raise AnnotatorError("need to set from array")
def getitem((s_array, s_index)):
ndim = pair(s_array, s_index).get_leftover_dim()
Modified: pypy/dist/pypy/rpython/numpy/rarray.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/rarray.py (original)
+++ pypy/dist/pypy/rpython/numpy/rarray.py Fri Sep 7 01:18:49 2007
@@ -1,4 +1,4 @@
-from pypy.rpython.rmodel import Repr, inputconst
+from pypy.rpython.rmodel import Repr, FloatRepr, inputconst
from pypy.rpython.rrange import AbstractRangeRepr
from pypy.rpython.rint import IntegerRepr
from pypy.rpython.rlist import AbstractBaseListRepr
@@ -70,7 +70,7 @@
("index", NPY_INTP),
("size", NPY_INTP),
("coordinates", INDEXARRAY),
- ("dims_m1", INDEXARRAY),
+ ("dims_m1", INDEXARRAY), # array of dimensions - 1
("strides", INDEXARRAY),
("backstrides", INDEXARRAY),
#("factors", INDEXARRAY),
@@ -115,21 +115,22 @@
return it
ll_iter_new._always_inline_ = True
- def ll_iter_broadcast_to_shape(ITER, ao, shape, iter_reset=ll_iter_reset):
- if ao.ndim > ndim:
- raise Exception("array is not broadcastable to correct shape") # XX raise here ?
- diff = j = ndim - ao.ndim
+ def ll_iter_broadcast_to_shape(ITER, ao, target_ao, iter_reset=ll_iter_reset):
+ "iterate over <ao> but broadcast to the shape of <target_ao>"
+ assert target_ao.ndim == ndim
+ delta = j = ndim - ao.ndim
+ shape = target_ao.shape
for i in range(ao.ndim):
if ao.shape[i] != 1 and ao.shape[i] != shape[j]:
- raise Exception("array is not broadcastable to correct shape") # XX raise here ?
+ raise Exception("array is not broadcastable to correct shape")
j += 1
it = malloc(ITER)
- it.size = ll_mul_list(ao.shape, ndim)
+ it.size = ll_mul_list(target_ao.shape, ndim)
it.nd_m1 = ndim - 1
- #it.factors[nd-1] = 1
+ #it.factors[ndim-1] = 1
for i in unroll_ndim:
- it.dims_m1[i] = ao.shape[i]-1
- k = i - diff
+ it.dims_m1[i] = shape[i]-1
+ k = i - delta
if k<0 or ao.shape[k] != shape[i]:
#it.contiguous = False
it.strides[i] = 0
@@ -137,7 +138,7 @@
it.strides[i] = ao.strides[k]
it.backstrides[i] = it.strides[i] * it.dims_m1[i]
#if i > 0:
- #it.factors[nd-i-1] = it.factors[nd-i]*shape[nd-i]
+ #it.factors[ndim-i-1] = it.factors[nd-i]*shape[ndim-i]
iter_reset(it, ao.dataptr)
return it
ll_iter_broadcast_to_shape._always_inline_ = True
@@ -153,8 +154,7 @@
it.dataptr = direct_ptradd(it.dataptr, -it.backstrides[i])
ll_iter_next._always_inline_ = True
-# return ll_iter_new, ll_iter_broadcast_to_shape, ll_iter_next
- return ll_iter_new, ll_iter_next
+ return ll_iter_new, ll_iter_reset, ll_iter_broadcast_to_shape, ll_iter_next
def ll_unary_op(p0, p1, op=lambda x:x):
p0[0] = op(p1[0])
@@ -162,16 +162,19 @@
def ll_binary_op(p0, p1, p2, op=lambda x,y:x+y):
p0[0] = op(p1[0], p2[0])
-def ll_array_unary_op(iter_new, iter_next, ITER, array0, array1):
- assert array0.ndim == array1.ndim
- it0 = iter_new(ITER, array0)
- it1 = iter_new(ITER, array1)
+
+def ll_array_set(it0, it1, iter_next):
assert it0.size == it1.size
while it0.index < it0.size:
it0.dataptr[0] = it1.dataptr[0]
iter_next(it0)
iter_next(it1)
+def ll_array_set1(value, it, iter_next):
+ while it.index < it.size:
+ it.dataptr[0] = value
+ iter_next(it)
+
def dim_of_ITER(ITER):
return ITER.TO.coordinates.length
@@ -275,32 +278,101 @@
class __extend__(pairtype(ArrayRepr, ArrayRepr)):
- def rtype_add((r_arr1, r_arr2), hop):
- v_arr1, v_arr2 = hop.inputargs(r_arr1, r_arr2)
+ def rtype_add((r_array1, r_array2), hop):
+ v_arr1, v_arr2 = hop.inputargs(r_array1, r_array2)
cARRAY = hop.inputconst(Void, hop.r_result.ARRAY.TO)
return hop.gendirectcall(ll_add, cARRAY, v_arr1, v_arr2)
#class __extend__(pairtype(ArrayRepr, Repr)): # <------ USE THIS ??
class __extend__(pairtype(ArrayRepr, IntegerRepr)):
- def rtype_setitem((r_arr, r_int), hop):
- v_array, v_index, v_item = hop.inputargs(r_arr, Signed, r_arr.item_repr)
+ def rtype_setitem((r_array, r_int), hop):
+ assert r_array.ndim == 1, "NotImplemented"
+ v_array, v_index, v_item = hop.inputargs(r_array, Signed, r_array.item_repr)
return hop.gendirectcall(ll_setitem1, v_array, v_index, v_item)
- def rtype_getitem((r_arr, r_int), hop):
- v_array, v_index = hop.inputargs(r_arr, Signed)
+ def rtype_getitem((r_array, r_int), hop):
+ assert r_array.ndim == 1, "NotImplemented"
+ v_array, v_index = hop.inputargs(r_array, Signed)
return hop.gendirectcall(ll_getitem1, v_array, v_index)
+def gen_get_view_slc(r_array, r_slc, hop): # XX method on the pair type ?
+ ndim = r_array.ndim
+ rslice = hop.rtyper.type_system.rslice
+ def ll_get_view_slc(ARRAY, ao, slc):
+ array = ll_allocate(ARRAY, ndim)
+ dataptr = direct_arrayitems(ao.data)
+ src_i = 0
+ tgt_i = 0
+ if r_slc == rslice.startonly_slice_repr:
+ start = slc
+ size = ao.shape[src_i]
+ if start > size:
+ start = size
+ size -= start
+ dataptr = direct_ptradd(dataptr, start*ao.strides[src_i])
+ array.shape[tgt_i] = size
+ array.strides[tgt_i] = ao.strides[src_i]
+ tgt_i += 1
+ elif r_slc == rslice.startstop_slice_repr:
+ start = slc.start
+ stop = slc.stop
+ size = ao.shape[src_i]
+ if start > size:
+ start = size
+ dataptr = direct_ptradd(dataptr, start*ao.strides[src_i])
+ if stop < size:
+ size = stop
+ size -= start
+ if size < 0:
+ size = 0
+ array.shape[tgt_i] = size
+ array.strides[tgt_i] = ao.strides[src_i]
+ tgt_i += 1
+ else:
+ assert 0
+ src_i += 1
+ # consume the rest of ndim as if we found more slices
+ while tgt_i < ndim:
+ array.shape[tgt_i] = ao.shape[src_i]
+ array.strides[tgt_i] = ao.strides[src_i]
+ tgt_i += 1
+ src_i += 1
+ assert tgt_i == ndim
+ array.dataptr = dataptr
+ array.data = ao.data # keep a ref
+ return array
+ return ll_get_view_slc
+
class __extend__(pairtype(ArrayRepr, AbstractSliceRepr)):
- def rtype_setitem((r_arr, r_slc), hop):
+ def rtype_setitem((r_array, r_slc), hop):
r_item = hop.args_r[2]
- v_array, v_slc, v_item = hop.inputargs(r_arr, r_slc, r_item)
- cITER = hop.inputconst(Void, r_arr.ITER.TO)
- iter_new, iter_next = gen_iter_funcs(r_arr.ndim)
- cnew = hop.inputconst(Void, iter_new)
+ v_array, v_slc, v_item = hop.inputargs(r_array, r_slc, r_item)
+ cITER = hop.inputconst(Void, r_array.ITER.TO)
+ cARRAY = hop.inputconst(Void, r_array.ARRAY.TO)
+ iter_new, iter_reset, iter_broadcast, iter_next = gen_iter_funcs(r_array.ndim)
cnext = hop.inputconst(Void, iter_next)
- assert r_arr.ndim == r_item.ndim
- return hop.gendirectcall(ll_array_unary_op, cnew, cnext, cITER, v_array, v_item)
+ creset = hop.inputconst(Void, iter_reset)
+## Blech... it would be nice to reuse gen_get_view
+## r_tuple = TupleRepr(hop.rtyper, [r_item]) # XX how to get this from rtyper ?
+## get_view = gen_get_view(r_array, r_tuple, hop)
+## # make a v_tuple here...
+## v_view = hop.gendirectcall(get_view, cARRAY, v_array, v_tuple)
+ get_view = gen_get_view_slc(r_array, r_slc, hop)
+ v_view = hop.gendirectcall(get_view, cARRAY, v_array, v_slc)
+ v_it0 = hop.gendirectcall(iter_new, cITER, v_view, creset)
+ if isinstance(r_item, ArrayRepr):
+ if r_array.ndim == r_item.ndim:
+ v_it1 = hop.gendirectcall(iter_new, cITER, v_item, creset)
+ else:
+ v_it1 = hop.gendirectcall(iter_broadcast, cITER, v_item, v_array, creset)
+ assert r_array.ndim >= r_item.ndim
+ return hop.gendirectcall(ll_array_set, v_it0, v_it1, cnext)
+ elif isinstance(r_item, FloatRepr):
+ # setitem from scalar
+ return hop.gendirectcall(ll_array_set1, v_item, v_it0, cnext)
+ else:
+ raise TypeError("can't setitem from %s"%r_item)
def gen_getset_item(ndim):
unrolling_dims = unrolling_iterable(range(ndim))
@@ -320,26 +392,57 @@
return ll_get_item, ll_set_item
-def get_view_ndim(r_tpl):
- return len([r_item for r_item in r_tpl.items_r if isinstance(r_item, AbstractSliceRepr)])
-
-def gen_get_view(r_tpl):
- ndim = get_view_ndim(r_tpl)
- unroll_r_tpl = unrolling_iterable(enumerate(r_tpl.items_r))
+def get_view_ndim(r_array, r_tuple): # XX method on the pair type ?
+ ndim = len([r_item for r_item in r_tuple.items_r if isinstance(r_item, AbstractSliceRepr)])
+ ndim += r_array.ndim - len(r_tuple.items_r)
+ return ndim
+
+def gen_get_view(r_array, r_tuple, hop): # XX method on the pair type ?
+ ndim = get_view_ndim(r_array, r_tuple)
+ unroll_r_tuple = unrolling_iterable(enumerate(r_tuple.items_r))
+ rslice = hop.rtyper.type_system.rslice
def ll_get_view(ARRAY, ao, tpl):
array = ll_allocate(ARRAY, ndim)
dataptr = direct_arrayitems(ao.data)
src_i = 0
tgt_i = 0
- for src_i, r_item in unroll_r_tpl:
- if isinstance(r_item, IntegerRepr):
- r_int = r_item
+ for src_i, r_key in unroll_r_tuple:
+ if isinstance(r_key, IntegerRepr):
dataptr = direct_ptradd(dataptr, getattr(tpl, 'item%d'%src_i)*ao.strides[src_i])
- else:
- r_slice = r_item
- array.shape[tgt_i] = ao.shape[src_i]
+ elif r_key == rslice.startonly_slice_repr:
+ start = getattr(tpl, 'item%d'%src_i)
+ size = ao.shape[src_i]
+ if start > size:
+ start = size
+ size -= start
+ dataptr = direct_ptradd(dataptr, start*ao.strides[src_i])
+ array.shape[tgt_i] = size
+ array.strides[tgt_i] = ao.strides[src_i]
+ tgt_i += 1
+ elif r_key == rslice.startstop_slice_repr:
+ start = getattr(tpl, 'item%d'%src_i).start
+ stop = getattr(tpl, 'item%d'%src_i).stop
+ size = ao.shape[src_i]
+ if start > size:
+ start = size
+ dataptr = direct_ptradd(dataptr, start*ao.strides[src_i])
+ if stop < size:
+ size = stop
+ size -= start
+ if size < 0:
+ size = 0
+ array.shape[tgt_i] = size
array.strides[tgt_i] = ao.strides[src_i]
tgt_i += 1
+ else:
+ assert 0
+ src_i += 1
+ # consume the rest of ndim as if we found more slices
+ while tgt_i < ndim:
+ array.shape[tgt_i] = ao.shape[src_i]
+ array.strides[tgt_i] = ao.strides[src_i]
+ tgt_i += 1
+ src_i += 1
assert tgt_i == ndim
array.dataptr = dataptr
array.data = ao.data # keep a ref
@@ -348,63 +451,70 @@
class __extend__(pairtype(ArrayRepr, AbstractTupleRepr)):
- def rtype_getitem((r_arr, r_tpl), hop):
- v_array, v_tuple = hop.inputargs(r_arr, r_tpl)
- ndim = get_view_ndim(r_tpl)
+ def rtype_getitem((r_array, r_tpl), hop):
+ v_array, v_tuple = hop.inputargs(r_array, r_tpl)
+ ndim = get_view_ndim(r_array, r_tpl)
if ndim == 0:
# return a scalar
- cARRAY = hop.inputconst(Void, r_arr.ARRAY.TO)
- get_item, set_item = gen_getset_item(r_arr.ndim)
+ cARRAY = hop.inputconst(Void, r_array.ARRAY.TO)
+ get_item, set_item = gen_getset_item(r_array.ndim)
return hop.gendirectcall(get_item, cARRAY, v_array, v_tuple)
r_result = hop.r_result
ARRAY = r_result.ARRAY
assert dim_of_ARRAY(ARRAY) == ndim
cARRAY = hop.inputconst(Void, ARRAY.TO)
- ll_get_view = gen_get_view(r_tpl)
+ ll_get_view = gen_get_view(r_array, r_tpl, hop)
return hop.gendirectcall(ll_get_view, cARRAY, v_array, v_tuple)
- def rtype_setitem((r_arr, r_tpl), hop):
+ def rtype_setitem((r_array, r_tuple), hop):
r_item = hop.args_r[2]
- v_array, v_tuple, v_item = hop.inputargs(r_arr, r_tpl, r_item)
- ndim = get_view_ndim(r_tpl)
- if isinstance(r_item, ArrayRepr):
- s_view = SomeArray(r_arr.s_array.typecode, ndim)
+ v_array, v_tuple, v_item = hop.inputargs(r_array, r_tuple, r_item)
+ ndim = get_view_ndim(r_array, r_tuple)
+ assert len(r_tuple.items_r) <= r_array.ndim
+ if ndim == 0:
+ # Set from scalar
+ assert isinstance(r_item, FloatRepr)
+ cARRAY = hop.inputconst(Void, r_array.ARRAY.TO)
+ get_item, set_item = gen_getset_item(r_array.ndim)
+ return hop.gendirectcall(set_item, cARRAY, v_array, v_tuple, v_item)
+ elif isinstance(r_item, ArrayRepr):
+ s_view = SomeArray(r_array.s_array.typecode, ndim)
r_view = hop.rtyper.getrepr(s_view)
cARRAY = hop.inputconst(Void, r_view.ARRAY.TO)
- get_view = gen_get_view(r_tpl)
+ get_view = gen_get_view(r_array, r_tuple, hop)
v_view = hop.gendirectcall(get_view, cARRAY, v_array, v_tuple)
- iter_new, iter_next = gen_iter_funcs(ndim)
- assert ndim == r_item.ndim
- assert dim_of_ITER(r_item.ITER) == dim_of_ITER(r_view.ITER)
- cnew = hop.inputconst(Void, iter_new)
+ iter_new, iter_reset, iter_broadcast, iter_next = gen_iter_funcs(ndim)
+ creset = hop.inputconst(Void, iter_reset)
cnext = hop.inputconst(Void, iter_next)
- cITER = hop.inputconst(Void, r_item.ITER.TO)
- return hop.gendirectcall(ll_array_unary_op, cnew, cnext, cITER, v_view, v_item)
+ cITER = hop.inputconst(Void, r_view.ITER.TO)
+ v_it0 = hop.gendirectcall(iter_new, cITER, v_view, creset)
+ assert r_item.ndim <= ndim
+ if ndim == r_item.ndim:
+ v_it1 = hop.gendirectcall(iter_new, cITER, v_item, creset)
+ else:
+ v_it1 = hop.gendirectcall(iter_broadcast, cITER, v_item, v_view, creset)
+ return hop.gendirectcall(ll_array_set, v_it0, v_it1, cnext)
else:
- # Set from scalar
- assert ndim == 0
- cARRAY = hop.inputconst(Void, r_arr.ARRAY.TO)
- get_item, set_item = gen_getset_item(r_arr.ndim)
- return hop.gendirectcall(set_item, cARRAY, v_array, v_tuple, v_item)
+ assert 0
class __extend__(pairtype(ArrayRepr, ArrayRepr)):
- def convert_from_to((r_arr0, r_arr1), v, llops):
+ def convert_from_to((r_array0, r_array1), v, llops):
assert 0
class __extend__(pairtype(AbstractBaseListRepr, ArrayRepr)):
- def convert_from_to((r_lst, r_arr), v, llops):
+ def convert_from_to((r_lst, r_array), v, llops):
if r_lst.listitem is None:
return NotImplemented
- if r_lst.item_repr != r_arr.item_repr:
- assert 0, (r_lst, r_arr.item_repr)
+ if r_lst.item_repr != r_array.item_repr:
+ assert 0, (r_lst, r_array.item_repr)
return NotImplemented
- cARRAY = inputconst(lltype.Void, r_arr.lowleveltype.TO)
+ cARRAY = inputconst(lltype.Void, r_array.lowleveltype.TO)
return llops.gendirectcall(ll_build_from_list, cARRAY, v)
class __extend__(pairtype(AbstractRangeRepr, ArrayRepr)):
- def convert_from_to((r_rng, r_arr), v, llops):
- cARRAY = inputconst(lltype.Void, r_arr.lowleveltype.TO)
+ def convert_from_to((r_rng, r_array), v, llops):
+ cARRAY = inputconst(lltype.Void, r_array.lowleveltype.TO)
return llops.gendirectcall(ll_build_from_list, cARRAY, v)
def ll_allocate(ARRAY, ndim):
Modified: pypy/dist/pypy/rpython/numpy/test/test_array.py
==============================================================================
--- pypy/dist/pypy/rpython/numpy/test/test_array.py (original)
+++ pypy/dist/pypy/rpython/numpy/test/test_array.py Fri Sep 7 01:18:49 2007
@@ -187,6 +187,19 @@
assert type(s_array) == SomeArray
assert s_array.ndim == 1
+ def test_annotate_broadcast(self):
+ def f():
+ a = numpy.empty((4,3), dtype='i')
+ b = numpy.array([33])
+ a[:] = b
+ return a
+ t = TranslationContext()
+ a = t.buildannotator()
+ s_array = a.build_types(f, [])
+ assert type(s_array) == SomeArray
+ assert s_array.ndim == 2
+
+
from pypy.objspace.flow.model import checkgraph, flatten, Block, mkentrymap
from pypy.translator.backendopt.malloc import LLTypeMallocRemover
@@ -341,6 +354,66 @@
assert interpret(f, [0, 0]) == 0
assert interpret(f, [3, 4]) == 12
+ def test_specialize_slice_1_0(self):
+ def f():
+ a = numpy.zeros((12,), dtype='i')
+ a[:2] = 1
+ a[5:9] = 2
+ a[10:] = 3
+ a[12:] = 99
+ a[12:0] = 999
+ return a
+ res = interpret(f, [])
+ data = [1,1,0,0,0,2,2,2,2,0,3,3]
+ for i in range(12):
+ assert res.dataptr[i] == data[i]
+
+ def test_specialize_slice_1_1(self):
+ py.test.skip('this involves a runtime test to see if we need a broadcast iterator')
+ def f():
+ a = numpy.zeros((6,), dtype='i')
+ a[:2] = numpy.array([1])
+ a[5:9] = numpy.array([2])
+ return a
+ res = interpret(f, [])
+ data = [1,1,0,0,0,2]
+ for i in range(6):
+ assert res.dataptr[i] == data[i]
+
+ def test_specialize_slice_2_0(self):
+ py.test.skip('not implemented')
+ def f():
+ a = numpy.zeros((12,), dtype='i').reshape((3,4))
+ a[:2, 0] = 1
+ return a
+ res = interpret(f, [])
+ data = [1,0,0,0,1,0,0,0,0,0,0,0]
+ for i in range(12):
+ assert res.dataptr[i] == data[i]
+
+ def test_specialize_slice_2_1(self):
+ def f():
+ a = numpy.zeros((12,), dtype='i').reshape((3,4))
+ a[:2, 0] = numpy.array([1,2])
+ a[1, 1:3] = numpy.array([4,5])
+ a[0:1, 3:] = numpy.array([6,])
+ return a
+ res = interpret(f, [])
+ data = [1, 0, 0, 6, 2, 4, 5, 0, 0, 0, 0, 0]
+ for i in range(12):
+ assert res.dataptr[i] == data[i]
+
+ def test_specialize_slice_2_2(self):
+ def f():
+ a = numpy.zeros((12,), dtype='i').reshape((3,4))
+ b = numpy.array([1,2,3,4]).reshape((2,2))
+ a[1:3, 2:] = b
+ return a
+ res = interpret(f, [])
+ data = [0, 0, 0, 0, 0, 0, 1, 2, 0, 0, 3, 4]
+ for i in range(12):
+ assert res.dataptr[i] == data[i]
+
def test_specialize_view(self):
def f(ii, jj):
a = numpy.zeros((4, 5))
@@ -354,6 +427,45 @@
assert interpret(f, [2, 3]) == 2
+ def test_specialize_view_implicit_slice(self):
+ def f():
+ a = numpy.array(range(12)).reshape((3,4))
+ b = a[0,]
+ return b
+
+ res = interpret(f, [])
+ for i in range(4):
+ assert res.dataptr[i] == i
+
+ def test_specialize_broadcast(self):
+ def f():
+ a = numpy.empty((4,3), dtype='i')
+ b = numpy.array([33])
+ a[:,:] = b
+ return a
+ res = interpret(f, [])
+ for i in range(4*3):
+ assert res.dataptr[i] == 33
+
+ def f():
+ a = numpy.empty((4,3), dtype='i')
+ b = numpy.array([33])
+ a[:,] = b
+ return a
+ res = interpret(f, [])
+ for i in range(4*3):
+ assert res.dataptr[i] == 33
+
+ def f():
+ a = numpy.empty((4,3,2), dtype='i')
+ a[:] = numpy.array([33])
+ a[0,:] = numpy.array([22])
+ return a
+ res = interpret(f, [])
+ data = [22]*6 + [33]*18
+ for i in range(3*4*2):
+ assert res.dataptr[i] == data[i]
+
def test_malloc_remove(self):
py.test.skip('this test requires _always_inline_ magic hook')
def f():
More information about the Pypy-commit
mailing list