conv3d

Description

Computes 3D convolution of the float16 type with the given 6HD data and FracZ weight.

This API supports bias.

The data tensor has a 6HD shape of (N, D, C1, H, W, C0). The weight tensor has a FracZ shape of (KD x C1 x KH x KW, Cout//C0_out, C0_out, C0).

Prototype

conv3d(x, filter, filter_size, para_dict)

Parameters

  • x: a 6HD Tensor of type float16 or int8, for the feature map of 3D convolution.
  • filter: a FracZ Tensor of type float16 or int8, for the weight of 3D convolution.
  • filter_size: a weight for 3D convolution.
  • para_dict: a dictionary for the key-value pairs, including the following keys:
    • dsl_flag: whether automatic UB fusion is support. If set to True, enables UB fusion; if set to False, disables UB fusion. Defaults to False.
    • bias_tensor: whether computation is performed without bias. Defaults to None.
    • pads: a list of the padding lines along the D, H, and W directions of the feature map.
    • strides: a list of the strides along the D, H, and W directions of the feature map.
    • dilations: a list of the dilations along the D, H, and W directions of the filter.
    • res_dtype: output data type.
    • mad_dtype: data type of the mean absolute deviation (MAD).
    • kernel_name: operator name.
    • group_dict: a dictionary of 3d group convolution arguments, including:

      1. fmap_c, C dimension size of the feature map.

      2. cout, batch size of the weight.

      3. groups, argument for group convolution.

      4. cout0, which is tbe_platform.C0_SIZE. Defaults to 16.

      5. cin0, which is tbe_platform.C0_SIZE. Defaults to 16.

      The computation formula is as follows.

      lcm(param1, param2): calculates the least common multiple.

      mag_factor0 = lcm(fmap_c // groups, cin0) // (fmap_c // groups)

      mag_factor1 = lcm(cout // groups, cout0) // (cout // groups)

      mag_factor = min(lcm(mag_factor0, mag_factor1), groups)

      cin1_g = (mag_factor * fmap_c // groups + cin0 - 1) // cin0

      cout_g = (mag_factor * cout // groups + cout0 - 1) // cout0 * cout0

      group_dict = {"real_g": (groups + mag_factor - 1) // mag_factor,

      "mag_factor": mag_factor,

      "cin1_g": cin1_g,

      "cout_g": cout_g,

      "cin_ori": fmap_c,

      "cout_ori": cout}

Returns

res_tensor: result tensor.

Restrictions

This API cannot be used in conjunction with other TBE DSL APIs.

Applicability

Atlas 200/300/500 Inference Product

Atlas Training Series Product

Example

from tbe import tvm
from tbe import dsl

shape_fmp_ndc1hwc0 = (1, 32, 1, 240, 352, 16)
fmp_dtype = "float16"
shape_filter = [16, 16, 3, 3, 3]
shape_w_frac_z = (27, 1, 16, 16)
w_dtype = "float16"

data = tvm.placeholder(shape_fmp_ndc1hwc0, name='Fmap', dtype=fmp_dtype)
weight = tvm.placeholder(shape_w_frac_z, name='Filter', dtype=w_dtype)

bias_tensor = None
pads = [1, 1, 1, 1, 1, 1]
stride_dhw = [1, 1, 1]
res_dtype = "float16"
mad_dtype = "float32"
kernel_name = "conv3d_1_32_240_352_1_3_3_3_1_16_1_1_1_SAME_NDHWC_1_0"
group_dict = {'real_g': 1, 'mag_factor': 1, 'cin1_g': 1, 'cout_g': 16, 'cin_ori': 1, 'cout_ori': 16}
dilation_dhw = [1, 1, 1]

para_dict = {
    "dsl_flag": False,
    "bias_tensor": bias_tensor,
    "pads": pads,
    "strides": stride_dhw,
    "res_dtype": res_dtype,
    "mad_dtype": mad_dtype,
    "kernel_name": kernel_name,
    "group_dict": group_dict,
    "dilations": dilation_dhw
}

conv_res = dsl.conv3d(data, weight, shape_filter, para_dict)