| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # pylint: disable=invalid-name, unused-variable, too-many-locals |
| # pylint: disable=unused-argument, redefined-builtin |
| """Generic convolution schedules""" |
| from tvm import te |
| from tvm import autotvm |
| from tvm import relay |
| from tvm.autotvm.task.space import SplitEntity, OtherOptionEntity |
| from ..utils import get_const_tuple, traverse_inline |
| from ..nn.utils import get_pad_tuple |
| |
| |
| def fallback_schedule_cpu_common_int8(cfg, wkl, int32_lanes, num_int8_elements): |
| """Fallback schedule for conv2d int8 on cpu. |
| Normally the inner most pattern takes two int8/uint8 tensors |
| data[num_int8_elements] and kernel[int32_lanes, num_int8_elements], |
| produces a dot product int32/uint32 output[int32_lanes]. |
| |
| Parameters |
| ---------- |
| int32_lanes : int |
| How many numbers of int32/uint32 will be produced using intrinsic. |
| This is related to output channel. |
| num_int8_elements : int |
| How many numbers of input int32/uint32 will be multiplied and reduced. |
| This is related to input channel. |
| """ |
| pt, pl, pb, pr = wkl.padt, wkl.padl, wkl.padb, wkl.padr |
| HSTR, WSTR = wkl.stride_h, wkl.stride_w |
| dilated_kernel_w = (wkl.kernel_w - 1) * wkl.dilation_w + 1 |
| out_width = (wkl.width + pl + pr - dilated_kernel_w) // WSTR + 1 |
| |
| assert ( |
| wkl.out_filter % int32_lanes == 0 |
| ), f"wkl.out_filter={wkl.out_filter}, int32_lanes={int32_lanes}" |
| assert ( |
| wkl.in_filter % num_int8_elements == 0 |
| ), f"wkl.in_filter={wkl.in_filter}, num_int8_elements={num_int8_elements}" |
| |
| oc_bn = int32_lanes if int32_lanes >= num_int8_elements else num_int8_elements |
| ic_bn = 1 |
| for bn in range(oc_bn, 0, -4): |
| if wkl.in_filter % bn == 0: |
| ic_bn = bn |
| break |
| |
| reg_n = 1 |
| for n in range(31, 0, -1): |
| if out_width % n == 0: |
| reg_n = n |
| break |
| |
| cfg["tile_ic"] = SplitEntity([wkl.in_filter // ic_bn, ic_bn]) |
| cfg["tile_oc"] = SplitEntity([wkl.out_filter // oc_bn, oc_bn]) |
| cfg["tile_ow"] = SplitEntity([out_width // reg_n, reg_n]) |
| cfg["unroll_kw"] = OtherOptionEntity(False) |
| |
| |
| def fallback_schedule_cpu_1x1_int8(cfg, wkl, int32_lanes, num_int8_elements): |
| """Fallback schedule for 1x1 conv2d int8 on cpu. |
| Normally the inner most pattern takes two int8/uint8 tensors |
| data[num_int8_elements] and kernel[int32_lanes, num_int8_elements], |
| produces a dot product int32/uint32 output[int32_lanes]. |
| |
| Parameters |
| ---------- |
| int32_lanes : int |
| How many numbers of int32/uint32 will be produced using intrinsic. |
| This is related to output channel. |
| num_int8_elements : int |
| How many numbers of input int32/uint32 will be multiplied and reduced. |
| This is related to input channel. |
| """ |
| pt, pl, pb, pr = wkl.padt, wkl.padl, wkl.padb, wkl.padr |
| HSTR, WSTR = wkl.stride_h, wkl.stride_w |
| out_height = (wkl.height + pt + pb - wkl.kernel_h) // HSTR + 1 |
| out_width = (wkl.width + pl + pr - wkl.kernel_w) // WSTR + 1 |
| |
| assert ( |
| wkl.out_filter % int32_lanes == 0 |
| ), f"wkl.out_filter={wkl.out_filter}, int32_lanes={int32_lanes}" |
| assert ( |
| wkl.in_filter % num_int8_elements == 0 |
| ), f"wkl.in_filter={wkl.in_filter}, num_int8_elements={num_int8_elements}" |
| |
| oc_bn = int32_lanes if int32_lanes >= num_int8_elements else num_int8_elements |
| ic_bn = 1 |
| for bn in range(oc_bn, 0, -4): |
| if wkl.in_filter % bn == 0: |
| ic_bn = bn |
| break |
| |
| for ow_factor in range(out_width, 0, -1): |
| if out_width % ow_factor == 0: |
| for oh_factor in range(out_height, 0, -1): |
| if out_height % oh_factor == 0 and ow_factor * oh_factor < 32: |
| cfg["tile_ic"] = SplitEntity([wkl.in_filter // ic_bn, ic_bn]) |
| cfg["tile_oc"] = SplitEntity([wkl.out_filter // oc_bn, oc_bn]) |
| cfg["tile_oh"] = OtherOptionEntity(oh_factor) |
| cfg["tile_ow"] = SplitEntity([out_width // ow_factor, ow_factor]) |
| return |
| raise ValueError(f"cannot decide default schedule for workload: {wkl}") |
| |
| |
| def schedule_conv_NCHWc_cpu_common_int8( |
| s, |
| cfg, |
| data_vec, |
| kernel_vec, |
| conv_out, |
| last, |
| int32_lanes=16, |
| int8_elems=4, |
| intrin=None, |
| inline_fused=True, |
| mem_scope="global", |
| ): |
| """ |
| Defines the schedule for INT8 for Intel and ARM machines |
| Uses the Intel/ARM intrinsics to use INT8 operations |
| More details - https://software.intel.com/en-us/articles/ |
| lower-numerical-precision-deep-learning-inference-and-training |
| """ |
| if isinstance(cfg["tile_ow"], int): |
| reg_n = cfg["tile_ow"] |
| else: |
| reg_n = cfg["tile_ow"].size[-1] |
| |
| if isinstance(cfg["unroll_kw"], (int, bool)): |
| unroll_kw = cfg["unroll_kw"] |
| else: |
| unroll_kw = cfg["unroll_kw"].val |
| |
| _, _, _, _, ic_bn = get_const_tuple(data_vec.shape) |
| _, _, _, _, oc_bn = get_const_tuple(conv_out.shape) |
| |
| # schedule pad |
| if isinstance(s[data_vec].op, te.tensor.ComputeOp) and "pad" in data_vec.op.tag: |
| batch, ic_chunk, ih, iw, ic_block = s[data_vec].op.axis |
| parallel_axis = s[data_vec].fuse(batch, ic_chunk, ih) |
| s[data_vec].parallel(parallel_axis) |
| data_vec = data_vec.op.input_tensors[0] |
| |
| if autotvm.GLOBAL_SCOPE.in_tuning: |
| # only in autotuning, input data of conv2d_NCHWc will be 4-D. |
| # skip this part during tuning to make records accurate. |
| # this part will be folded during Relay fold_constant pass. |
| if isinstance(data_vec.op, te.tensor.ComputeOp): |
| s[data_vec].pragma(s[data_vec].op.axis[0], "debug_skip_region") |
| if isinstance(kernel_vec.op, te.tensor.ComputeOp): |
| s[kernel_vec].pragma(s[kernel_vec].op.axis[0], "debug_skip_region") |
| elif isinstance(kernel_vec.op, te.tensor.ComputeOp) and kernel_vec.name == "kernel_vec": |
| # data and kernel are not pre-computed, schedule layout transform here. |
| # this should only be used by x86 conv2d_nchw, which is for |
| # testing purpose. |
| batch, ic_chunk, ih, iw, ic_block = s[data_vec].op.axis |
| parallel_axis = s[data_vec].fuse(batch, ic_chunk, ih) |
| s[data_vec].parallel(parallel_axis) |
| |
| # conv2d_nchwc_int8 has 7D kernel |
| oc_chunk, ic_chunk, oh, ow, ic_block, oc_block, _ = s[kernel_vec].op.axis |
| s[kernel_vec].reorder(oc_chunk, oh, ic_chunk, ow, ic_block, oc_block) |
| oc_bn = cfg["tile_oc"].size[-1] |
| if oc_bn > 1: |
| s[kernel_vec].vectorize(oc_block) |
| parallel_axis = s[kernel_vec].fuse(oc_chunk, oh) |
| s[kernel_vec].parallel(parallel_axis) |
| |
| # schedule 5-D NCHW[x]c conv |
| C, O = conv_out, last |
| CC = s.cache_write(C, mem_scope) |
| |
| batch, oc_chunk, oh, ow, oc_block = s[C].op.axis |
| ow_chunk, ow_block = s[C].split(ow, factor=reg_n) |
| s[C].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block) |
| parallel_axis = s[C].fuse(batch, oc_chunk, oh) |
| s[C].vectorize(oc_block) |
| if C == O: |
| s[C].parallel(parallel_axis) |
| |
| s[CC].compute_at(s[C], parallel_axis) |
| _, oc_chunk, oh, ow, oc_block = s[CC].op.axis |
| kh, kw, ic_outer, ic_f_inner, ic_s_inner = s[CC].op.reduce_axis |
| |
| ow_chunk, ow_block = s[CC].split(ow, factor=reg_n) |
| |
| assert oc_bn % int32_lanes == 0, f"oc_bn={oc_bn} % int32_lanes={int32_lanes} != 0" |
| assert ( |
| ic_bn % int8_elems == 0 |
| ), f"ic_bn={ic_bn} % int8_elems={int8_elems} != 0" # (u)int8 elements in (u)int32 |
| |
| oc_f_inner, oc_s_inner = s[CC].split(oc_block, factor=int32_lanes) |
| |
| if unroll_kw: |
| s[CC].reorder( |
| oc_chunk, |
| oh, |
| ow_chunk, |
| ic_outer, |
| kh, |
| ic_f_inner, |
| kw, |
| ow_block, |
| oc_f_inner, |
| oc_s_inner, |
| ic_s_inner, |
| ) |
| s[CC].unroll(kw) |
| else: |
| s[CC].reorder( |
| oc_chunk, |
| oh, |
| ow_chunk, |
| ic_outer, |
| kh, |
| kw, |
| ic_f_inner, |
| ow_block, |
| oc_f_inner, |
| oc_s_inner, |
| ic_s_inner, |
| ) |
| |
| if intrin is not None: |
| s[CC].tensorize(oc_s_inner, intrin) |
| s[CC].unroll(ow_block) |
| s[CC].unroll(oc_f_inner) |
| |
| if C != O: |
| out_ndim = len(s[O].op.axis) |
| if out_ndim == 5: |
| batch, oc_chunk, oh, ow, oc_block = s[O].op.axis |
| ow_chunk, ow_block = s[O].split(ow, factor=reg_n) |
| s[O].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block) |
| elif out_ndim == 4: |
| batch, oc, oh, ow = s[O].op.axis |
| ow_chunk, ow_block = s[O].split(ow, factor=reg_n) |
| oc_chunk, oc_block = s[O].split(oc, factor=oc_bn) |
| s[O].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block) |
| else: |
| raise ValueError(f"Unsupported output ndim: {out_ndim}") |
| parallel_axis = s[O].fuse(batch, oc_chunk, oh) |
| if inline_fused: |
| s[C].compute_at(s[O], ow_block) |
| else: |
| s[C].compute_at(s[O], parallel_axis) |
| s[O].vectorize(oc_block) |
| s[O].parallel(parallel_axis) |
| |
| return s |
| |
| |
| def schedule_conv_NCHWc_cpu_1x1_int8( |
| s, |
| cfg, |
| data_vec, |
| kernel_vec, |
| conv_out, |
| last, |
| int32_lanes=16, |
| int8_elems=4, |
| intrin=None, |
| inline_fused=False, |
| mem_scope="global", |
| ): |
| """ |
| Defines the 1x1 conv schedule for INT8 for Intel and ARM machines |
| Uses the Intel/ARM intrinsics to use INT8 operations |
| More details - https://software.intel.com/en-us/articles/ |
| lower-numerical-precision-deep-learning-inference-and-training |
| """ |
| oh_factor, ow_factor = cfg["tile_oh"].val, cfg["tile_ow"].size[-1] |
| _, _, _, _, ic_bn = get_const_tuple(data_vec.shape) |
| _, _, _, _, oc_bn = get_const_tuple(conv_out.shape) |
| |
| # schedule pad |
| if isinstance(s[data_vec].op, te.tensor.ComputeOp) and "pad" in data_vec.op.tag: |
| batch, ic_chunk, ih, iw, ic_block = s[data_vec].op.axis |
| parallel_axis = s[data_vec].fuse(batch, ic_chunk, ih) |
| s[data_vec].parallel(parallel_axis) |
| data_vec = data_vec.op.input_tensors[0] |
| |
| if autotvm.GLOBAL_SCOPE.in_tuning: |
| # only in autotuning, input data of conv2d_NCHWc will be 4-D. |
| # skip this part during tuning to make records accurate. |
| # this part will be folded during Relay fold_constant pass. |
| if isinstance(data_vec.op, te.tensor.ComputeOp): |
| s[data_vec].pragma(s[data_vec].op.axis[0], "debug_skip_region") |
| if isinstance(kernel_vec.op, te.tensor.ComputeOp): |
| s[kernel_vec].pragma(s[kernel_vec].op.axis[0], "debug_skip_region") |
| elif isinstance(kernel_vec.op, te.tensor.ComputeOp) and kernel_vec.name == "kernel_vec": |
| # data and kernel are not pre-computed, schedule layout transform here. |
| # this should only be used by x86 conv2d_nchw, which is for |
| # testing purpose. |
| batch, ic_chunk, ih, iw, ic_block = s[data_vec].op.axis |
| parallel_axis = s[data_vec].fuse(batch, ic_chunk, ih) |
| s[data_vec].parallel(parallel_axis) |
| |
| # Conv2d int8 schedule has 7D kernel |
| oc_chunk, ic_chunk, oh, ow, ic_block, oc_block, _ = s[kernel_vec].op.axis |
| s[kernel_vec].reorder(oc_chunk, oh, ic_chunk, ow, ic_block, oc_block) |
| oc_bn = cfg["tile_oc"].size[-1] |
| if oc_bn > 1: |
| s[kernel_vec].vectorize(oc_block) |
| parallel_axis = s[kernel_vec].fuse(oc_chunk, oh) |
| s[kernel_vec].parallel(parallel_axis) |
| |
| C, O = conv_out, last |
| CC = s.cache_write(C, mem_scope) |
| |
| batch, oc_chunk, oh, ow, oc_block = s[C].op.axis |
| oh_outer, oh_inner = s[C].split(oh, factor=oh_factor) |
| ow_outer, ow_inner = s[C].split(ow, factor=ow_factor) |
| s[C].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block) |
| s[C].vectorize(oc_block) |
| |
| parallel_axis = s[C].fuse(batch, oc_chunk, oh_outer) |
| if C == O: |
| s[C].parallel(parallel_axis) |
| s[CC].compute_at(s[C], parallel_axis) # good perf on mobilenet, but not on individuals? |
| |
| _, oc_chunk, oh, ow, oc_block = s[CC].op.axis |
| kh, kw, ic_outer, ic_f_inner, ic_s_inner = s[CC].op.reduce_axis |
| |
| assert oc_bn % int32_lanes == 0 |
| assert ic_bn % int8_elems == 0 # (u)int8 elements in (u)int32 |
| |
| oc_f_inner, oc_s_inner = s[CC].split(oc_block, factor=int32_lanes) |
| |
| oh_outer, oh_inner = s[CC].split(oh, factor=oh_factor) |
| ow_outer, ow_inner = s[CC].split(ow, factor=ow_factor) |
| |
| s[CC].reorder( |
| oc_chunk, |
| oh_outer, |
| ow_outer, |
| kh, |
| kw, |
| ic_outer, |
| ic_f_inner, |
| oh_inner, |
| ow_inner, |
| oc_f_inner, |
| oc_s_inner, |
| ic_s_inner, |
| ) |
| s[CC].fuse(oc_chunk, oh_outer) |
| |
| if intrin is not None: |
| s[CC].tensorize(oc_s_inner, intrin) |
| s[CC].unroll(ow_inner) |
| s[CC].unroll(oh_inner) |
| |
| if C != O: |
| out_ndim = len(s[O].op.axis) |
| if out_ndim == 5: |
| batch, oc_chunk, oh, ow, oc_block = s[O].op.axis |
| oh_outer, oh_inner = s[O].split(oh, factor=oh_factor) |
| ow_outer, ow_inner = s[O].split(ow, factor=ow_factor) |
| elif out_ndim == 4: |
| batch, oc, oh, ow = s[O].op.axis |
| oc_chunk, oc_block = s[O].split(oc, factor=oc_bn) |
| oh_outer, oh_inner = s[O].split(oh, factor=oh_factor) |
| ow_outer, ow_inner = s[O].split(ow, factor=ow_factor) |
| else: |
| raise ValueError(f"Unsupported output ndim: {out_ndim}") |
| |
| s[O].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block) |
| parallel_axis = s[O].fuse(batch, oc_chunk, oh_outer) |
| if inline_fused: |
| s[C].compute_at(s[O], ow_inner) |
| else: |
| s[C].compute_at(s[O], parallel_axis) |
| s[O].vectorize(oc_block) |
| s[O].parallel(parallel_axis) |
| |
| return s |
| |
| |
| def schedule_depthwise_conv2d_nhwc(outs): |
| """Create schedule for depthwise conv2d in NHWC layout. |
| Parameters |
| ---------- |
| outs : list[te.tensor.Tensor] |
| The output tensors. |
| Returns |
| ------- |
| s : tvm.te.schedule.Schedule |
| The computation schedule for depthwise conv2d. |
| """ |
| outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs |
| s = te.create_schedule([x.op for x in outs]) |
| |
| def _callback(op): |
| """Traverse operators from computation graph""" |
| if "depthwise_conv2d_nhwc" in op.tag: |
| out = outs[0] |
| depthwise_conv2d_out = op.output(0) |
| data_pad = depthwise_conv2d_out.op.input_tensors[0] |
| s[data_pad].compute_inline() |
| if depthwise_conv2d_out != out: |
| s[depthwise_conv2d_out].compute_at(s[out], s[out].op.axis[3]) |
| s[out].fuse(*s[out].op.axis) |
| |
| traverse_inline(s, outs[0].op, _callback) |
| return s |
| |
| |
| def conv2d_alter_int8_common( |
| data, |
| data_tensor, |
| kernel, |
| kernel_tensor, |
| output_tensor, |
| attrs, |
| data_dtype: str, |
| in_channel_vector_length: int, |
| out_channel_vector_length: int, |
| ): |
| """ |
| Convert TE inputs/outputs so that they are suitable for fast Int8 instructions. |
| |
| Int8 instructions require input channels and output channels to be a |
| multiple of the vector length. For input channels, we pad both the inputs |
| and weights channels. For output channels, we pad the weight and |
| stride_slice the output. |
| |
| Arguments |
| --------- |
| data: Expr |
| Data Expr |
| data_tensor: Tensor |
| Data tensor |
| kernel: Expr |
| Kernel Expr |
| kernel_tensor: Tensor |
| Kernel tensor |
| output_tensor: Tensor |
| Output tensor |
| attrs: Conv2dAttrs |
| Attributes of the computation |
| data_dtype: "int8" or "uint8" |
| Desired dtype of data. Data will be converted to this dtype before the main computation. |
| in_channel_vector_length: int |
| Length of vector units on target hardware. Input channels are padded to this length. |
| out_channel_vector_length: int |
| Output size of vector instruction. Output channels are padded to this length. |
| |
| Returns |
| ------- |
| out : Tensor |
| Conv2d computation with inputs in the correct order for tensorization. |
| """ |
| # Dilation not supported yet. Return None if dilation is not (1, 1) |
| dilation = attrs.get_int_tuple("dilation") |
| if not (dilation[0] == 1 and dilation[1] == 1): |
| return None |
| |
| # No legalization for depthwise convolutions yet. |
| groups = attrs.get_int("groups") |
| if groups != 1: |
| return None |
| |
| # Get the conv attrs |
| new_attrs = {k: attrs[k] for k in attrs.keys()} |
| |
| padding = attrs.get_int_tuple("padding") |
| kh, kw = attrs.get_int_tuple("kernel_size") |
| pt, pl, pb, pr = get_pad_tuple(padding, (kh, kw)) |
| |
| if data_tensor.dtype != data_dtype: |
| # How to convert data to uint8 |
| # Original --> C = A (conv) B |
| # A and B are int8 |
| # C = (A + 128 - 128) (conv) B |
| # C = (A' conv B) - 128 (conv) B |
| # where A' = A + 128 |
| # and 128 (conv) B is basically a reduce on CRS axis for weights. |
| # |
| # How to convert data to int8 |
| # C = (A - 128 + 128) (conv) B |
| # C = (A' conv B) + 128 (conv) B |
| # where A' = A - 128 |
| if data_dtype == "uint8": |
| # shift data to uint8 |
| before_shift = relay.add |
| after_shift = relay.subtract |
| pad_value = 128 |
| else: |
| # shift data to int8 |
| before_shift = relay.subtract |
| after_shift = relay.add |
| pad_value = -128 |
| |
| if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO": |
| adjust_shift = relay.sum(relay.cast(kernel, dtype="int32"), axis=(0, 1, 2)) |
| pad_width = ((0, 0), (pt, pb), (pl, pr), (0, 0)) |
| elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW": |
| pad_width = ((0, 0), (0, 0), (pt, pb), (pl, pr)) |
| adjust_shift = relay.sum(relay.cast(kernel, dtype="int32"), axis=(1, 2, 3)) |
| adjust_shift = relay.expand_dims(adjust_shift, axis=1, num_newaxis=2) |
| else: |
| return None |
| |
| data = relay.cast(data, "int32") |
| data = before_shift(data, relay.const(128, "int32")) |
| data = relay.cast(data, data_dtype) |
| |
| # Do external padding as pad value has to be 128. |
| if any(padding): |
| data = relay.nn.pad(data, pad_width=pad_width, pad_value=pad_value) |
| |
| new_attrs["padding"] = (0, 0) |
| |
| # Multiply 128 to adjust shift. |
| adjust_shift = relay.multiply(adjust_shift, relay.const(128, "int32")) |
| |
| # Flags to remember if the expr is modified |
| ic_modified = False |
| oc_modified = False |
| |
| # Find the value of input and output channel. |
| in_channel = -1 |
| out_channel = -1 |
| if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO": |
| in_channel = data_tensor.shape[3].value |
| out_channel = kernel_tensor.shape[3].value |
| elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW": |
| in_channel = data_tensor.shape[1].value |
| out_channel = kernel_tensor.shape[0].value |
| else: |
| return None |
| |
| if in_channel % in_channel_vector_length != 0: |
| new_in_channel = ( |
| (in_channel + in_channel_vector_length) // in_channel_vector_length |
| ) * in_channel_vector_length |
| diff = new_in_channel - in_channel |
| if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO": |
| data = relay.nn.pad(data, pad_width=((0, 0), (0, 0), (0, 0), (0, diff))) |
| kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, diff), (0, 0))) |
| ic_modified = True |
| elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW": |
| pad_width = ((0, 0), (0, diff), (0, 0), (0, 0)) |
| data = relay.nn.pad(data, pad_width=pad_width) |
| kernel = relay.nn.pad(kernel, pad_width=pad_width) |
| ic_modified = True |
| else: |
| return None |
| |
| new_out_channel = out_channel |
| if out_channel % out_channel_vector_length != 0: |
| new_out_channel = ( |
| (out_channel + out_channel_vector_length) // out_channel_vector_length |
| ) * out_channel_vector_length |
| diff = new_out_channel - out_channel |
| if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO": |
| kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, 0), (0, diff))) |
| oc_modified = True |
| elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW": |
| kernel = relay.nn.pad(kernel, pad_width=((0, diff), (0, 0), (0, 0), (0, 0))) |
| oc_modified = True |
| else: |
| return None |
| |
| if oc_modified: |
| new_attrs["channels"] = new_out_channel |
| out = relay.nn.conv2d(data, kernel, **new_attrs) |
| original_out_shape = [x.value for x in output_tensor.shape] |
| out = relay.strided_slice(out, begin=[0, 0, 0, 0], end=original_out_shape) |
| else: |
| out = relay.nn.conv2d(data, kernel, **new_attrs) |
| |
| if data_tensor.dtype != data_dtype: |
| out = after_shift(out, adjust_shift) |
| |
| return out |
