From 461caaa5b31b6edb8ee29094ea30ebb815d09fea Mon Sep 17 00:00:00 2001 From: Cody Yu Date: Tue, 2 Mar 2021 22:17:52 +0000 Subject: [PATCH 1/2] [AutoScheduler] Query in task extraction --- python/tvm/auto_scheduler/relay_integration.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/tvm/auto_scheduler/relay_integration.py b/python/tvm/auto_scheduler/relay_integration.py index b39aba227a88..68f53125c7ae 100644 --- a/python/tvm/auto_scheduler/relay_integration.py +++ b/python/tvm/auto_scheduler/relay_integration.py @@ -283,10 +283,13 @@ def auto_schedule_topi(outs): key = register_workload_tensors(dag.workload_key(), io_tensors) target = tvm.target.Target.current() + dispatch_ctx = DispatchContext.current + state = dispatch_ctx.query(target, key, has_complex_op, dag) + schedule = None + env = TracingEnvironment.current if env is None: # in the final build mode - state = DispatchContext.current.query(target, key, has_complex_op, dag) if state is None: return None @@ -303,8 +306,6 @@ def auto_schedule_topi(outs): LayoutRewriteOption.get_target_default(target, True) != LayoutRewriteOption.NO_REWRITE and has_layout_free ): - dispatch_ctx = DispatchContext.current - state = dispatch_ctx.query(target, key, has_complex_op, dag) if state is None: return None @@ -316,7 +317,7 @@ def auto_schedule_topi(outs): else: raise ValueError("Invalid tracing mode: " + env.tracing_mode) - return None + return schedule def tensor_no_check_call(self, *indices): From 366a34a83776d6177b777287a7d3a688da73e945 Mon Sep 17 00:00:00 2001 From: Cody Yu Date: Wed, 3 Mar 2021 18:09:29 +0000 Subject: [PATCH 2/2] trigger ci