Create an extern op to do inference convolution of 4D tensor data and 4D tensor kernel and 1D tensor bias with nnpack. Parameters ---------- data : Tensor data 4D tensor input[batch][input_channels][input_height][input_width] of FP32 elements. kernel : Tensor
(
data, kernel, bias, padding, stride, nthreads=1, algorithm=ConvolutionAlgorithm.AUTO
)
| 73 | |
| 74 | |
| 75 | def convolution_inference( |
| 76 | data, kernel, bias, padding, stride, nthreads=1, algorithm=ConvolutionAlgorithm.AUTO |
| 77 | ): |
| 78 | """Create an extern op to do inference convolution of 4D tensor data and |
| 79 | 4D tensor kernel and 1D tensor bias with nnpack. |
| 80 | |
| 81 | Parameters |
| 82 | ---------- |
| 83 | data : Tensor |
| 84 | data 4D tensor input[batch][input_channels][input_height][input_width] of |
| 85 | FP32 elements. |
| 86 | kernel : Tensor |
| 87 | kernel 4D tensor kernel[output_channels][input_channels][kernel_height] |
| 88 | [kernel_width] of FP32 elements. |
| 89 | bias : Tensor |
| 90 | bias 1D array bias[output_channels][input_channels][kernel_height] |
| 91 | [kernel_width] of FP32 elements. |
| 92 | padding : list |
| 93 | padding A 4-dim list of [pad_top, pad_bottom, pad_left, pad_right], |
| 94 | which indicates the padding around the feature map. |
| 95 | stride : list |
| 96 | stride A 2-dim list of [stride_height, stride_width], which indicates |
| 97 | the stride. |
| 98 | |
| 99 | Returns |
| 100 | ------- |
| 101 | output : Tensor |
| 102 | output 4D tensor output[batch][output_channels][output_height][output_width] |
| 103 | of FP32 elements. |
| 104 | """ |
| 105 | |
| 106 | assert isinstance(padding, list) and len(padding) == 4 |
| 107 | assert isinstance(stride, list) and len(stride) == 2 |
| 108 | batch, _, input_height, input_width = data.shape |
| 109 | output_channels, _, kernel_height, kernel_width = kernel.shape |
| 110 | idxdiv = te.indexdiv |
| 111 | output_height = idxdiv(input_height + padding[0] + padding[1] - kernel_height, stride[0]) + 1 |
| 112 | output_width = idxdiv(input_width + padding[0] + padding[1] - kernel_width, stride[1]) + 1 |
| 113 | |
| 114 | return te.extern( |
| 115 | (batch, output_channels, output_height, output_width), |
| 116 | [data, kernel, bias] if bias is not None else [data, kernel], |
| 117 | lambda ins, outs: tvm.tirx.call_packed( |
| 118 | "tvm.contrib.nnpack.convolution_inference", |
| 119 | ins[0], |
| 120 | ins[1], |
| 121 | ins[2] if bias is not None else 0, |
| 122 | outs[0], |
| 123 | padding[0], |
| 124 | padding[1], |
| 125 | padding[2], |
| 126 | padding[3], |
| 127 | stride[0], |
| 128 | stride[1], |
| 129 | nthreads, |
| 130 | algorithm, |
| 131 | ), |
| 132 | name="C", |
nothing calls this directly
no test coverage detected
searching dependent graphs…