From a5847eb60d13c08d9e64a107971f22e80d85f117 Mon Sep 17 00:00:00 2001 From: mengceng he Date: Mon, 13 Jun 2022 13:44:13 +0800 Subject: [PATCH 1/2] fix autotvm int8 dense vnni extraction --- python/tvm/autotvm/task/relay_integration.py | 4 ++++ tests/python/relay/test_autotvm_task_extraction.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py index 11f40ed62756..c5e7694017a1 100644 --- a/python/tvm/autotvm/task/relay_integration.py +++ b/python/tvm/autotvm/task/relay_integration.py @@ -130,6 +130,10 @@ def extract_from_multiple_program(mods, params, target, target_host=None, ops=No assert isinstance( mod, tvm.IRModule ), "only support relay Module or Function to be tuned" + + with target: + mod = tvm.relay.transform.AlterOpLayout()(mod) + relay.backend.te_compiler.get().clear() # wrap build call in thread to avoid multiprocessing problems build_thread = threading.Thread(target=_lower, args=(mod, target, param)) diff --git a/tests/python/relay/test_autotvm_task_extraction.py b/tests/python/relay/test_autotvm_task_extraction.py index 83480a044f45..1ebbdb3fd294 100644 --- a/tests/python/relay/test_autotvm_task_extraction.py +++ b/tests/python/relay/test_autotvm_task_extraction.py @@ -121,7 +121,21 @@ def get_net(batch, in_dim, out_dim, dtype, out_dtype): tasks = autotvm.task.extract_from_program(mod, target=target, params=params, ops=(dense,)) assert len(tasks) == 1 and tasks[0].name == "dense_int8.cuda" +def test_task_extraction_for_dense_int8_vnni(): + target = "llvm -mcpu=cascadelake" + + def get_net(batch, in_dim, out_dim, dtype_input, dtype_weight, out_dtype): + data = tvm.relay.var("data", shape=[batch, in_dim], dtype=dtype_input) + weight = tvm.relay.var("weight", shape=[out_dim, in_dim], dtype=dtype_weight) + 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, "uint8", "int8", "int32") + tasks = autotvm.task.extract_from_program(mod, target=target, params=params) + assert len(tasks) == 1 and tasks[0].name == "dense_vnni.x86" if __name__ == "__main__": test_task_extraction() test_task_extraction_for_dense_int8_cuda() + test_task_extraction_for_dense_int8_vnni() From 26337c5e99ba742b45f21d437b0f96359da73197 Mon Sep 17 00:00:00 2001 From: mengceng he Date: Tue, 14 Jun 2022 11:35:52 +0800 Subject: [PATCH 2/2] pylint --- tests/python/relay/test_autotvm_task_extraction.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/python/relay/test_autotvm_task_extraction.py b/tests/python/relay/test_autotvm_task_extraction.py index 1ebbdb3fd294..7092d71849b5 100644 --- a/tests/python/relay/test_autotvm_task_extraction.py +++ b/tests/python/relay/test_autotvm_task_extraction.py @@ -121,6 +121,7 @@ def get_net(batch, in_dim, out_dim, dtype, out_dtype): tasks = autotvm.task.extract_from_program(mod, target=target, params=params, ops=(dense,)) assert len(tasks) == 1 and tasks[0].name == "dense_int8.cuda" + def test_task_extraction_for_dense_int8_vnni(): target = "llvm -mcpu=cascadelake" @@ -135,6 +136,7 @@ def get_net(batch, in_dim, out_dim, dtype_input, dtype_weight, out_dtype): tasks = autotvm.task.extract_from_program(mod, target=target, params=params) assert len(tasks) == 1 and tasks[0].name == "dense_vnni.x86" + if __name__ == "__main__": test_task_extraction() test_task_extraction_for_dense_int8_cuda()