From 32c490dbdeead86ca72c83c97fb13a063d606137 Mon Sep 17 00:00:00 2001 From: Haozheng Fan Date: Fri, 18 Dec 2020 12:28:57 +0800 Subject: [PATCH 1/4] complex reduce --- python/tvm/topi/cuda/reduction.py | 6 +++-- tests/python/topi/python/test_topi_reduce.py | 26 ++++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/python/tvm/topi/cuda/reduction.py b/python/tvm/topi/cuda/reduction.py index ee868ac9b639..c194fd611513 100644 --- a/python/tvm/topi/cuda/reduction.py +++ b/python/tvm/topi/cuda/reduction.py @@ -130,12 +130,14 @@ def traverse_after_reduce(operator): for tensor in operator.input_tensors: traverse_after_reduce(tensor.op) elif operator.tag == "comm_reduce": - _schedule_reduce(operator, sch, is_idx_reduce=False) + if operator not in scheduled_ops: + _schedule_reduce(operator, sch, is_idx_reduce=False) for tensor in operator.input_tensors: if tensor.op not in scheduled_ops: traverse_before_reduce(tensor.op) elif operator.tag == "comm_reduce_idx": - _schedule_reduce(operator, sch, is_idx_reduce=True) + if operator not in scheduled_ops: + _schedule_reduce(operator, sch, is_idx_reduce=True) input_tensors = operator.input_tensors[0].op.input_tensors for tensor in input_tensors: if tensor.op not in scheduled_ops: diff --git a/tests/python/topi/python/test_topi_reduce.py b/tests/python/topi/python/test_topi_reduce.py index daf380d9f4e6..24d5968268bb 100644 --- a/tests/python/topi/python/test_topi_reduce.py +++ b/tests/python/topi/python/test_topi_reduce.py @@ -152,5 +152,31 @@ def test_reduce_map(): ) +@tvm.testing.uses_gpu +def test_complex_reduce(): + in_shape = (2, 3) + dtype = "float32" + axis = 0 + keep_dims = False + A = te.placeholder(shape=in_shape, name="A", dtype=dtype) + B = topi.sum(A1, axis=axis, keepdims=keep_dims) + C = topi.add(B, B) + D = topi.multiply(B, B) + E = topi.add(C, D) + for device, ctx in tvm.testing.enabled_targets(): + print("Running on target: %s" % device) + with tvm.target.Target(device): + s = tvm.topi.testing.get_reduce_schedule(device)(B) + foo = tvm.build(s, [A, E], device, name="sum") + in_npy = np.random.uniform(-1, 1, size=in_shape).astype(dtype) + sum_npy = in_npy.sum(axis=axis, keepdims=keepdims) + out_npy = sum_npy * 2 + sum_npy * sum_npy + data_tvm = tvm.nd.array(in_npy, ctx=ctx) + out_tvm = tvm.nd.empty(shape=out_npy.shape, ctx=ctx, dtype=dtype) + foo(data_tvm, out_tvm) + tvm.testing.assert_allclose(out_tvm.asnumpy(), out_npy, 1e-3, 1e-3) + + if __name__ == "__main__": test_reduce_map() + test_complex_reduce() From b7aad1f65e5a2a5bb75fc56b4a5bd36d54ff1819 Mon Sep 17 00:00:00 2001 From: Haozheng Fan Date: Fri, 18 Dec 2020 04:49:00 +0000 Subject: [PATCH 2/4] fix --- tests/python/topi/python/test_topi_reduce.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/python/topi/python/test_topi_reduce.py b/tests/python/topi/python/test_topi_reduce.py index 24d5968268bb..a484d416e595 100644 --- a/tests/python/topi/python/test_topi_reduce.py +++ b/tests/python/topi/python/test_topi_reduce.py @@ -157,16 +157,16 @@ def test_complex_reduce(): in_shape = (2, 3) dtype = "float32" axis = 0 - keep_dims = False + keepdims = False A = te.placeholder(shape=in_shape, name="A", dtype=dtype) - B = topi.sum(A1, axis=axis, keepdims=keep_dims) + B = topi.sum(A, axis=axis, keepdims=keepdims) C = topi.add(B, B) D = topi.multiply(B, B) E = topi.add(C, D) for device, ctx in tvm.testing.enabled_targets(): print("Running on target: %s" % device) with tvm.target.Target(device): - s = tvm.topi.testing.get_reduce_schedule(device)(B) + s = tvm.topi.testing.get_reduce_schedule(device)(E) foo = tvm.build(s, [A, E], device, name="sum") in_npy = np.random.uniform(-1, 1, size=in_shape).astype(dtype) sum_npy = in_npy.sum(axis=axis, keepdims=keepdims) @@ -178,5 +178,5 @@ def test_complex_reduce(): if __name__ == "__main__": - test_reduce_map() + # test_reduce_map() test_complex_reduce() From f4f5ca5ef7dfe66e4b6969df80f7b42302e965da Mon Sep 17 00:00:00 2001 From: Haozheng Fan Date: Fri, 18 Dec 2020 14:33:48 +0800 Subject: [PATCH 3/4] fix --- python/tvm/topi/cuda/reduction.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/topi/cuda/reduction.py b/python/tvm/topi/cuda/reduction.py index c194fd611513..ceab71640533 100644 --- a/python/tvm/topi/cuda/reduction.py +++ b/python/tvm/topi/cuda/reduction.py @@ -149,5 +149,6 @@ def traverse_after_reduce(operator): scheduled_ops.append(operator) - traverse_after_reduce(outs[0].op) + for out in outs: + traverse_after_reduce(out.op) return sch From 29beba516300b2cbd2788411d1e1c46e3419ff84 Mon Sep 17 00:00:00 2001 From: Haozheng Fan Date: Fri, 18 Dec 2020 14:40:59 +0800 Subject: [PATCH 4/4] fix --- tests/python/topi/python/test_topi_reduce.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/topi/python/test_topi_reduce.py b/tests/python/topi/python/test_topi_reduce.py index a484d416e595..9ddcb0d3884b 100644 --- a/tests/python/topi/python/test_topi_reduce.py +++ b/tests/python/topi/python/test_topi_reduce.py @@ -178,5 +178,5 @@ def test_complex_reduce(): if __name__ == "__main__": - # test_reduce_map() + test_reduce_map() test_complex_reduce()