From b20f3b695b046d61ee215e11735a8c411112aa70 Mon Sep 17 00:00:00 2001
From: Tianqi Chen <tqchen@users.noreply.github.com>
Date: Tue, 7 Jul 2020 20:55:44 -0700
Subject: [PATCH] [VTA] Move compiler related registry items to
 vta/build_module.py (#6012)

---
 vta/python/vta/build_module.py | 12 ++++++++++++
 vta/python/vta/environment.py  | 12 ------------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/vta/python/vta/build_module.py b/vta/python/vta/build_module.py
index 2d67edb7051f..5becc5e037c4 100644
--- a/vta/python/vta/build_module.py
+++ b/vta/python/vta/build_module.py
@@ -125,3 +125,15 @@ def build(*args, **kwargs):
         with build_config():
             return tvm.build(*args, **kwargs)
     return tvm.build(*args, **kwargs)
+
+
+# Register key ops
+tvm.ir.register_op_attr("tir.vta.coproc_sync", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
+tvm.ir.register_op_attr("tir.vta.coproc_dep_push", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
+tvm.ir.register_op_attr("tir.vta.coproc_dep_pop", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
+
+tvm.ir.register_op_attr("tir.vta.uop_push", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
+tvm.ir.register_op_attr("tir.vta.uop_push", "TGlobalSymbol", "VTAUopPush")
+
+tvm.ir.register_op_attr("tir.vta.command_handle", "TGlobalSymbol", "VTATLSCommandHandle")
+tvm.ir.register_op_attr("tir.vta.command_handle", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
diff --git a/vta/python/vta/environment.py b/vta/python/vta/environment.py
index 3d7423848501..3a18cf74bc1d 100644
--- a/vta/python/vta/environment.py
+++ b/vta/python/vta/environment.py
@@ -313,18 +313,6 @@ def coproc_dep_pop(op):
         get_env().dev.command_handle,
         op.args[0], op.args[1])
 
-# register a dummy into to trigger registration of the ops
-# change the info to lowering rule later.
-tvm.ir.register_op_attr("tir.vta.coproc_sync", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
-tvm.ir.register_op_attr("tir.vta.coproc_dep_push", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
-tvm.ir.register_op_attr("tir.vta.coproc_dep_pop", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
-
-tvm.ir.register_op_attr("tir.vta.uop_push", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
-tvm.ir.register_op_attr("tir.vta.uop_push", "TGlobalSymbol", "VTAUopPush")
-
-tvm.ir.register_op_attr("tir.vta.command_handle", "TGlobalSymbol", "VTATLSCommandHandle")
-tvm.ir.register_op_attr("tir.vta.command_handle", "TCallEffectKind", tvm.tir.CallEffectKind.Opaque)
-
 
 def _init_env():
     """Initialize the default global env"""