From 1f9c76ba2b10f092fdc3b6e8358fdf744d2c70c4 Mon Sep 17 00:00:00 2001 From: yuanfz <42092999+FZYUAN-1@users.noreply.github.com> Date: Sun, 30 Jan 2022 01:50:53 +0100 Subject: [PATCH] [Op][Topi] 5 ops can accept unsigned integers as indices (#10098) * tests passed * reformat * add uint test for unravel_index --- src/relay/op/tensor/transform.cc | 18 ++- tests/python/relay/test_op_level3.py | 127 +++++++++--------- .../python/topi/python/test_topi_transform.py | 30 +++-- 3 files changed, 95 insertions(+), 80 deletions(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index c3e36d4cbe07..645c7dc0071f 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1105,7 +1105,8 @@ bool ScatterRel(const Array& types, int num_inputs, const Attrs& attrs, if (updates == nullptr) { return false; } - ICHECK(indices->dtype.is_int()) << "indices of scatter must be tensor of integer"; + ICHECK(indices->dtype.is_int() || indices->dtype.is_uint()) + << "indices of scatter must be tensor of integer"; const auto param = attrs.as(); ICHECK(param != nullptr); reporter->Assign(types[3], TensorType(data->shape, data->dtype)); @@ -1152,7 +1153,8 @@ bool ScatterAddRel(const Array& types, int num_inputs, const Attrs& attrs, if (updates == nullptr) { return false; } - ICHECK(indices->dtype.is_int()) << "indices of scatter_add must be tensor of integer"; + ICHECK(indices->dtype.is_int() || indices->dtype.is_uint()) + << "indices of scatter_add must be tensor of integer"; const auto param = attrs.as(); ICHECK(param != nullptr); reporter->Assign(types[3], TensorType(data->shape, data->dtype)); @@ -1204,7 +1206,8 @@ bool ScatterNDRel(const Array& types, int num_inputs, const Attrs& attrs, << "ScatterND: expect updates type to be TensorType but got " << types[2]; return false; } - ICHECK(indices->dtype.is_int()) << "ScatterND: indices must be a tensor of integers."; + ICHECK(indices->dtype.is_int() || indices->dtype.is_uint()) + << "ScatterND: indices must be a tensor of integers."; const auto out_shape = data->shape; const IntImmNode* mdim = indices->shape[0].as(); @@ -3656,7 +3659,8 @@ bool UnRavelIndexRel(const Array& types, int num_inputs, const Attrs& attr << "unravel_index: expect input type to be TensorType but get " << types[0]; return false; } - ICHECK(indices->dtype.is_int()) << "indices of unravel_index must be tensor of integer"; + ICHECK(indices->dtype.is_int() || indices->dtype.is_uint()) + << "indices of unravel_index must be tensor of integer"; const auto* shape = types[1].as(); if (shape == nullptr) { @@ -3664,7 +3668,8 @@ bool UnRavelIndexRel(const Array& types, int num_inputs, const Attrs& attr << "unravel_index: expect input type to be TensorType but get " << types[1]; return false; } - ICHECK(indices->dtype.is_int()) << "shape of unravel_index must be tensor of integer"; + ICHECK(shape->dtype.is_int() || shape->dtype.is_uint()) + << "shape of unravel_index must be tensor of integer"; Array indices_shape; Array shape_shape; @@ -3894,7 +3899,8 @@ bool AdvIndexRel(const Array& types, int num_inputs, const Attrs& attrs, if (index_type == nullptr) { return false; } - ICHECK(index_type->dtype.is_int()) << "indices must be tensor of integers"; + ICHECK(index_type->dtype.is_int() || index_type->dtype.is_uint()) + << "indices must be tensor of integers"; int64_t flatten_len = 1; bool has_dyn_shape = false; diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index cff1c3ba3f93..a6eeaa6a8d2d 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -999,9 +999,9 @@ def ref_scatter(data, indices, updates, axis=0): def test_scatter(target, dev, executor_kind): - def verify_scatter(dshape, ishape, axis=0): + def verify_scatter(dshape, ishape, axis=0, indices_dtype="int64"): d = relay.var("d", relay.TensorType(dshape, "float32")) - i = relay.var("i", relay.TensorType(ishape, "int64")) + i = relay.var("i", relay.TensorType(ishape, indices_dtype)) u = relay.var("u", relay.TensorType(ishape, "float32")) z = relay.op.scatter(d, i, u, axis) @@ -1009,7 +1009,7 @@ def verify_scatter(dshape, ishape, axis=0): data_np = np.random.uniform(size=dshape).astype("float32") updates_np = np.random.uniform(size=ishape).astype("float32") - indices_np = np.random.randint(-dshape[axis], dshape[axis] - 1, ishape).astype("int64") + indices_np = np.random.randint(0, dshape[axis] - 1, ishape).astype(indices_dtype) ref_res = ref_scatter(data_np, indices_np, updates_np, axis) @@ -1031,6 +1031,7 @@ def verify_scatter(dshape, ishape, axis=0): verify_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1) verify_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2) verify_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3) + verify_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3, indices_dtype="uint32") class TestDynamicScatter: @@ -1073,27 +1074,28 @@ def test_dynamic_scatter(self, target, dev, executor_kind, dshape, ishape, axis) class TestScatterAdd: - dshape, ishape, axis, dtype = tvm.testing.parameters( - ((10,), (10,), 0, "int32"), - ((1000,), (1000,), 0, "int32"), - ((10, 5), (10, 5), -2, "float32"), - ((10, 5), (10, 5), -1, "float32"), - ((10, 5), (3, 5), 0, "float32"), - ((12, 4), (7, 2), 1, "float32"), - ((2, 3, 4), (1, 3, 4), 0, "float32"), - ((2, 3, 4), (2, 1, 4), 1, "float32"), - ((2, 3, 4), (2, 3, 1), 2, "float32"), - ((2, 3, 4, 5), (1, 3, 4, 5), 0, "float32"), - ((6, 3, 4, 5), (2, 3, 4, 5), 1, "float32"), - ((2, 3, 8, 5), (2, 3, 1, 1), 2, "float32"), - ((16, 16, 4, 5), (16, 16, 4, 5), 3, "float32"), + dshape, ishape, axis, dtype, indice_dtype = tvm.testing.parameters( + ((10,), (10,), 0, "int32", "int64"), + ((1000,), (1000,), 0, "int32", "int64"), + ((10, 5), (10, 5), -2, "float32", "int64"), + ((10, 5), (10, 5), -1, "float32", "int64"), + ((10, 5), (3, 5), 0, "float32", "int64"), + ((12, 4), (7, 2), 1, "float32", "int64"), + ((2, 3, 4), (1, 3, 4), 0, "float32", "int64"), + ((2, 3, 4), (2, 1, 4), 1, "float32", "int64"), + ((2, 3, 4), (2, 3, 1), 2, "float32", "int64"), + ((2, 3, 4, 5), (1, 3, 4, 5), 0, "float32", "int64"), + ((6, 3, 4, 5), (2, 3, 4, 5), 1, "float32", "int64"), + ((2, 3, 8, 5), (2, 3, 1, 1), 2, "float32", "int64"), + ((16, 16, 4, 5), (16, 16, 4, 5), 3, "float32", "int64"), + ((16, 16, 4, 5), (16, 16, 4, 5), 3, "float32", "uint32"), ) @tvm.testing.fixture(cache_return_value=True) - def ref_data(self, dshape, ishape, axis, dtype): + def ref_data(self, dshape, ishape, axis, dtype, indice_dtype): data_np = np.random.uniform(size=dshape).astype(dtype) updates_np = np.random.uniform(size=ishape).astype(dtype) - indices_np = np.random.randint(-dshape[axis], dshape[axis] - 1, ishape).astype("int64") + indices_np = np.random.randint(0, dshape[axis] - 1, ishape).astype(indice_dtype) out_np = np.copy(data_np) for index in np.ndindex(*indices_np.shape): @@ -1105,9 +1107,11 @@ def ref_data(self, dshape, ishape, axis, dtype): # Optimization can produce tir.atomic_add, not currently supported # on vulkan runtime. @tvm.testing.known_failing_targets("vulkan") - def test_scatter_add(self, target, dev, ref_data, dshape, ishape, axis, dtype): + def test_scatter_add(self, target, dev, ref_data, dshape, ishape, axis, dtype, indice_dtype): d = relay.var("d", relay.TensorType(shape=[relay.Any() for _ in dshape], dtype=dtype)) - i = relay.var("i", relay.TensorType(shape=[relay.Any() for _ in ishape], dtype="int64")) + i = relay.var( + "i", relay.TensorType(shape=[relay.Any() for _ in ishape], dtype=indice_dtype) + ) u = relay.var("u", relay.TensorType(shape=[relay.Any() for _ in ishape], dtype=dtype)) z = relay.op.scatter_add(d, i, u, axis) @@ -1955,47 +1959,48 @@ def verify_scatter_nd_with_stack( ) tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=rtol, atol=atol) - data = np.zeros((2, 2)).astype("int64") - indices = np.array([[1, 1, 0], [0, 1, 0]]) - updates = np.array([2, 3, 0]) - out = np.array([[0, 0], [2, 3]]) - verify_scatter_nd(data, indices, updates, out) - verify_scatter_nd_with_stack(data, indices, updates, out) - - data = np.zeros((2, 2, 2, 2)).astype("int64") - indices = np.array([[0, 1], [1, 1]]) - updates = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) - out = np.array([[[[0, 0], [0, 0]], [[1, 2], [3, 4]]], [[[0, 0], [0, 0]], [[5, 6], [7, 8]]]]) - verify_scatter_nd(data, indices, updates, out) - verify_scatter_nd_with_stack(data, indices, updates, out) - - indices = np.array([[1, 0, 0]]) - updates = np.reshape(np.arange(1560 * 3), (3, 1560)).astype("float32") - shape = (2, 1560) - data = np.zeros(shape).astype("float32") - out = data.copy() - out[1, :] += updates[0, :] - out[0, :] += updates[1, :] - out[0, :] += updates[2, :] - verify_scatter_nd(data, indices, updates, out, mode="add") - verify_scatter_nd_with_stack(data, indices, updates, out) - - for mode in ["add", "update"]: - indices = np.stack((np.random.randint(2, size=5), np.random.randint(7, size=5))).astype( - "int64" - ) - updates = np.ones((5, 3)).astype("float64") - shape = (2, 7, 3) - data = np.random.random(shape).astype("float64") + for indice_dtype in ["uint8", "uint16", "uint32"]: + data = np.zeros((2, 2)).astype("int64") + indices = np.array([[1, 1, 0], [0, 1, 0]]).astype(indice_dtype) + updates = np.array([2, 3, 0]) + out = np.array([[0, 0], [2, 3]]) + verify_scatter_nd(data, indices, updates, out) + verify_scatter_nd_with_stack(data, indices, updates, out) + + data = np.zeros((2, 2, 2, 2)).astype("int64") + indices = np.array([[0, 1], [1, 1]]).astype(indice_dtype) + updates = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) + out = np.array([[[[0, 0], [0, 0]], [[1, 2], [3, 4]]], [[[0, 0], [0, 0]], [[5, 6], [7, 8]]]]) + verify_scatter_nd(data, indices, updates, out) + verify_scatter_nd_with_stack(data, indices, updates, out) + + indices = np.array([[1, 0, 0]]).astype(indice_dtype) + updates = np.reshape(np.arange(1560 * 3), (3, 1560)).astype("float32") + shape = (2, 1560) + data = np.zeros(shape).astype("float32") out = data.copy() - for i in range(indices.shape[1]): - for j in range(updates.shape[1]): - if mode == "add": - out[indices[0, i], indices[1, i], j] += updates[i, j] - elif mode == "update": - out[indices[0, i], indices[1, i], j] = updates[i, j] - verify_scatter_nd(data, indices, updates, out, mode) - verify_scatter_nd_with_stack(data, indices, updates, out, mode) + out[1, :] += updates[0, :] + out[0, :] += updates[1, :] + out[0, :] += updates[2, :] + verify_scatter_nd(data, indices, updates, out, mode="add") + verify_scatter_nd_with_stack(data, indices, updates, out) + + for mode in ["add", "update"]: + indices = np.stack((np.random.randint(2, size=5), np.random.randint(7, size=5))).astype( + indice_dtype + ) + updates = np.ones((5, 3)).astype("float64") + shape = (2, 7, 3) + data = np.random.random(shape).astype("float64") + out = data.copy() + for i in range(indices.shape[1]): + for j in range(updates.shape[1]): + if mode == "add": + out[indices[0, i], indices[1, i], j] += updates[i, j] + elif mode == "update": + out[indices[0, i], indices[1, i], j] = updates[i, j] + verify_scatter_nd(data, indices, updates, out, mode) + verify_scatter_nd_with_stack(data, indices, updates, out, mode) def test_unique(target, dev): diff --git a/tests/python/topi/python/test_topi_transform.py b/tests/python/topi/python/test_topi_transform.py index e9de148e7fcf..d500b6697937 100644 --- a/tests/python/topi/python/test_topi_transform.py +++ b/tests/python/topi/python/test_topi_transform.py @@ -674,15 +674,15 @@ def check_device(target, dev): check_device(target, dev) -def verify_unravel_index(indices, shape, dtype): - x_data = np.array(indices).astype(dtype) +def verify_unravel_index(indices, shape, dtype, indice_dtype="int64"): + x_data = np.array(indices).astype(indice_dtype) y_data = np.array(shape).astype(dtype) if len(x_data.shape) == 1: dst_shape = [y_data.shape[0], x_data.shape[0]] else: dst_shape = [y_data.shape[0]] - X = te.placeholder(shape=x_data.shape, dtype=dtype, name="X") + X = te.placeholder(shape=x_data.shape, dtype=indice_dtype, name="X") Y = te.placeholder(shape=y_data.shape, dtype=dtype, name="Y") Z = topi.unravel_index(X, Y) @@ -771,7 +771,7 @@ def check_device(target, dev): check_device(target, dev) -def verify_adv_index(data_shape, index_shapes): +def verify_adv_index(data_shape, index_shapes, indice_dtype="int64"): dtype = "float32" data = te.placeholder(shape=data_shape, name="data", dtype=dtype) indices = [] @@ -779,8 +779,10 @@ def verify_adv_index(data_shape, index_shapes): np_indices = [] for i, index_shape in enumerate(index_shapes): limit = data_shape[i] - np_indices.append(np.random.uniform(0, limit - 1, size=index_shape).astype("int64")) - indices.append(te.placeholder(shape=index_shape, name="index_{}".format(i), dtype="int64")) + np_indices.append(np.random.uniform(0, limit - 1, size=index_shape).astype(indice_dtype)) + indices.append( + te.placeholder(shape=index_shape, name="index_{}".format(i), dtype=indice_dtype) + ) np_out = np_data[tuple(np_indices)] out = topi.adv_index(data, indices) @@ -1207,10 +1209,11 @@ def test_one_hot(): @tvm.testing.uses_gpu def test_unravel_index(): for dtype in ["int32", "int64"]: - verify_unravel_index([0, 1, 2, 3], [2, 2], dtype) - verify_unravel_index([144], [5, 5, 5, 2], dtype) - verify_unravel_index(144, [5, 5, 5, 2], dtype) - verify_unravel_index([100, 13, 5], [5, 5, 5, 2], dtype) + for indice_dtype in ["int64", "uint8", "uint16", "uint32"]: + verify_unravel_index([0, 1, 2, 3], [2, 2], dtype, indice_dtype) + verify_unravel_index([144], [5, 5, 5, 2], dtype, indice_dtype) + verify_unravel_index(144, [5, 5, 5, 2], dtype, indice_dtype) + verify_unravel_index([100, 13, 5], [5, 5, 5, 2], dtype, indice_dtype) @tvm.testing.uses_gpu @@ -1254,9 +1257,10 @@ def test_matrix_set_diag(): @tvm.testing.uses_gpu def test_adv_index(): - verify_adv_index((3, 4, 5), [(2,), (2,), (1,)]) - verify_adv_index((10, 15, 5), [(1, 1), (2, 7)]) - verify_adv_index((10, 5, 15), [(1, 2, 1), (1, 2, 7)]) + for indice_dtype in ["int32", "int64", "uint8", "uint16", "uint32"]: + verify_adv_index((3, 4, 5), [(2,), (2,), (1,)], indice_dtype=indice_dtype) + verify_adv_index((10, 15, 5), [(1, 1), (2, 7)], indice_dtype=indice_dtype) + verify_adv_index((10, 5, 15), [(1, 2, 1), (1, 2, 7)], indice_dtype=indice_dtype) if __name__ == "__main__":