| # 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, no-else-return |
| """Conv3D operators""" |
| from collections import namedtuple |
| |
| import tvm |
| from tvm import autotvm, te |
| from tvm.autotvm.task.space import OtherOptionEntity, SplitEntity |
| from tvm.target.x86 import get_simd_32bit_lanes |
| |
| from ..nn.pad import pad |
| from ..nn.utils import get_pad_tuple3d, infer_pad3d |
| from ..utils import get_const_int, get_const_tuple, simplify, traverse_inline |
| |
| Workload3D = namedtuple( |
| "Workload", |
| [ |
| "in_dtype", |
| "out_dtype", |
| "depth", |
| "height", |
| "width", |
| "in_filter", |
| "groups", |
| "out_filter", |
| "dkernel", |
| "hkernel", |
| "wkernel", |
| "dpad", |
| "hpad", |
| "wpad", |
| "dstride", |
| "hstride", |
| "wstride", |
| ], |
| ) |
| |
| |
| @autotvm.register_topi_compute("conv3d_ndhwc.x86") |
| def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dtype): |
| """3D convolution forward operator. |
| |
| Parameters |
| ---------- |
| input : tvm.te.Tensor |
| 5-D input data with shapes: |
| [batch, in_depth, in_height, in_width, in_channel] for NDHWC layout |
| |
| filter : tvm.te.Tensor |
| 5-D filter with shape [kernel_depth, kernel_height, kernel_width, in_channels, out_channels] |
| |
| strides : int or a list/tuple of three ints |
| stride size, or [stride_depth, stride_height, stride_width] |
| |
| padding : int or a list/tuple of three ints |
| padding size, or [pad_depth, pad_height, pad_width] |
| |
| dilation: int or a list/tuple of three ints |
| dilation size, or [dilation_depth, dilation_height, dilation_width] |
| |
| groups: int |
| Number of groups |
| |
| Returns |
| ------- |
| output : tvm.te.Tensor |
| 5-D with shape [batch, out_depth, out_height, out_width, out_channel] for NDHWC layout |
| """ |
| layout = "NDHWC" |
| out_dtype = data.dtype if out_dtype is None else out_dtype |
| strides = strides if isinstance(strides, (tuple, list)) else (strides, strides, strides) |
| dilation = dilation if isinstance(dilation, (tuple, list)) else (dilation, dilation, dilation) |
| |
| _create_tuning_space(cfg, data, kernel, strides, padding, dilation, groups, layout) |
| if cfg.is_fallback: |
| _get_default_config(cfg, data, kernel, strides, padding, groups, out_dtype, layout) |
| return _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dtype) |
| |
| |
| @autotvm.register_topi_compute("conv3d_ncdhw.x86") |
| def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, groups, out_dtype): |
| """3D convolution forward operator. |
| |
| Parameters |
| ---------- |
| input : tvm.te.Tensor |
| 5-D input data with shapes: |
| [batch, in_channel, in_depth, in_height, in_width] for NCDHW layout |
| |
| filter : tvm.te.Tensor |
| 5-D filter with shape [out_channels, in_channels, kernel_depth, kernel_height, kernel_width] |
| |
| strides : int or a list/tuple of three ints |
| stride size, or [stride_depth, stride_height, stride_width] |
| |
| padding : int or a list/tuple of three ints |
| padding size, or [pad_depth, pad_height, pad_width] |
| |
| dilation: int or a list/tuple of three ints |
| dilation size, or [dilation_depth, dilation_height, dilation_width] |
| |
| groups: int |
| Number of groups |
| |
| Returns |
| ------- |
| output : tvm.te.Tensor |
| 5-D with shape [batch, out_channel, out_depth, out_height, out_width] for NCDHW layout |
| """ |
| # assert groups == 1, "conv3d_ncdhw.x86 does not support groups" |
| layout = "NCDHW" |
| out_dtype = data.dtype if out_dtype is None else out_dtype |
| strides = strides if isinstance(strides, (tuple, list)) else (strides, strides, strides) |
| dilation = dilation if isinstance(dilation, (tuple, list)) else (dilation, dilation, dilation) |
| |
| _create_tuning_space(cfg, data, kernel, strides, padding, dilation, groups, layout) |
| if cfg.is_fallback: |
| _get_default_config(cfg, data, kernel, strides, padding, groups, out_dtype, layout) |
| return _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, groups, out_dtype) |
| |
| |
| @autotvm.register_topi_schedule("conv3d_ndhwc.x86") |
| def schedule_conv3d_ndhwc(cfg, outs): |
| """TOPI schedule callback for conv3d |
| |
| Parameters |
| ---------- |
| outs: Array of Tensor |
| The computation graph description of conv3d |
| in the format of an array of tensors. |
| |
| Returns |
| ------- |
| s: Schedule |
| The computation schedule for conv3d. |
| """ |
| s = te.create_schedule([x.op for x in outs]) |
| |
| def _traverse(op): |
| if "conv3d_ndhwc" in op.tag: |
| output = op.output(0) |
| conv_out = op.input_tensors[0] |
| kernel_vec = conv_out.op.input_tensors[1] |
| kernel = kernel_vec.op.input_tensors[0] |
| if isinstance(kernel.op, tvm.te.ComputeOp) and "dilate" in kernel.op.tag: |
| s[kernel].compute_inline() |
| data_vec = conv_out.op.input_tensors[0] |
| data = data_vec.op.input_tensors[0] |
| data_pad = None |
| if isinstance(data.op, tvm.te.ComputeOp) and "pad" in data.op.tag: |
| data_pad = data |
| data = data_pad.op.input_tensors[0] |
| |
| kd, kh, kw, i, o = get_const_tuple(kernel.shape) |
| args = [s, cfg, data, data_pad, data_vec, kernel_vec, conv_out, output, outs[0]] |
| _schedule_conv3d_ndhwc(*args) |
| |
| traverse_inline(s, outs[0].op, _traverse) |
| return s |
| |
| |
| @autotvm.register_topi_schedule("conv3d_ncdhw.x86") |
| def schedule_conv3d_ncdhw(cfg, outs): |
| """TOPI schedule callback for conv3d |
| |
| Parameters |
| ---------- |
| outs: Array of Tensor |
| The computation graph description of conv3d |
| in the format of an array of tensors. |
| |
| Returns |
| ------- |
| s: Schedule |
| The computation schedule for conv3d. |
| """ |
| s = te.create_schedule([x.op for x in outs]) |
| |
| def _traverse(op): |
| if "conv3d_ncdhw" in op.tag: |
| output = op.output(0) |
| conv_out = op.input_tensors[0] |
| kernel_vec = conv_out.op.input_tensors[1] |
| kernel = kernel_vec.op.input_tensors[0] |
| if isinstance(kernel.op, tvm.te.ComputeOp) and "dilate" in kernel.op.tag: |
| s[kernel].compute_inline() |
| data_vec = conv_out.op.input_tensors[0] |
| data = data_vec.op.input_tensors[0] |
| data_pad = None |
| if isinstance(data.op, tvm.te.ComputeOp) and "pad" in data.op.tag: |
| data_pad = data |
| data = data_pad.op.input_tensors[0] |
| |
| kd, kh, kw, i, o = get_const_tuple(kernel.shape) |
| args = [s, cfg, data, data_pad, data_vec, kernel_vec, conv_out, output, outs[0]] |
| _schedule_conv3d_ncdhw(*args) |
| |
| traverse_inline(s, outs[0].op, _traverse) |
| return s |
| |
| |
| def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dtype): |
| out_dtype = data.dtype if out_dtype is None else out_dtype |
| |
| assert isinstance(dilation, int) or len(dilation) == 3 |
| if isinstance(dilation, int): |
| dilation_d, dilation_h, dilation_w = (dilation, dilation, dilation) |
| else: |
| dilation_d, dilation_h, dilation_w = dilation |
| |
| DSTR, HSTR, WSTR = strides |
| batch_size, in_depth, in_height, in_width, in_channel = get_const_tuple(data.shape) |
| kernel_depth, kernel_height, kernel_width, _, num_filter = get_const_tuple(kernel.shape) |
| |
| assert in_channel % groups == 0, "input channels must be a multiple of group size" |
| assert num_filter % groups == 0, "number of filters must be a multiple of group size" |
| |
| dilated_kernel_d = (kernel_depth - 1) * dilation_d + 1 |
| dilated_kernel_h = (kernel_height - 1) * dilation_h + 1 |
| dilated_kernel_w = (kernel_width - 1) * dilation_w + 1 |
| |
| pad_front, pad_top, pad_left, pad_back, pad_down, pad_right = get_pad_tuple3d( |
| padding, (dilated_kernel_d, dilated_kernel_h, dilated_kernel_w) |
| ) |
| |
| pad_d = pad_front + pad_back |
| pad_h = pad_top + pad_down |
| pad_w = pad_left + pad_right |
| |
| pad_depth = in_depth + pad_d |
| pad_height = in_height + pad_h |
| pad_width = in_width + pad_w |
| |
| out_depth = simplify((in_depth + pad_d - dilated_kernel_d) // DSTR + 1) |
| out_height = simplify((in_height + pad_h - dilated_kernel_h) // HSTR + 1) |
| out_width = simplify((in_width + pad_w - dilated_kernel_w) // WSTR + 1) |
| |
| # pack data |
| DOPAD = pad_d != 0 or pad_h != 0 or pad_w != 0 |
| if DOPAD: |
| data_pad = pad( |
| data, |
| (0, pad_front, pad_top, pad_left, 0), |
| (0, pad_back, pad_down, pad_right, 0), |
| name="data_pad", |
| ) |
| else: |
| data_pad = data |
| |
| # fetch schedule |
| ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1] |
| assert groups == 1 or ic_bn <= groups |
| assert groups == 1 or oc_bn <= groups |
| shape = (batch_size, in_channel // ic_bn, pad_depth, pad_height, ic_bn, pad_width) |
| data_vec = te.compute( |
| shape, lambda n, C, d, h, c, w: data_pad[n, d, h, w, C * ic_bn + c], name="data_vec" |
| ) |
| |
| ci_tile = in_channel // groups // ic_bn |
| if ci_tile == 0 or ci_tile * ic_bn * groups < in_channel: |
| ci_tile += 1 |
| |
| # pack kernel |
| shape = (num_filter // oc_bn, ci_tile, kernel_depth, kernel_height, kernel_width, ic_bn, oc_bn) |
| kernel_vec = te.compute( |
| shape, |
| lambda CO, CI, d, h, w, ci, co: kernel[d, h, w, CI * ic_bn + ci, CO * oc_bn + co], |
| name="kernel_vec", |
| ) |
| |
| # convolution |
| oshape = (batch_size, num_filter // oc_bn, out_depth, out_height, out_width, oc_bn) |
| unpack_shape = (batch_size, out_depth, out_height, out_width, num_filter) |
| |
| ic = te.reduce_axis((0, in_channel // groups), name="ic") |
| kh = te.reduce_axis((0, kernel_height), name="kh") |
| kw = te.reduce_axis((0, kernel_width), name="kw") |
| kd = te.reduce_axis((0, kernel_depth), name="kd") |
| idxmod = tvm.tir.indexmod |
| idxdiv = tvm.tir.indexdiv |
| |
| conv = te.compute( |
| oshape, |
| lambda n, oc_chunk, od, oh, ow, oc_block: te.sum( |
| data_vec[ |
| n, |
| idxdiv( |
| (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups) |
| + ic, |
| ic_bn, |
| ), |
| od * DSTR + kd * dilation_d, |
| oh * HSTR + kh * dilation_h, |
| idxmod( |
| (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups) |
| + ic, |
| ic_bn, |
| ), |
| ow * WSTR + kw * dilation_w, |
| ].astype(out_dtype) |
| * kernel_vec[ |
| oc_chunk, idxdiv(ic, ic_bn), kd, kh, kw, idxmod(ic, ic_bn), oc_block |
| ].astype(out_dtype), |
| axis=[kd, kh, kw, ic], |
| ), |
| name="conv", |
| ) |
| conv_unpacked = te.compute( |
| unpack_shape, |
| lambda n, d, h, w, c: conv[n, idxdiv(c, oc_bn), d, h, w, idxmod(c, oc_bn)].astype( |
| out_dtype |
| ), |
| name="output_unpack", |
| tag="conv3d_ndhwc", |
| ) |
| return conv_unpacked |
| |
| |
| def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, groups, out_dtype): |
| out_dtype = data.dtype if out_dtype is None else out_dtype |
| |
| assert isinstance(dilation, int) or len(dilation) == 3 |
| if isinstance(dilation, int): |
| dilation_d, dilation_h, dilation_w = (dilation, dilation, dilation) |
| else: |
| dilation_d, dilation_h, dilation_w = dilation |
| |
| DSTR, HSTR, WSTR = strides |
| batch_size, in_channel, in_depth, in_height, in_width = get_const_tuple(data.shape) |
| num_filter, _, kernel_depth, kernel_height, kernel_width = get_const_tuple(kernel.shape) |
| |
| dilated_kernel_d = (kernel_depth - 1) * dilation_d + 1 |
| dilated_kernel_h = (kernel_height - 1) * dilation_h + 1 |
| dilated_kernel_w = (kernel_width - 1) * dilation_w + 1 |
| |
| pad_front, pad_top, pad_left, pad_back, pad_down, pad_right = get_pad_tuple3d( |
| padding, (dilated_kernel_d, dilated_kernel_h, dilated_kernel_w) |
| ) |
| |
| pad_d = pad_front + pad_back |
| pad_h = pad_top + pad_down |
| pad_w = pad_left + pad_right |
| |
| pad_depth = in_depth + pad_d |
| pad_height = in_height + pad_h |
| pad_width = in_width + pad_w |
| |
| out_depth = simplify((in_depth + pad_d - dilated_kernel_d) // DSTR + 1) |
| out_height = simplify((in_height + pad_h - dilated_kernel_h) // HSTR + 1) |
| out_width = simplify((in_width + pad_w - dilated_kernel_w) // WSTR + 1) |
| |
| # pack data |
| DOPAD = pad_d != 0 or pad_h != 0 or pad_w != 0 |
| if DOPAD: |
| data_pad = pad( |
| data, |
| (0, 0, pad_front, pad_top, pad_left), |
| (0, 0, pad_back, pad_down, pad_right), |
| name="data_pad", |
| ) |
| else: |
| data_pad = data |
| |
| # fetch schedule |
| ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1] |
| |
| shape = (batch_size, in_channel // ic_bn, pad_depth, pad_height, ic_bn, pad_width) |
| data_vec = te.compute( |
| shape, lambda n, C, d, h, c, w: data_pad[n, C * ic_bn + c, d, h, w], name="data_vec" |
| ) |
| |
| ci_tile = in_channel // groups // ic_bn |
| if ci_tile == 0 or ci_tile * ic_bn * groups < in_channel: |
| ci_tile += 1 |
| |
| # pack kernel |
| shape = (num_filter // oc_bn, ci_tile, kernel_depth, kernel_height, kernel_width, ic_bn, oc_bn) |
| kernel_vec = te.compute( |
| shape, |
| lambda CO, CI, d, h, w, ci, co: kernel[CO * oc_bn + co, CI * ic_bn + ci, d, h, w], |
| name="kernel_vec", |
| ) |
| |
| # convolution |
| oshape = (batch_size, num_filter // oc_bn, out_depth, out_height, out_width, oc_bn) |
| unpack_shape = (batch_size, num_filter, out_depth, out_height, out_width) |
| |
| ic = te.reduce_axis((0, in_channel // groups), name="ic") |
| kh = te.reduce_axis((0, kernel_height), name="kh") |
| kw = te.reduce_axis((0, kernel_width), name="kw") |
| kd = te.reduce_axis((0, kernel_depth), name="kd") |
| idxmod = tvm.tir.indexmod |
| idxdiv = tvm.tir.indexdiv |
| |
| conv = te.compute( |
| oshape, |
| lambda n, oc_chunk, od, oh, ow, oc_block: te.sum( |
| data_vec[ |
| n, |
| idxdiv( |
| (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups) |
| + ic, |
| ic_bn, |
| ), |
| od * DSTR + kd * dilation_d, |
| oh * HSTR + kh * dilation_h, |
| idxmod( |
| (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups) |
| + ic, |
| ic_bn, |
| ), |
| ow * WSTR + kw * dilation_w, |
| ].astype(out_dtype) |
| * kernel_vec[ |
| oc_chunk, idxdiv(ic, ic_bn), kd, kh, kw, idxmod(ic, ic_bn), oc_block |
| ].astype(out_dtype), |
| axis=[ic, kd, kh, kw], |
| ), |
| name="conv", |
| ) |
| conv_unpacked = te.compute( |
| unpack_shape, |
| lambda n, c, d, h, w: conv[n, idxdiv(c, oc_bn), d, h, w, idxmod(c, oc_bn)].astype( |
| out_dtype |
| ), |
| name="output_unpack", |
| tag="conv3d_ncdhw", |
| ) |
| return conv_unpacked |
| |
| |
| def _create_tuning_space(cfg, data, kernel, strides, padding, dilation, groups, layout): |
| """Create schedule configuration from input arguments""" |
| dshape = get_const_tuple(data.shape) |
| kshape = get_const_tuple(kernel.shape) |
| if layout == "NDHWC": |
| n, d, h, w, ic = dshape |
| kd, kh, kw, _, oc = kshape |
| elif layout == "NCDHW": |
| n, ic, d, h, w = dshape |
| oc, _, kd, kh, kw = kshape |
| else: |
| raise ValueError(f"Not support this layout {layout} with schedule template.") |
| |
| # pad_front, pad_top, pad_left, pad_back, pad_down(bottom), pad_right |
| pf, pt, pl, pb, pd, pr = get_pad_tuple3d(padding, (kd, kh, kw)) |
| sd, sh, sw = strides if isinstance(strides, (tuple, list)) else (strides, strides, strides) |
| od = (d - kd + pf + pb) // sd + 1 |
| oh = (h - kh + pt + pd) // sh + 1 |
| ow = (w - kw + pl + pr) // sw + 1 |
| |
| # Create schedule config |
| cfg.define_split("tile_ic", ic, num_outputs=2) |
| cfg.define_split("tile_oc", oc, num_outputs=2) |
| cfg.define_split("tile_ow", ow, num_outputs=2, filter=lambda y: y.size[-1] <= 8) |
| cfg.define_knob("unroll_kw", [True, False]) |
| |
| |
| def _get_default_config(cfg, data, kernel, strides, padding, groups, out_dtype, layout): |
| """ |
| Get default schedule config for the workload |
| """ |
| if layout not in ["NDHWC", "NCDHW"]: |
| raise ValueError(f"Layout {layout} is not supported") |
| |
| static_data_shape = [] |
| for dim in get_const_tuple(data.shape): |
| if isinstance(dim, tvm.tir.Var): |
| static_data_shape.append(1) |
| else: |
| static_data_shape.append(dim) |
| data = te.placeholder(static_data_shape, dtype=data.dtype) |
| wkl = _get_conv3d_workload(data, kernel, strides, padding, groups, out_dtype, layout) |
| _fallback_schedule(cfg, wkl) |
| |
| |
| def _get_conv3d_workload(data, kernel, stride, padding, groups, out_dtype, data_layout="NCHW"): |
| """Get the workload structure.""" |
| if data_layout == "NCDHW": |
| _, CI, ID, IH, IW = get_const_tuple(data.shape) |
| CO, CIG, KD, KH, KW = get_const_tuple(kernel.shape) |
| elif data_layout == "NDHWC": |
| _, ID, IH, IW, CI = get_const_tuple(data.shape) |
| KD, KH, KW, CIG, CO = get_const_tuple(kernel.shape) |
| else: |
| raise ValueError(f"not support this layout {data_layout} yet") |
| |
| pad_front, pad_top, pad_left, pad_back, pad_down, pad_right = get_pad_tuple3d( |
| padding, (get_const_int(KD), get_const_int(KH), get_const_int(KW)) |
| ) |
| DPAD = pad_front + pad_back |
| HPAD = pad_top + pad_down |
| WPAD = pad_left + pad_right |
| if isinstance(stride, (tuple, list)): |
| DSTR, HSTR, WSTR = stride |
| else: |
| DSTR, HSTR, WSTR = stride, stride, stride |
| assert (data.dtype == kernel.dtype) or ( |
| data.dtype == "uint8" and kernel.dtype == "int8" |
| ), f"Do not support inputs with different data types now. {data.dtype} vs. {kernel.dtype}" |
| return Workload3D( |
| data.dtype, |
| out_dtype, |
| ID, |
| IH, |
| IW, |
| CI, |
| groups, |
| CO, |
| KD, |
| KH, |
| KW, |
| DPAD, |
| HPAD, |
| WPAD, |
| DSTR, |
| HSTR, |
| WSTR, |
| ) |
| |
| |
| def _fallback_schedule(cfg, wkl): |
| simd_width = get_simd_32bit_lanes() |
| DPAD, HPAD, WPAD = wkl.dpad, wkl.hpad, wkl.wpad |
| DSTR, HSTR, WSTR = wkl.dstride, wkl.hstride, wkl.wstride |
| out_width = (wkl.width + 2 * WPAD - wkl.wkernel) // WSTR + 1 |
| |
| oc_bn = 1 |
| for bn in range(simd_width, 0, -1): |
| if wkl.out_filter % bn == 0: |
| oc_bn = bn |
| break |
| |
| ic_bn = 1 |
| for bn in range(oc_bn, 0, -1): |
| if wkl.in_filter % bn == 0: |
| ic_bn = bn |
| break |
| |
| reg_n = 1 |
| for n in range(7, 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 _schedule_conv3d_ndhwc(s, cfg, data, data_pad, data_vec, kernel_vec, conv_out, output, last): |
| # fetch schedule |
| ic_bn, oc_bn, reg_n, unroll_kw = ( |
| cfg["tile_ic"].size[-1], |
| cfg["tile_oc"].size[-1], |
| cfg["tile_ow"].size[-1], |
| cfg["unroll_kw"].val, |
| ) |
| |
| # get padding size |
| padding = infer_pad3d(data, data_pad, "NDHWC") |
| DPAD, HPAD, WPAD = padding |
| DOPAD = DPAD != 0 or HPAD != 0 or WPAD != 0 |
| |
| A, W = data, kernel_vec |
| A0, A1 = data_pad, data_vec |
| |
| # schedule data |
| if DOPAD: |
| s[A0].compute_inline() |
| batch, ic_chunk, idd, ih, ic_block, iw = s[A1].op.axis |
| parallel_axis = s[A1].fuse(batch, ic_chunk, idd, ih) |
| s[A1].parallel(parallel_axis) |
| |
| # schedule kernel pack |
| oc_chunk, ic_chunk, od, oh, ow, ic_block, oc_block = s[W].op.axis |
| s[W].reorder(oc_chunk, od, oh, ic_chunk, ow, ic_block, oc_block) |
| if oc_bn > 1: |
| s[W].vectorize(oc_block) |
| parallel_axis = s[W].fuse(oc_chunk, od, oh) |
| s[W].parallel(parallel_axis) |
| |
| # schedule conv |
| C, O0, O = conv_out, output, last |
| CC = s.cache_write(C, "global") |
| |
| _, oc_chunk, od, oh, ow, oc_block = s[C].op.axis |
| ow_chunk, ow_block = s[C].split(ow, factor=reg_n) |
| s[C].reorder(oc_chunk, od, oh, ow_chunk, ow_block, oc_block) |
| s[C].fuse(oc_chunk, od, oh) |
| s[C].vectorize(oc_block) |
| |
| s[CC].compute_at(s[C], ow_chunk) |
| _, oc_chunk, od, oh, ow, oc_block = s[CC].op.axis |
| kd, kh, kw, ic = s[CC].op.reduce_axis |
| |
| ow_chunk, ow_block = s[CC].split(ow, factor=reg_n) |
| ic_chunk, ic_block = s[CC].split(ic, factor=ic_bn) |
| |
| if unroll_kw: |
| s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kd, kh, ic_block, kw, ow_block, oc_block) |
| s[CC].unroll(kw) |
| else: |
| s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kd, kh, kw, ic_block, ow_block, oc_block) |
| |
| s[CC].fuse(oc_chunk, od, oh) |
| s[CC].vectorize(oc_block) |
| s[CC].unroll(ow_block) |
| |
| if O0 != O: |
| s[O0].compute_inline() |
| |
| # unpacking |
| batch, od, oh, ow, oc = 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, od, oh, ow_chunk, ow_block, oc_block) |
| parallel_axis = s[O].fuse(batch, oc_chunk, od, oh) |
| s[C].compute_at(s[O], parallel_axis) |
| s[O].vectorize(oc_block) |
| s[O].parallel(parallel_axis) |
| return s |
| |
| |
| def _schedule_conv3d_ncdhw(s, cfg, data, data_pad, data_vec, kernel_vec, conv_out, output, last): |
| # fetch schedule |
| ic_bn, oc_bn, reg_n, unroll_kw = ( |
| cfg["tile_ic"].size[-1], |
| cfg["tile_oc"].size[-1], |
| cfg["tile_ow"].size[-1], |
| cfg["unroll_kw"].val, |
| ) |
| |
| # get padding size |
| padding = infer_pad3d(data, data_pad, "NCDHW") |
| DPAD, HPAD, WPAD = padding |
| DOPAD = DPAD != 0 or HPAD != 0 or WPAD != 0 |
| |
| A, W = data, kernel_vec |
| A0, A1 = data_pad, data_vec |
| |
| # schedule data |
| if DOPAD: |
| s[A0].compute_inline() |
| batch, ic_chunk, idd, ih, ic_block, iw = s[A1].op.axis |
| parallel_axis = s[A1].fuse(batch, ic_chunk, idd, ih) |
| s[A1].parallel(parallel_axis) |
| |
| # schedule kernel pack |
| oc_chunk, ic_chunk, od, oh, ow, ic_block, oc_block = s[W].op.axis |
| s[W].reorder(oc_chunk, od, oh, ic_chunk, ow, ic_block, oc_block) |
| if oc_bn > 1: |
| s[W].vectorize(oc_block) |
| parallel_axis = s[W].fuse(oc_chunk, od, oh) |
| s[W].parallel(parallel_axis) |
| |
| # schedule conv |
| C, O0, O = conv_out, output, last |
| CC = s.cache_write(C, "global") |
| |
| _, oc_chunk, od, oh, ow, oc_block = s[C].op.axis |
| ow_chunk, ow_block = s[C].split(ow, factor=reg_n) |
| s[C].reorder(oc_chunk, od, oh, ow_chunk, ow_block, oc_block) |
| s[C].fuse(oc_chunk, od, oh) |
| s[C].vectorize(oc_block) |
| |
| s[CC].compute_at(s[C], ow_chunk) |
| _, oc_chunk, od, oh, ow, oc_block = s[CC].op.axis |
| ic, kd, kh, kw = s[CC].op.reduce_axis |
| |
| ow_chunk, ow_block = s[CC].split(ow, factor=reg_n) |
| ic_chunk, ic_block = s[CC].split(ic, factor=ic_bn) |
| |
| if unroll_kw: |
| s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kd, kh, ic_block, kw, ow_block, oc_block) |
| s[CC].unroll(kw) |
| else: |
| s[CC].reorder(oc_chunk, oh, ow_chunk, ic_chunk, kd, kh, kw, ic_block, ow_block, oc_block) |
| |
| s[CC].fuse(oc_chunk, od, oh) |
| s[CC].vectorize(oc_block) |
| s[CC].unroll(ow_block) |
| |
| if O0 != O: |
| s[O0].compute_inline() |
| |
| # unpacking |
| batch, oc, od, 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, od, oh, ow_chunk, ow_block, oc_block) |
| parallel_axis = s[O].fuse(batch, oc_chunk, od, oh) |
| s[C].compute_at(s[O], parallel_axis) |
| s[O].vectorize(oc_block) |
| s[O].parallel(parallel_axis) |
| |
| return s |
