From 2ad9ebdd5fce45771524ddbb9420bb23a3077729 Mon Sep 17 00:00:00 2001 From: liuliang Date: Thu, 4 Mar 2021 16:33:01 +0800 Subject: [PATCH 1/3] [Relay] Fix relay op strategy for cuda dense int8 --- python/tvm/relay/op/strategy/cuda.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index 85bbab692574..e38a0e3576bd 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -655,7 +655,8 @@ def dense_strategy_cuda(attrs, inputs, out_type, target): data, weights = inputs b, i = get_const_tuple(data.shape) o, _ = get_const_tuple(weights.shape) - if out_type.dtype == "int8": + if data.dtype in ("int8", "uint8") and weights.dtype in ("int8", "uint8"): + assert data.dtype == weights.dtype strategy.add_implementation( wrap_compute_dense(topi.cuda.dense_int8), wrap_topi_schedule(topi.cuda.schedule_dense_int8), From 49c89853a3013ceff3184635e1f21e7e5cf10782 Mon Sep 17 00:00:00 2001 From: liuliang Date: Fri, 5 Mar 2021 13:28:45 +0800 Subject: [PATCH 2/3] Remove uint8 && Add autotvm task extraction test for relay graph that contains dense op (int8 * int8 -> int32) --- python/tvm/relay/op/strategy/cuda.py | 3 +-- .../relay/test_autotvm_task_extraction.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index e38a0e3576bd..bd25704d2e74 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -655,8 +655,7 @@ def dense_strategy_cuda(attrs, inputs, out_type, target): data, weights = inputs b, i = get_const_tuple(data.shape) o, _ = get_const_tuple(weights.shape) - if data.dtype in ("int8", "uint8") and weights.dtype in ("int8", "uint8"): - assert data.dtype == weights.dtype + if data.dtype == "int8" and weights.dtype == "int8" and out_type.dtype == "int32": strategy.add_implementation( wrap_compute_dense(topi.cuda.dense_int8), wrap_topi_schedule(topi.cuda.schedule_dense_int8), diff --git a/tests/python/relay/test_autotvm_task_extraction.py b/tests/python/relay/test_autotvm_task_extraction.py index d6bfd8d0ec11..712f87a70a08 100644 --- a/tests/python/relay/test_autotvm_task_extraction.py +++ b/tests/python/relay/test_autotvm_task_extraction.py @@ -102,5 +102,30 @@ def test_task_extraction(): assert len(tasks) == 31 +def test_task_extraction_for_dense_int8_cuda(): + target = 'cuda' + dense = relay.op.get("nn.dense") + + def get_net(batch, in_dim, out_dim, dtype, out_dtype): + data = tvm.relay.var("data", shape=[batch, in_dim], dtype=dtype) + weight = tvm.relay.var("weight", shape=[out_dim, in_dim], dtype=dtype) + out = relay.nn.dense(data, weight, out_dtype=out_dtype) + mod, params = relay.testing.create_workload(out) + return mod, params + + mod, params = get_net(1, 16, 32, 'float32', 'float32') + tasks = autotvm.task.extract_from_program( + mod, target=target, params=params, ops=(dense,) + ) + assert len(tasks) == 1 and tasks[0].name == 'dense_small_batch.cuda' + + mod, params = get_net(1, 16, 32, 'int8', 'int32') + tasks = autotvm.task.extract_from_program( + mod, target=target, params=params, ops=(dense,) + ) + assert len(tasks) == 1 and tasks[0].name == 'dense_int8.cuda' + + if __name__ == "__main__": test_task_extraction() + test_task_extraction_for_dense_int8_cuda() From 4baac18720a6e07fa52b353c86b7c0baa59a3280 Mon Sep 17 00:00:00 2001 From: liuliang Date: Fri, 5 Mar 2021 15:55:30 +0800 Subject: [PATCH 3/3] Reformat the code of test case --- .../relay/test_autotvm_task_extraction.py | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/tests/python/relay/test_autotvm_task_extraction.py b/tests/python/relay/test_autotvm_task_extraction.py index 712f87a70a08..b3f1868969cc 100644 --- a/tests/python/relay/test_autotvm_task_extraction.py +++ b/tests/python/relay/test_autotvm_task_extraction.py @@ -103,27 +103,23 @@ def test_task_extraction(): def test_task_extraction_for_dense_int8_cuda(): - target = 'cuda' + target = "cuda" dense = relay.op.get("nn.dense") def get_net(batch, in_dim, out_dim, dtype, out_dtype): - data = tvm.relay.var("data", shape=[batch, in_dim], dtype=dtype) + data = tvm.relay.var("data", shape=[batch, in_dim], dtype=dtype) weight = tvm.relay.var("weight", shape=[out_dim, in_dim], dtype=dtype) out = relay.nn.dense(data, weight, out_dtype=out_dtype) mod, params = relay.testing.create_workload(out) return mod, params - mod, params = get_net(1, 16, 32, 'float32', 'float32') - tasks = autotvm.task.extract_from_program( - mod, target=target, params=params, ops=(dense,) - ) - assert len(tasks) == 1 and tasks[0].name == 'dense_small_batch.cuda' + mod, params = get_net(1, 16, 32, "float32", "float32") + tasks = autotvm.task.extract_from_program(mod, target=target, params=params, ops=(dense,)) + assert len(tasks) == 1 and tasks[0].name == "dense_small_batch.cuda" - mod, params = get_net(1, 16, 32, 'int8', 'int32') - tasks = autotvm.task.extract_from_program( - mod, target=target, params=params, ops=(dense,) - ) - assert len(tasks) == 1 and tasks[0].name == 'dense_int8.cuda' + mod, params = get_net(1, 16, 32, "int8", "int32") + tasks = autotvm.task.extract_from_program(mod, target=target, params=params, ops=(dense,)) + assert len(tasks) == 1 and tasks[0].name == "dense_int8.cuda" if __name__ == "__main__":