update strategy

icemelon · Jan 8, 2020 · bfff218 · bfff218
1 parent 3d3984d
commit bfff218
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py
@@ -645,11 +645,7 @@ def compute_bitpack(attrs, inputs, out_dtype):
                               name)
     return [out]
 
-# @reg.register_schedule("nn.bitpack")
-# def schedule_bitpack(attrs, outs, target):
-#     with target:
-#         return topi.generic.schedule_bitpack(outs)
-
+reg.register_schedule("nn.bitpack", strategy.schedule_bitpack)
 reg.register_pattern("nn.bitpack", OpPattern.INJECTIVE)
 
 

diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py
@@ -84,6 +84,12 @@ def schedule_l2_normalize(attrs, outs, target):
     with target:
         return topi.generic.schedule_l2_normalize(outs)
 
+# bitpack
+@generic_func
+def schedule_bitpack(attrs, outs, target):
+    with target:
+        return topi.generic.schedule_bitpack(outs)
+
 # conv2d
 def wrap_compute_conv2d(topi_func, has_group=False):
     def compute_conv2d(attrs, inputs, out_type):

diff --git a/topi/python/topi/generic/nn.py b/topi/python/topi/generic/nn.py
@@ -550,7 +550,6 @@ def schedule_binarize_pack(outs):
     return _default_schedule(outs, False)
 
 
-@tvm.target.override_native_generic_func("schedule_bitpack")
 def schedule_bitpack(outs):
     """Schedule for bitpack
     Parameters