From b454c7523f8880d1934b8c183402e01d7e2608ed Mon Sep 17 00:00:00 2001 From: ezw2 Date: Thu, 5 Aug 2021 12:44:33 -0400 Subject: [PATCH] initial attempt for supporting tensor slices --- python/heterocl/api.py | 24 +- python/heterocl/compute_api.py | 2 +- python/heterocl/tensor.py | 209 +++++--- python/heterocl/util.py | 30 +- tests/test_dsl_tensorslice.py | 839 +++++++++++++++++++++++++++++++++ 5 files changed, 1024 insertions(+), 80 deletions(-) create mode 100644 tests/test_dsl_tensorslice.py diff --git a/python/heterocl/api.py b/python/heterocl/api.py index 8986fafa8..a122b88ab 100644 --- a/python/heterocl/api.py +++ b/python/heterocl/api.py @@ -446,16 +446,34 @@ def print_val(val): stage = Stage.get_current() if isinstance(val, (Scalar, _expr.Expr, numbers.Number)): stage.emit(_make.Print([val], get_format(val) + "\n")) + elif isinstance(val, TensorSlice) \ + and sum(isinstance(x, slice) for x in val.indices) > 0: + nshape = len(val.tensor.shape) + startdim = len(val.indices) - sum(isinstance(x, slice) for x in val.indices) + ndim = nshape - startdim + args = ["print_"+str(n) for n in range(0, ndim)] + ivs = [_IterVar((0, val._shape[n]), args[n-startdim], 0) for n in range(startdim, nshape)] + if ndim >= 1: + import builtins + stage.emit(print_tensor(val, ivs, ndim-1, ndim)) + stage.emit(_make.Print([], "\n")) + else: + stage.emit(_make.Print([val], get_format(val) + "\n")) elif isinstance(val, TensorSlice) \ and len(val.indices) == len(val.tensor.shape): stage.emit(_make.Print([val], get_format(val) + "\n")) - else: # we are dealing with tensors + else: nshape = len(val.tensor.shape) ndim = nshape if isinstance(val, TensorSlice): ndim = nshape - len(val.indices) - args = ["print_"+str(n) for n in range(0, ndim)] - ivs = [_IterVar((0, val.tensor.shape[nshape-n-1]), args[n], 0) \ + startdim = len(val.indices) + args = ["print_"+str(n) for n in range(startdim, nshape)] + ivs = [_IterVar((0, val.tensor.shape[n]), args[n-startdim], 0) \ + for n in range(startdim, nshape)] + else: + args = ["print_"+str(n) for n in range(0, ndim)] + ivs = [_IterVar((0, val.tensor.shape[n]), args[n], 0) \ for n in range(0, ndim)] import builtins stage.emit(print_tensor(val, ivs, ndim-1, ndim)) diff --git a/python/heterocl/compute_api.py b/python/heterocl/compute_api.py index 620766fd5..65fb3c843 100644 --- a/python/heterocl/compute_api.py +++ b/python/heterocl/compute_api.py @@ -445,7 +445,7 @@ def copy(tensor, name=None, dtype=None): B3 = hcl.copy(nA, "B3") """ name = get_name("copy", name) - if isinstance(tensor, Tensor): + if isinstance(tensor, (Tensor, TensorSlice)): return compute( tensor.shape, lambda *args: tensor[args], diff --git a/python/heterocl/tensor.py b/python/heterocl/tensor.py index 89f602f8f..4fb5750b3 100644 --- a/python/heterocl/tensor.py +++ b/python/heterocl/tensor.py @@ -10,34 +10,28 @@ from . import util from . import debug from . import types +from .tvm.api import _IterVar class Scalar(NodeGeneric, _expr.ExprOp): """A non-mutable scalar. - This should be used by `heterocl.placeholder` only. Valid usages of accessing a scalar include direct access and bit operations. - Parameters ---------- var : Var A TVM variable - Attributes ---------- var : Var The wrapped TVM variable - dtype : Type The data type of the scalar - See Also -------- heterocl.placeholder - Examples -------- .. code-block:: python - # use () to specify it is a non-mutable scalar a = hcl.placeholder((), "a") # direct access @@ -77,42 +71,34 @@ def asnode(self): class TensorSlice(NodeGeneric, _expr.ExprOp): """A helper class for tensor operations. - Valid tensor accesses include: 1. getting an element from a tensor 2. bit operations on the element. We **do not** support operations on a slice of tensor. - Parameters ---------- tensor : Tensor The target tensor - indices : int or tuple of int The indices to access the tensor - Attributes ---------- tensor : Tensor The target tensor - indices : int or tuple of int The indices to access the tensor - dtype : Type The data type of the tensor - Examples -------- .. code-block:: python - A = hcl.placeholder((10,), "A") # get a single element a = A[5] + # get a slice of the tensor + a = A[1:6] # bit operations on a single element b = A[5][2] c = A[5][3:7] - - # not allowed: A[5:7] """ def __init__(self, tensor, indices, dtype=None): if not isinstance(indices, tuple): @@ -120,8 +106,45 @@ def __init__(self, tensor, indices, dtype=None): self.tensor = tensor self.indices = indices self._dtype = dtype if dtype is not None else self.tensor.dtype + nshape = len(self.tensor.shape) + self._shape = list(self.tensor.shape) + indices_slice = [] + shape_dict = {} + # the key is the tensor dimension being altered by the slice + # the value is a list of indices of the slices in self.indices + # corresponding to the tensor dimension + shape_idx = 0 + x = 0 + while shape_idx < nshape and x < len(self.indices): + i = self.indices[x] + if isinstance(i, slice): + if shape_idx not in shape_dict: + shape_dict[shape_idx] = [] + shape_dict[shape_idx].append(x) + indices_slice.append(i) + diff = i.stop - i.start + self._shape[shape_idx] = diff + if diff <= 0: + raise TensorError("Invalid index range") + else: + shape_idx += 1 + x += 1 + self._shape = tuple(self._shape) + # is true when individual elements of the tensor are being accessed + self.slice_offset = len(self.indices) - nshape - len(indices_slice) == 0 + + # offsets the dimensions that are sliced + if len(indices_slice) > 0 and self.slice_offset: + i = 0 + for key, value in shape_dict.items(): + for x in range(0, len(value)): + index_slice = value[x] + len(value) - x + tmp = list(self.indices) + tmp[index_slice] = self.indices[index_slice] + indices_slice[i].start + self.indices = tuple(tmp) + i += 1 + index, bit, _ = util.get_index(self.tensor.shape, self.indices, 0) # check if we have bit slicing - index, bit, _ = util.get_index(self.tensor.shape, indices, 0) if isinstance(bit, slice) and not isinstance(self.tensor.type, types.Struct): diff = bit.start - bit.stop if not isinstance(diff, int): @@ -147,16 +170,35 @@ def __getitem__(self, indices): def __setitem__(self, indices, expr): if not isinstance(indices, tuple): indices = (indices,) - indices = self.indices + indices indices = util.CastRemover().mutate(indices) - index, bit, _ = util.get_index(self.tensor.shape, indices, 0) + indices = self.indices + indices + off_shift = 0 + num_slice = 0 + # if the indices have not yet been shifted + if not self.slice_offset: + shift_index = False + idx = 0 + tmp = list(indices) + for i in indices: + if shift_index and not isinstance(i, slice): + tmp[idx] = indices[idx] + off_shift + shift_index = False + elif shift_index: + off_shift += i.start + num_slice += 1 + elif isinstance(i, slice): + shift_index = True + off_shift = i.start + num_slice += 1 + idx += 1 + indices = tuple(tmp) + index, bit, index_acc = util.get_index(self.tensor.shape, indices, 0) if not Stage.get_len(): raise TensorError("Cannot set tensor elements without compute APIs") builder = Stage.get_current() + emit_dtype = None if bit is None: - builder.emit(_make.Store(self.tensor.buf.data, - _make.Cast(self._dtype, expr), - index)) + emit_dtype = self._dtype elif isinstance(bit, slice): load = _make.Load(self.tensor.dtype, self.tensor.buf.data, index) # special handle for struct: we need to make sure the bitwidths @@ -167,14 +209,46 @@ def __setitem__(self, indices, expr): expr = _make.Call(ty, "bitcast", [expr], _expr.Call.PureIntrinsic, None, 0) expr = _make.SetSlice(load, expr, bit.start, bit.stop) - builder.emit(_make.Store(self.tensor.buf.data, - _make.Cast(self.tensor.dtype, expr), - index)) + emit_dtype = self.tensor.dtype else: load = _make.Load(self.tensor.dtype, self.tensor.buf.data, index) expr = _make.SetBit(load, expr, bit) + emit_dtype = self._dtype + + set_dim = len(indices) - num_slice + if isinstance(index, slice): + index = 0 + if isinstance(indices[-1], slice) and set_dim < len(self.tensor.shape)\ + and isinstance(expr, (TensorSlice, Tensor)): + st = _make.Add(_make.Mul(index, index_acc, False), + _make.Div(_make.Mul(off_shift, index_acc, False), + self.tensor.shape[set_dim], False), False) + acc = 1 + for x in self.tensor.shape[(set_dim + 1):]: + acc *= x + iv = [0 for n in range(set_dim, len(self.tensor.shape) - 1)] + iv.append(_IterVar((0, (indices[-1].stop - indices[-1].start) * acc), "set_array_var", 0)) + stmt = _make.Store(self.tensor._buf.data, + _make.Cast(emit_dtype, expr[tuple(iv)]), + iv[-1] + st) + for_stmt = util.make_for([iv[-1]], stmt, 0, "tensorslice_set_array") + builder.emit(for_stmt) + elif isinstance(indices[-1], slice) and set_dim < len(self.tensor.shape): + st = _make.Add(_make.Mul(index, index_acc, False), + _make.Div(_make.Mul(off_shift, index_acc, False), + self.tensor.shape[set_dim], False), False) + acc = 1 + for x in self.tensor.shape[(set_dim + 1):]: + acc *= x + iv = [_IterVar((0, (indices[-1].stop - indices[-1].start) * acc), "broadcast_var", 0)] + stmt = _make.Store(self.tensor._buf.data, + _make.Cast(emit_dtype, expr), + iv[0] + st) + for_stmt = util.make_for(iv, stmt, 0, "tensorslice_broadcast") + builder.emit(for_stmt) + else: builder.emit(_make.Store(self.tensor.buf.data, - _make.Cast(self._dtype, expr), + _make.Cast(emit_dtype, expr), index)) def __getattr__(self, key): @@ -204,6 +278,8 @@ def __getattr__(self, key): def __setattr__(self, key, expr): if key in ("tensor", "indices", "_dtype"): super().__setattr__(key, expr) + elif key == "_shape" or key == "slice_offset": + self.__dict__[key] = expr else: hcl_dtype = self.tensor.hcl_dtype if not isinstance(hcl_dtype, types.Struct): @@ -230,14 +306,14 @@ def dtype(self): @property def shape(self): - if len(self.indices) > len(self.tensor.shape): + num_slice = sum(isinstance(x, slice) for x in self.indices) + idx = len(self.indices) - num_slice + if idx > len(self.tensor.shape): raise TensorError("Shape is not defined when the length of indices" + " is greater than the number of dimensions") - return self.tensor.shape[len(self.indices):] + return self._shape[idx:] def asnode(self): - if len(self.indices) < len(self.tensor.shape): - raise TensorError("Accessing a slice of tensor is not allowed") self.indices = util.CastRemover().mutate(self.indices) index, bit, _ = util.get_index(self.tensor.shape, self.indices, 0) if bit is None: @@ -266,59 +342,42 @@ def asnode(self): class Tensor(NodeGeneric, _expr.ExprOp): """A HeteroCL tensor. - This is a wrapper for a TVM tensor. It should be generated from HeteroCL compute APIs. - Parameters ---------- shape : tuple of int The shape of the tensor - dtype : Type, optional The data type of the tensor - name : str, optional The name of the tensor - buf : Buffer, optional The TVM buffer of the tensor - Attributes ---------- dtype : Type The data type of the tensor - name : str The name of the tensor - var_dict : dict(str, Var) A dictionary that maps between a name and a variable - first_update : Stage The first stage that updates the tensor - last_update : Stage The last stage that updates the tensor - tensor : Operation The TVM tensor - buf : Buffer The TVM buffer - type : Type The data type in HeteroCL format - op : Stmt The operation statement - axis : list of IterVar A list of axes of the tensor - v : Expr Syntactic sugar to access the element of an single-element tensor - See Also -------- heterocl.placeholder, heterocl.compute @@ -357,29 +416,37 @@ def __setitem__(self, indices, expr): if not isinstance(indices, tuple): indices = (indices,) indices = util.CastRemover().mutate(indices) - if len(indices) < len(self.shape): - raise TensorError("Accessing a slice of tensor is not allowed") - else: - index, bit, _ = util.get_index(self.shape, indices, 0) - if not Stage.get_len(): + index, bit, index_acc = util.get_index(self.shape, indices, 0) + if not Stage.get_len(): raise TensorError("Cannot set tensor elements without compute APIs") - builder = Stage.get_current() - if bit is None: - builder.emit(_make.Store(self.buf.data, - _make.Cast(self.dtype, expr), - index)) - elif isinstance(bit, slice): - load = _make.Load(self.tensor.dtype, self.tensor.buf.data, index) - expr = _make.SetSlice(load, expr, bit.start, bit.stop) - builder.emit(_make.Store(self.tensor.buf.data, - _make.Cast(self.tensor.dtype, expr), - index)) + builder = Stage.get_current() + if isinstance(index, slice): + r_acc = _make.Div(index_acc, self.shape[0], False) + iv_bound = _make.Mul((index.stop-index.start), r_acc, False) + ivs = [_IterVar((0, iv_bound), "tensor_broadcast", 0)] + stmt_index = _make.Add(ivs[0], _make.Mul(index.start, r_acc, False), False) + if isinstance(expr, (TensorSlice, Tensor)): + stmt = _make.Store(self._buf.data, expr[ivs[0]], stmt_index) else: - load = _make.Load(self.tensor.dtype, self.tensor.buf.data, index) - expr = _make.SetBit(load, expr, bit) - builder.emit(_make.Store(self.tensor.buf.data, - _make.Cast(self.tensor.dtype, expr), - index)) + stmt = _make.Store(self._buf.data, expr, stmt_index) + for_stmt = util.make_for(ivs, stmt, 0, "tensor_set_array_loop") + builder.emit(for_stmt) + elif bit is None: + builder.emit(_make.Store(self.buf.data, + _make.Cast(self.dtype, expr), + index)) + elif isinstance(bit, slice): + load = _make.Load(self.tensor.dtype, self.tensor.buf.data, index) + expr = _make.SetSlice(load, expr, bit.start, bit.stop) + builder.emit(_make.Store(self.tensor.buf.data, + _make.Cast(self.tensor.dtype, expr), + index)) + else: + load = _make.Load(self.tensor.dtype, self.tensor.buf.data, index) + expr = _make.SetBit(load, expr, bit) + builder.emit(_make.Store(self.tensor.buf.data, + _make.Cast(self.tensor.dtype, expr), + index)) @property def tensor(self): @@ -411,7 +478,6 @@ def v(self): @buf.setter def buf(self, buf): """Set the TVM buffer. - Parameters ---------- buf : Buffer @@ -422,7 +488,6 @@ def buf(self, buf): @tensor.setter def tensor(self, tensor): """Set the TVM tensor. - Parameters ---------- tensor : Tensor @@ -432,9 +497,7 @@ def tensor(self, tensor): @v.setter def v(self, value): """A syntactic sugar for setting the value of a single-element tensor. - This is the same as using `a[0]=value`, where a is a single-element tensor. - Parameters ---------- value : Expr diff --git a/python/heterocl/util.py b/python/heterocl/util.py index 9ef069007..873400480 100644 --- a/python/heterocl/util.py +++ b/python/heterocl/util.py @@ -92,18 +92,42 @@ def make_for(indices, body, level, stage_name=""): body = _make.AttrStmt(iter_var, "loop_scope", iter_var.var, make_for(indices, body, level+1, stage_name)) return _make.For(iter_var.var, iter_var.dom.min, iter_var.dom.extent, 0, 0, body, ["stage_name"], [stage_name]) -# return (index, bit, _) +# return (index, bit, index_acc) def get_index(shape, args, level): + if any(map(lambda i: isinstance(i, slice), args[:len(shape)])): + tmp = [] + acc = 1 + idx = 0 + slice_idx = 0 + for a in args: + if not isinstance(a, slice) or len(tmp) == len(shape): + tmp.append(args[idx]) + elif isinstance(a, slice): + slice_idx = idx + idx += 1 + num_slices = len(args) - len(tmp) + for x in range(slice_idx - num_slices + 1, len(shape)): + acc = _make.Mul(acc, shape[x], False) + if len(tmp) == 0: + return (args[-1], None, acc) + ret = list(get_index_inner(shape, tuple(tmp), level)) + ret[2] = acc + return tuple(ret) + else: + return get_index_inner(shape, args, level) + +# return (index, bit, _) +def get_index_inner(shape, args, level): if level == len(args) - 1: # the last arg if level == len(shape): # bit-selection return (0, args[level], 1) else: return (args[level], None, shape[level]) else: - index = get_index(shape, args, level+1) + index = get_index_inner(shape, args, level+1) new_arg = args[level] new_index = _make.Add(index[0], - _make.Mul(new_arg, index[2], False), False) + _make.Mul(new_arg, index[2], False), False) new_acc = _make.Mul(index[2], shape[level], False) return (new_index, index[1], new_acc) diff --git a/tests/test_dsl_tensorslice.py b/tests/test_dsl_tensorslice.py new file mode 100644 index 000000000..0d5536e35 --- /dev/null +++ b/tests/test_dsl_tensorslice.py @@ -0,0 +1,839 @@ +import heterocl as hcl +import numpy as np + +def test_1D_basic(): + + hcl.init() + + def kernel(A): + matrix_C = A[2:7] + return hcl.compute((5,), lambda x: matrix_C[x]) + + A = hcl.placeholder((10,)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10,)) + np_B = np.zeros(5) + golden = np_A[2:7] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_1D_copy(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[2:7]) + A[2:7][0] = 123 + hcl.assert_(A[2] == 123) + return hcl.compute((5,), lambda x: matrix_C[x]) + + A = hcl.placeholder((10,)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10,)) + np_B = np.zeros(5) + golden = np_A[2:7] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_1D_broadcast(): + + hcl.init() + + def kernel(A): + A[2:7] = 999 + return hcl.compute((10,), lambda x: A[x]) + + A = hcl.placeholder((10,)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10,)) + np_B = np.zeros(10) + + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + golden = np_A + for x in range(2, 7): + golden[x] = 999 + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_1D_slice(): + + hcl.init() + + def kernel(A): + matrix_C = A[7][3][2:6] + return hcl.compute((4,), lambda x: matrix_C[x]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros(4) + golden = np_A[7][3][2:6] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_2D_slice(): + + hcl.init() + + def kernel(A): + matrix_C = A[7][2:6] + return hcl.compute((4,8), lambda x, y: matrix_C[x][y]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((4, 8)) + golden = np_A[7][2:6] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_3D_slice(): + + hcl.init() + + def kernel(A): + matrix_C = A[2:6] + return hcl.compute((4, 9, 8), lambda x, y, z: matrix_C[x][y][z]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((4, 9, 8)) + golden = np_A[2:6] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_3D_copy_dim2(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[8][1][2:7]) + A[8][1][2:7][0] = 123 + hcl.assert_(A[8][1][2] == 123) + return hcl.compute((5,), lambda x: matrix_C[x]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros(5) + golden = np_A[8][1][2:7] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_3D_copy_dim1_setslice(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[8][2:7]) + A[8][3][3:8] = A[2][6][1:6] + with hcl.for_(0, 5) as i: + hcl.assert_(A[8][3][i + 3] == A[2][6][i + 1]) + return hcl.compute((5,8), lambda x, y: matrix_C[x][y]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((5, 8)) + golden = np_A[8][2:7] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_3D_copy_dim0_broadcast(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[2:7]) + A[1:4] = 999 + with hcl.for_(0, 3) as x: + with hcl.for_(0, 9) as y: + with hcl.for_(0, 8) as z: + hcl.assert_(A[x + 1][y][z] == 999) + return hcl.compute((5, 9, 8), lambda x, y, z: matrix_C[x][y][z]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((5, 9, 8)) + golden = np_A[2:7] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_3D_copy_dim0(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[0]) + A[1:4] = 999 + with hcl.for_(0, 3) as x: + with hcl.for_(0, 9) as y: + with hcl.for_(0, 8) as z: + hcl.assert_(A[x + 1][y][z] == 999) + return hcl.compute((9, 8), lambda x, y: matrix_C[x][y]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((9, 8)) + golden = np_A[0] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_5D_broadcast(): + + hcl.init() + + def kernel(A): + A[1:3] = 999 + return hcl.compute((4, 3, 2, 3, 2), lambda x, y, z, a, b: A[x][y][z][a][b]) + + A = hcl.placeholder((4, 3, 2, 3, 2)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(4, 3, 2, 3, 2)) + np_B = np.zeros((4, 3, 2, 3, 2)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + golden = np_A + for x in range(1, 3): + for y in range(0, 3): + for z in range(0, 2): + for a in range(0, 3): + for b in range(0, 2): + golden[x][y][z][a][b] = 999 + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_nested_slice_dim1(): + + hcl.init() + + def kernel(A): + return hcl.compute((10, 3, 8), lambda x, y, z: A[x][2:9][3:6][y][z]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((10, 3, 8)) + golden = np.zeros((10, 3, 8)) + for x in range(0, 10): + for y in range(0, 3): + for z in range(0, 8): + golden[x][y][z] = np_A[x][y+5][z] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + golden = np.zeros((10, 3, 8)) + for x in range(0, 10): + for y in range(0, 3): + for z in range(0, 8): + golden[x][y][z] = np_A[x][y+5][z] + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_nested_slice_dim2(): + + hcl.init() + + def kernel(A): + return hcl.compute((10, 9, 3), lambda x, y, z: A[x][y][2:8][3:6][z]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((10, 9, 3)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + golden = np.zeros((10, 9, 3)) + for x in range(0, 10): + for y in range(0, 9): + for z in range(0, 3): + golden[x][y][z] = np_A[x][y][z+5] + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_nested_slice_dim0(): + + hcl.init() + + def kernel(A): + return hcl.compute((3, 9, 8), lambda x, y, z: A[2:8][3:6][x][y][z]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((3, 9, 8)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + golden = np.zeros((3, 9, 8)) + for x in range(0, 3): + for y in range(0, 9): + for z in range(0, 8): + golden[x][y][z] = np_A[x+5][y][z] + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_double_slice(): + + hcl.init() + + def kernel(A): + matrix_B = A[1][2:7][3:5] + return hcl.compute((2, 8), lambda x, y: matrix_B[x][y]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((2, 8)) + golden = np_A[1][2:7][3:5] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_get_bitslice_1D_tensor(): + + hcl.init() + + def kernel(A): + matrix_B = A[3:7] + return hcl.compute((4,), lambda x: matrix_B[x][2:0]) + + A = hcl.placeholder((10,)) + s = hcl.create_schedule(A, kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10,)) + np_B = np.zeros(4) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + golden = np_A[3:7] & 0b11 + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_get_bitslice_3D_tensorslice(): + + hcl.init() + + def kernel(A): + return hcl.compute((4, 9, 8), lambda x, y, z: A[3:7][x][y][z][3:1]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule(A, kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((4, 9, 8)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + golden = (np_A[3:7] & 0b110) >> 1 + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_get_bitslice_1D_tensorslice(): + + hcl.init() + + def kernel(A): + matrix_B = A[6][1][2:7] + return hcl.compute((5,), lambda x: matrix_B[x][0:8]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule(A, kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((5,)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + golden = np_A[6][1][2:7] & 0xFF + golden = golden.astype('uint8') + + ret = hcl_B.asnumpy() + ret = ret.astype('uint8') + for i in range(0, 5): + x = np.unpackbits(golden[i]) + x = np.flip(x) + y = np.unpackbits(ret[i]) + assert np.array_equal(x, y) + +def test_set_bitslice_1D_tensorslice(): + + hcl.init() + + def kernel(A, B): + matrix_C = B[1][2:5][2][4][1:11] + with hcl.for_(0, 10) as i: + matrix_C[i][2:0] = A[i] + return hcl.compute((10,), lambda x: matrix_C[x]) + + A = hcl.placeholder((10,)) + B = hcl.placeholder((3, 6, 5, 13)) + C = hcl.placeholder((10,)) + s = hcl.create_schedule([A, B], kernel) + f = hcl.build(s) + + np_A = np.random.randint(2, size=(10,)) + np_B = np.random.randint(10, size=(3, 6, 5, 13)) + np_C = np.zeros(10) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + hcl_C = hcl.asarray(np_C) + + f(hcl_A, hcl_B, hcl_C) + + golden = (np_B[1][2:5][2][4][1:11] & 0b1100) | np_A + ret = hcl_C.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_bitslice_2D_tensorslice(): + + hcl.init() + + def kernel(A, B): + matrix_C = B[1][2:5][2][3:7] + with hcl.for_(0, 5) as x: + with hcl.for_(0, 10) as y: + matrix_C[x][y][2:0] = A[x][y] + return hcl.compute((5,10), lambda x, y: matrix_C[x][y]) + + A = hcl.placeholder((5, 10)) + B = hcl.placeholder((3, 6, 8, 10)) + C = hcl.placeholder((5, 10)) + s = hcl.create_schedule([A, B], kernel) + f = hcl.build(s) + + np_A = np.random.randint(2, size=(5, 10)) + np_B = np.random.randint(10, size=(3, 6, 8, 10)) + np_C = np.zeros((5, 10)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + hcl_C = hcl.asarray(np_C) + + f(hcl_A, hcl_B, hcl_C) + + golden = (np_B[1][2:5][2][3:8] & 0b1100) | np_A + ret = hcl_C.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_1D(): + + hcl.init() + + def kernel(A): + matrix_B = A[2:8] + matrix_B[3] = 999 + return hcl.compute((6,), lambda x: matrix_B[x]) + + A = hcl.placeholder((10,)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10,)) + np_B = np.zeros(6) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + golden = np_A[2:8] + golden[3] = 999 + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_3D(): + + hcl.init() + + def kernel(A): + matrix_B = A[3:8] + matrix_B[2][7][4] = 999 + return hcl.compute((5, 9, 8), lambda x, y, z: matrix_B[x][y][z]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((5, 9, 8)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + golden = np_A[3:8] + golden[2][7][4] = 999 + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_2D(): + + hcl.init() + + def kernel(A): + matrix_B = A[0][3:8] + matrix_B[2][4] = 999 + return hcl.compute((5, 8), lambda x, y: matrix_B[x][y]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((5, 8)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + golden = np_A[0][3:8] + golden[2][4] = 999 + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_1D_tensorslice(): + + hcl.init() + + def kernel(A): + matrix_B = A[0][1][3:8] + matrix_B[4] = 999 + return hcl.compute((5,), lambda x: matrix_B[x]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros(5) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + golden = np_A[0][1][3:8] + golden[4] = 999 + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_multi_slice(): + + hcl.init() + + def kernel(A): + return hcl.compute((3,8), lambda x, y: A[2:7][1][3:6][x][y]) + + A = hcl.placeholder((10, 9, 8)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8)) + np_B = np.zeros((3, 8)) + golden = np_A[2:7][1][3:6] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_multislice_copy_3D(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[5][2:7][4][1:3]) + A[5][2:7][4][1:3] = 999 + with hcl.for_(0, 2) as x: + with hcl.for_(0, 11) as y: + with hcl.for_(0, 12) as z: + hcl.assert_(A[5][6][x+1][y][z] == 999) + return hcl.compute((2,11,12), lambda x,y,z: matrix_C[x][y][z]) + + A = hcl.placeholder((6, 9, 8, 11, 12)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(10, 9, 8, 11, 12)) + np_B = np.zeros((2,11,12)) + golden = np_A[5][2:7][4][1:3] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_multislice_copy_4D(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[5][2:7][1:3]) + A[5][2:7][1:3] = 999 + with hcl.for_(0, 2) as x: + with hcl.for_(0, 4) as y: + with hcl.for_(0, 3) as z: + with hcl.for_(0, 2) as a: + hcl.assert_(A[5][x+3][y][z][a] == 999) + return hcl.compute((2, 4, 3, 2), lambda x,y,z,a: matrix_C[x][y][z][a]) + + A = hcl.placeholder((8, 7, 4, 3, 2)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(8, 7, 4, 3, 2)) + np_B = np.zeros((2, 4, 3, 2)) + golden = np_A[5][2:7][1:3] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_multislice_copy_5D(): + + hcl.init() + + def kernel(A): + matrix_C = hcl.copy(A[2:7][1:3]) + A[2:7][1:3] = 999 + with hcl.for_(0, 2) as x: + with hcl.for_(0, 3) as y: + with hcl.for_(0, 4) as z: + with hcl.for_(0, 3) as a: + with hcl.for_(0, 2) as b: + hcl.assert_(A[x+3][y][z][a][b] == 999) + return hcl.compute((2, 3, 4, 3, 2), lambda x,y,z,a,b: matrix_C[x][y][z][a][b]) + + A = hcl.placeholder((8, 3, 4, 3, 2)) + s = hcl.create_schedule([A], kernel) + f = hcl.build(s) + + np_A = np.random.randint(10, size=(8, 3, 4, 3, 2)) + np_B = np.zeros((2, 3, 4, 3, 2)) + golden = np_A[2:7][1:3] + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + + f(hcl_A, hcl_B) + + ret = hcl_B.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_tensorslice_1D_nested(): + + hcl.init() + + def kernel(A, B): + A[1:5] = B[1][3:5][0][2][4:8] + return hcl.compute((10,), lambda x: A[x]) + + A = hcl.placeholder((10,)) + B = hcl.placeholder((3, 6, 5, 13)) + C = hcl.placeholder((10,)) + s = hcl.create_schedule([A, B], kernel) + f = hcl.build(s) + + np_A = np.random.randint(20, size=(10,)) + np_B = np.random.randint(10, size=(3, 6, 5, 13)) + np_C = np.zeros(10) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + hcl_C = hcl.asarray(np_C) + f(hcl_A, hcl_B, hcl_C) + + golden = np_A + golden[1:5] = np_B[1][3:5][0][2][4:8] + ret = hcl_C.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_tensorslice_3D(): + + hcl.init() + + def kernel(A, B): + A[1][2][2:6] = B[1:5] + return hcl.compute((2,4,8,3,4), lambda x,y,z,a,b: A[x][y][z][a][b]) + + A = hcl.placeholder((2, 4, 8, 3, 4)) + B = hcl.placeholder((6,3, 4)) + C = hcl.placeholder((2,4,8,3, 4)) + s = hcl.create_schedule([A, B], kernel) + f = hcl.build(s) + + np_A = np.random.randint(20, size=(2, 4, 8, 3, 4)) + np_B = np.random.randint(10, size=(6,3, 4)) + np_C = np.zeros((2, 4, 8, 3, 4)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + hcl_C = hcl.asarray(np_C) + + f(hcl_A, hcl_B, hcl_C) + + golden = np_A + golden[1][2][2:6] = np_B[1:5] + ret = hcl_C.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_tensorslice_3D_nested(): + + hcl.init() + + def kernel(A, B): + A[1][2][2:6][3][1:2] = B[2:5][1][1:2] + return hcl.compute((2,4,8,3,4), lambda x,y,z,a,b: A[x][y][z][a][b]) + + A = hcl.placeholder((2, 4, 8, 3, 4)) + B = hcl.placeholder((6,3, 4)) + C = hcl.placeholder((2,4,8,3, 4)) + s = hcl.create_schedule([A, B], kernel) + f = hcl.build(s) + + np_A = np.random.randint(20, size=(2, 4, 8, 3, 4)) + np_B = np.random.randint(10, size=(6,3, 4)) + np_C = np.zeros((2, 4, 8, 3, 4)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + hcl_C = hcl.asarray(np_C) + + f(hcl_A, hcl_B, hcl_C) + + golden = np_A + golden[1][2][2:6][3][1:2] = np_B[2:5][1][1:2] + ret = hcl_C.asnumpy() + assert np.array_equal(golden, ret) + +def test_set_tensorslice_2D_tensor(): + + hcl.init() + + def kernel(A, B): + A[1][2][2:6][3][1:2] = B + return hcl.compute((2,4,8,3,4), lambda x,y,z,a,b: A[x][y][z][a][b]) + + A = hcl.placeholder((2, 4, 8, 3, 4)) + B = hcl.placeholder((1, 4)) + C = hcl.placeholder((2,4,8,3, 4)) + s = hcl.create_schedule([A, B], kernel) + f = hcl.build(s) + + np_A = np.random.randint(20, size=(2, 4, 8, 3, 4)) + np_B = np.random.randint(10, size=(1, 4)) + np_C = np.zeros((2, 4, 8, 3, 4)) + hcl_A = hcl.asarray(np_A) + hcl_B = hcl.asarray(np_B) + hcl_C = hcl.asarray(np_C) + + f(hcl_A, hcl_B, hcl_C) + + golden = np_A + golden[1][2][2:6][3][1:2] = np_B + ret = hcl_C.asnumpy() + assert np.array_equal(golden, ret)