""" Clip operator. """ x = inputs[0] a_min = attrs.get_float("a_min") a_max = attrs.get_float("a_max") const_min = tvm.const(a_min, x.dtype) const_max = tvm.const(a_max, x.dtype) with tvm.tag_scope(topi.tag.ELEMWISE): x = tvm.compute( x.shape, lambda *i: tvm.min(x(*i), const_max), name="clipA") x = tvm.compute( x.shape, lambda *i: tvm.max(x(*i), const_min), name="clipB") return x # override to force partition at copy reg.register_pattern("copy", OpPattern.INJECTIVE, level=15) def is_packed_layout(layout): """Check if layout is packed layout""" if layout == "NCHW": return False if "n" in layout and "c" in layout: return True return False @reg.register_alter_op_layout("conv2d", level=15) def alter_conv2d_layout(attrs, inputs, out): layout = attrs['layout'] if is_packed_layout(layout): return None return _nn.alter_conv2d_layout(attrs, inputs, out)
# pylint: disable=broad-except try: f = tvm.lower(sch, inputs, name=func_name) if "quantized_conv2d" in func_name: logging.info(graph.ir(join_entry_attrs=["shape"])) except Exception: msg = traceback.format_exc() msg += "Error during compile graph\n" msg += "--------------------------\n" msg += graph.ir(join_entry_attrs=["shape"]) raise RuntimeError(msg) return f if isinstance( f, (tvm.container.Array, tuple, list)) else [f] # override to force partition at copy reg.register_pattern("copy", OpPattern.INJECTIVE, level=15) @reg.register_compute("clip", level=15) def compute_clip(attrs, inputs, _): """ Clip operator. """ x = inputs[0] a_min = attrs.get_float("a_min") a_max = attrs.get_float("a_max") const_min = tvm.const(a_min, x.dtype) const_max = tvm.const(a_max, x.dtype) with tvm.tag_scope(topi.tag.ELEMWISE): x = tvm.compute( x.shape, lambda *i: tvm.min(x(*i), const_max), name="clipA") x = tvm.compute( x.shape, lambda *i: tvm.max(x(*i), const_min), name="clipB") return x
def _bitpack(*indices): ret = None mask = tvm.const((1 << bits) - 1, pack_type) for k in range(lanes): idx = list(indices) idx[-1] = idx[-1] * lanes + k elem = data(*idx).astype(pack_type) if k == 0: ret = elem & mask else: val = (elem & mask) << tvm.const(k * bits, pack_type) ret = ret | val return ret return tvm.compute(oshape, _bitpack, name=name, tag='bitpack') @reg.register_compute("bitpack", level=15) def compute_bitpack(attrs, inputs, out): lanes = attrs.get_int("lanes") dtype = inputs[0].dtype assert dtype == "int8" width = 8 assert width % lanes == 0 bits = 8 // lanes return bitpack(inputs[0], bits, dtype) reg.register_schedule("bitpack", _fschedule_broadcast) reg.register_pattern("bitpack", OpPattern.INJECTIVE)