bitpack strategy

icemelon · Jan 9, 2020 · cf4fcec · cf4fcec
1 parent ff0ca78
commit cf4fcec
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
@@ -94,6 +94,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_compute, has_group=False):
     """Wrap conv2d topi compute"""

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