def multidim_dot(space, left, right, result, dtype, right_critical_dim):
''' assumes left, right are concrete arrays
given left.shape == [3, 5, 7],
right.shape == [2, 7, 4]
then
result.shape == [3, 5, 2, 4]
broadcast shape should be [3, 5, 2, 7, 4]
result should skip dims 3 which is len(result_shape) - 1
(note that if right is 1d, result should
skip len(result_shape))
left should skip 2, 4 which is a.ndims-1 + range(right.ndims)
except where it==(right.ndims-2)
right should skip 0, 1
'''
broadcast_shape = left.shape[:-1] + right.shape
shapelen = len(broadcast_shape)
left_skip = [len(left.shape) - 1 + i for i in range(len(right.shape))
if i != right_critical_dim]
right_skip = range(len(left.shape) - 1)
result_skip = [len(result.shape) - (len(right.shape) > 1)]
_r = calculate_dot_strides(result.strides, result.backstrides,
broadcast_shape, result_skip)
outi = ViewIterator(result.start, _r[0], _r[1], broadcast_shape)
_r = calculate_dot_strides(left.strides, left.backstrides,
broadcast_shape, left_skip)
lefti = ViewIterator(left.start, _r[0], _r[1], broadcast_shape)
_r = calculate_dot_strides(right.strides, right.backstrides,
broadcast_shape, right_skip)
righti = ViewIterator(right.start, _r[0], _r[1], broadcast_shape)
while not outi.done():
dot_driver.jit_merge_point(left=left,
right=right,
shapelen=shapelen,
lefti=lefti,
righti=righti,
outi=outi,
result=result,
dtype=dtype,
)
lval = left.getitem(lefti.offset).convert_to(dtype)
rval = right.getitem(righti.offset).convert_to(dtype)
outval = result.getitem(outi.offset).convert_to(dtype)
v = dtype.itemtype.mul(lval, rval)
value = dtype.itemtype.add(v, outval).convert_to(dtype)
result.setitem(outi.offset, value)
outi = outi.next(shapelen)
righti = righti.next(shapelen)
lefti = lefti.next(shapelen)
return result
def test_C_viewiterator(self):
#Let's get started, simple iteration in C order with
#contiguous layout => strides[-1] is 1
start = 0
shape = [3, 5]
strides = [5, 1]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [10, 4]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next(2)
i = i.next(2)
i = i.next(2)
assert i.offset == 3
assert not i.done()
assert i.indices == [0,3]
#cause a dimension overflow
i = i.next(2)
i = i.next(2)
assert i.offset == 5
assert i.indices == [1,0]
#Now what happens if the array is transposed? strides[-1] != 1
# therefore layout is non-contiguous
strides = [1, 3]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [2, 12]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next(2)
i = i.next(2)
i = i.next(2)
assert i.offset == 9
assert not i.done()
assert i.indices == [0,3]
#cause a dimension overflow
i = i.next(2)
i = i.next(2)
assert i.offset == 1
assert i.indices == [1,0]
def allocate_iter(self, arr, transforms):
return ViewIterator(arr.start, arr.strides, arr.backstrides,
arr.shape).apply_transformations(arr, transforms)
def test_C_viewiterator_step(self):
#iteration in C order with #contiguous layout => strides[-1] is 1
#skip less than the shape
start = 0
shape = [3, 5]
strides = [5, 1]
backstrides = [x * (y - 1) for x, y in zip(strides, shape)]
assert backstrides == [10, 4]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next_skip_x(2, 2)
i = i.next_skip_x(2, 2)
i = i.next_skip_x(2, 2)
assert i.offset == 6
assert not i.done()
assert i.indices == [1, 1]
#And for some big skips
i = i.next_skip_x(2, 5)
assert i.offset == 11
assert i.indices == [2, 1]
i = i.next_skip_x(2, 5)
# Note: the offset does not overflow but recycles,
# this is good for broadcast
assert i.offset == 1
assert i.indices == [0, 1]
assert i.done()
#Now what happens if the array is transposed? strides[-1] != 1
# therefore layout is non-contiguous
strides = [1, 3]
backstrides = [x * (y - 1) for x, y in zip(strides, shape)]
assert backstrides == [2, 12]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next_skip_x(2, 2)
i = i.next_skip_x(2, 2)
i = i.next_skip_x(2, 2)
assert i.offset == 4
assert i.indices == [1, 1]
assert not i.done()
i = i.next_skip_x(2, 5)
assert i.offset == 5
assert i.indices == [2, 1]
assert not i.done()
i = i.next_skip_x(2, 5)
assert i.indices == [0, 1]
assert i.offset == 3
assert i.done()
def test_C_viewiterator(self):
#Let's get started, simple iteration in C order with
#contiguous layout => strides[-1] is 1
start = 0
shape = [3, 5]
strides = [5, 1]
backstrides = [x * (y - 1) for x, y in zip(strides, shape)]
assert backstrides == [10, 4]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next(2)
i = i.next(2)
i = i.next(2)
assert i.offset == 3
assert not i.done()
assert i.indices == [0, 3]
#cause a dimension overflow
i = i.next(2)
i = i.next(2)
assert i.offset == 5
assert i.indices == [1, 0]
#Now what happens if the array is transposed? strides[-1] != 1
# therefore layout is non-contiguous
strides = [1, 3]
backstrides = [x * (y - 1) for x, y in zip(strides, shape)]
assert backstrides == [2, 12]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next(2)
i = i.next(2)
i = i.next(2)
assert i.offset == 9
assert not i.done()
assert i.indices == [0, 3]
#cause a dimension overflow
i = i.next(2)
i = i.next(2)
assert i.offset == 1
assert i.indices == [1, 0]
def test_C_viewiterator_step(self):
#iteration in C order with #contiguous layout => strides[-1] is 1
#skip less than the shape
start = 0
shape = [3, 5]
strides = [5, 1]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [10, 4]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next_skip_x(2,2)
i = i.next_skip_x(2,2)
i = i.next_skip_x(2,2)
assert i.offset == 6
assert not i.done()
assert i.indices == [1,1]
#And for some big skips
i = i.next_skip_x(2,5)
assert i.offset == 11
assert i.indices == [2,1]
i = i.next_skip_x(2,5)
# Note: the offset does not overflow but recycles,
# this is good for broadcast
assert i.offset == 1
assert i.indices == [0,1]
assert i.done()
#Now what happens if the array is transposed? strides[-1] != 1
# therefore layout is non-contiguous
strides = [1, 3]
backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
assert backstrides == [2, 12]
i = ViewIterator(start, strides, backstrides, shape)
i = i.next_skip_x(2,2)
i = i.next_skip_x(2,2)
i = i.next_skip_x(2,2)
assert i.offset == 4
assert i.indices == [1,1]
assert not i.done()
i = i.next_skip_x(2,5)
assert i.offset == 5
assert i.indices == [2,1]
assert not i.done()
i = i.next_skip_x(2,5)
assert i.indices == [0,1]
assert i.offset == 3
assert i.done()