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 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]