From bfff218c2b1fa9968962597ab1ba3e5f1e78c078 Mon Sep 17 00:00:00 2001 From: Cody Yu Date: Wed, 8 Jan 2020 22:22:33 +0000 Subject: [PATCH] update strategy --- python/tvm/relay/op/nn/_nn.py | 6 +----- python/tvm/relay/op/strategy/generic.py | 6 ++++++ topi/python/topi/generic/nn.py | 1 - 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index e6e85d8e25898..9a8321b28b90f 100644 --- 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 index 2cf19135e05de..c2da5925ac372 100644 --- 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 index 40dd58f56b1b7..07e74d3b561ce 100644 --- 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