in tensorflow/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h [9430:9947]
inline void DepthwiseConvDotProduct3x3(
const DepthwiseParams& params, const RuntimeShape& input_shape,
const uint8* input_data, const RuntimeShape& filter_shape,
const uint8* filter_data, const RuntimeShape& bias_shape,
const int32* bias_data, const RuntimeShape& output_shape,
uint8* output_data, int thread_start, int thread_end, int thread_dim) {
// Check kernel restrictions.
constexpr int filter_size = 3;
constexpr int kMaxStride = 2;
constexpr int kMaxPadding = 1;
constexpr int kSymmetricZeroPoint = 128;
TFLITE_DCHECK_EQ(params.weights_offset, -kSymmetricZeroPoint);
TFLITE_DCHECK_LE(params.stride_width, kMaxStride);
TFLITE_DCHECK_EQ(params.stride_height, params.stride_width);
TFLITE_DCHECK_EQ(params.dilation_width_factor, 1);
TFLITE_DCHECK_EQ(params.dilation_height_factor, 1);
TFLITE_DCHECK_LE(params.padding_values.width, kMaxPadding);
TFLITE_DCHECK_LE(params.padding_values.height, kMaxPadding);
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
// Key kernel parameters (along with padding handled later).
const int stride = params.stride_width;
const int depth_multiplier = params.depth_multiplier;
const bool has_depth_multiplication = depth_multiplier > 1;
// Extract task dimensions.
const int input_depth = input_shape.Dims(3);
const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
const int input_height = input_shape.Dims(1);
const int input_width = input_shape.Dims(2);
const int output_height = output_shape.Dims(1);
const int output_width = output_shape.Dims(2);
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
TFLITE_DCHECK(!has_depth_multiplication || input_depth == 1);
TFLITE_DCHECK(has_depth_multiplication || input_depth == output_depth);
TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth);
TFLITE_DCHECK_EQ(input_depth * depth_multiplier, output_depth);
TFLITE_DCHECK_EQ(MatchingDim(filter_shape, 1, filter_shape, 2), filter_size);
TFLITE_DCHECK(thread_dim == 0 || thread_dim == 1);
// Return now if nothing to do.
if (output_width == 0 || output_height == 0) {
return;
}
// Kernel parameter structure: set basic fields.
//
// In asm it is easier to pass a structure than more than, say, 8 parameters.
DepthwiseConvDotProdParams function_params;
function_params.input_depth = input_depth;
function_params.output_depth = output_depth;
function_params.input_offset = params.input_offset;
function_params.output_offset = params.output_offset;
function_params.output_multiplier = params.output_multiplier;
function_params.output_shift = params.output_shift;
function_params.quantized_activation_min = params.quantized_activation_min;
function_params.quantized_activation_max = params.quantized_activation_max;
function_params.stride = stride;
// Handle inbound bias data.
//
// Note that this data is adjusted in a per-depth process before the main
// filters. The adjustment accounts for a non-symmetric input offset.
//
// Kernel subroutines need to be able to operate consistently on an bias
// array. Where there is no bias, we provide one filled with zeros.
constexpr int kMinBiasLoad = 8;
int32 zero_bias_data[kMinBiasLoad];
int32 bias_increment;
if (bias_data) {
bias_increment = 4;
} else {
memset(zero_bias_data, 0, sizeof(zero_bias_data));
bias_data = &zero_bias_data[0];
bias_increment = 0;
}
function_params.bias_increment = bias_increment;
TFLITE_DCHECK_LE(2 * function_params.bias_increment, kMinBiasLoad);
// Process multithreading.
int batch_start = 0;
int batch_end = batches;
int row_start = 0;
int row_end = output_height;
switch (thread_dim) {
case 0:
TFLITE_DCHECK_GE(thread_start, 0);
TFLITE_DCHECK_LE(thread_end, batches);
batch_start = thread_start;
batch_end = thread_end;
break;
case 1:
TFLITE_DCHECK_GE(thread_start, 0);
TFLITE_DCHECK_LE(thread_end, output_height);
row_start = thread_start;
row_end = thread_end;
break;
}
const int row_count = row_end - row_start;
// Process padding.
//
// Whether "correct" or not, this matches ComputeConvSizes. When there is
// stride > 1 there can be padding on the bottom or top, and therefore
// we need to consider padding. This is true even if one or other of the
// padding_values is 0.
const int padded_width = (output_width - 1) * stride + filter_size;
int full_padding_top;
{
const int padding_left = params.padding_values.width;
// Right padding would be -1 if discarding input because of stride.
const int padding_right =
std::max(padded_width - input_width - padding_left, 0);
int padding_top = params.padding_values.height;
const int padded_height = (output_height - 1) * stride + filter_size;
int padding_bottom =
std::max(padded_height - input_height - padding_top, 0);
TFLITE_DCHECK_LE(padding_left, padding_right);
TFLITE_DCHECK_LE(padding_top, padding_bottom);
full_padding_top = padding_top;
if (row_start != 0) {
padding_top = 0;
}
if (row_end != output_height) {
padding_bottom = 0;
}
function_params.padding_left = padding_left;
function_params.padding_right = padding_right;
function_params.padding_top = padding_top;
function_params.padding_bottom = padding_bottom;
}
// When stride == 1 left or top padding may only be non-zero.
// This is when padding is specified but not needed on a trailing dimension.
// When stride == 2 right or bottom padding may only be non-zero.
// This is a result of the details of the padding calculations.
const bool padding_required =
function_params.padding_left > 0 || function_params.padding_top > 0 ||
function_params.padding_right > 0 || function_params.padding_bottom > 0;
// Choose parameter-specific kernel subroutines.
//
// The main part of the kernel has two stages. First, a temporary workspace is
// filled with padded and permuted data. Second, the filter is applied to the
// workspace data to generate output.
//
// The workspace fill stage handles padding so that the filter stage does not
// need to account for it. The workspace fill stage does not need to
// understand striding, and implicitly handles striding through the parameters
// that it is given.
using pack_macro_block_func_t = decltype(
&PackMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kNoMultiplication,
0>::Run);
using kernel_macro_block_func_t = decltype(
&KernelMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kNoMultiplication,
1>::Run);
pack_macro_block_func_t pack_macro_block_func;
kernel_macro_block_func_t kernel_macro_block_func;
{
if (has_depth_multiplication) {
if (padding_required) {
pack_macro_block_func =
PackMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kUnitInputDepth,
/*max_padding=*/1>::Run;
} else {
pack_macro_block_func =
PackMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kUnitInputDepth,
/*max_padding=*/0>::Run;
}
if (stride == 1) {
kernel_macro_block_func =
KernelMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kUnitInputDepth,
/*stride=*/1>::Run;
} else {
kernel_macro_block_func =
KernelMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kUnitInputDepth,
/*stride=*/2>::Run;
}
} else {
if (padding_required) {
pack_macro_block_func =
PackMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kNoMultiplication,
/*max_padding=*/1>::Run;
} else {
pack_macro_block_func =
PackMacroBlock<implementation,
DepthwiseConvDepthMultiplication::kNoMultiplication,
/*max_padding=*/0>::Run;
}
if (stride == 1) {
kernel_macro_block_func = KernelMacroBlock<
implementation, DepthwiseConvDepthMultiplication::kNoMultiplication,
/*stride=*/1>::Run;
} else {
kernel_macro_block_func = KernelMacroBlock<
implementation, DepthwiseConvDepthMultiplication::kNoMultiplication,
/*stride=*/2>::Run;
}
}
}
// Stride-only variables.
//
const int row_count_per_macro = stride == 1 ? 4 : 2;
// row_count_per_macro * stride:
constexpr int input_height_per_macro = 4;
// Number of rows per micro block (= rows per macro block) is
// (row_count_per_macro - 1) * stride + 1 + (filter_size - 1)
const int height_block_size = stride == 1 ? 3 + filter_size : 2 + filter_size;
const int input_height_overlap = filter_size - stride;
// stride == 1 ? 4 : 2:
function_params.four_over_stride = row_count_per_macro;
TFLITE_DCHECK_EQ(stride * function_params.four_over_stride, 4);
TFLITE_DCHECK_EQ(height_block_size,
input_height_per_macro + input_height_overlap);
// Create workspaces.
//
// Filter workspace is for shuffle: only first depth/8 is used.
// indexed as [depth/8][sub-block][height][depth][width].
TFLITE_DCHECK_EQ(kDepthwiseConvAdjustedBiasLimit % 8, 0);
int8 macroblock_workspace[kDepthwiseConvScratchWorkspaceSize];
int32 adjusted_bias_data[kDepthwiseConvAdjustedBiasLimit];
int8 filter_workspace[kDepthwiseConvAdjustedBiasLimit >> 3][3][2][4][4];
// Output depth characterization.
//
const int depth_macro_count = output_depth / 64;
const int depth_overall_macro_count = (output_depth + 63) / 64;
// Number of micro blocks down the depth in a final incomplete macro block.
const int depth_trailing_micro_repeats = output_depth / 8 % 8;
// The output_depth may not have a remainder: it must be a multiple of 8.
TFLITE_DCHECK_EQ(output_depth,
64 * depth_macro_count + 8 * depth_trailing_micro_repeats);
// Characterize the first macro block depth, the largest.
//
// We base treatment of the width on the trailing macro block if there are
// no full blocks, in order to do more work together (that is, increase
// workspace_width_micro_repeats when largest_macro_depth < 64).
const int largest_macro_depth =
has_depth_multiplication
? 1
: (depth_macro_count > 0 ? 64 : 8 * depth_trailing_micro_repeats);
// Characterize width, consumption of input and generation of output.
//
// In the case of depth multiplication, we ensure that some of the workspace
// at the end remains unused. This enables the filter routines to load the
// "next" data, of at least 16 bytes, even when at the end of the workspace.
// It is relatively expensive to detect the end micro block. It is also very
// difficult to test for (to trigger) erroneous reads (past end of array) in
// the depth multplication case.
int workspace_width_micro_repeats =
(has_depth_multiplication
? kDepthwiseConvScratchWorkspaceSize - kWorkspaceExtension
: kDepthwiseConvScratchWorkspaceSize) /
(4 * largest_macro_depth * height_block_size);
// When there is no depth multiplication, the workspace depth is a multiple of
// 8, which ensures that workspace rows are 16-byte aligned. (Actually 32,
// because of the micro width of 4.) This is not necessarily the case under
// depth multiplication, so we adjust now to impose this restriction.
if (has_depth_multiplication) {
workspace_width_micro_repeats = (workspace_width_micro_repeats / 4) * 4;
}
TFLITE_DCHECK_EQ((workspace_width_micro_repeats * largest_macro_depth) % 4,
0);
// Discount 1 of the micro-block repeats in each macro block to account for
// overlap.
const int consumed_width_per_macro_block =
4 * (workspace_width_micro_repeats - 1);
const int output_width_per_macro_block =
function_params.four_over_stride * (workspace_width_micro_repeats - 1);
TFLITE_DCHECK_GT(workspace_width_micro_repeats, 1);
TFLITE_DCHECK_EQ(output_width_per_macro_block * stride,
consumed_width_per_macro_block);
// Width repetitions and residuals.
//
// Use of the workspace is characterized primarily in terms of *padded input*.
// Striding only matters in a few places.
//
// Simplifications: We require that there always be at least one full
// micro-block across the width. Since the maximum padding is 1, the trailing
// padding cannot span two micro blocks.
const int residual_micro_width = padded_width % 4;
// We base the count of macro blocks on the amount of padded input data each
// one consumes.
int width_overall_macro_count = (padded_width - residual_micro_width +
consumed_width_per_macro_block - 1) /
consumed_width_per_macro_block;
// Recall that we left a micro block at the end of each macro block for use as
// overlap. There is a special case in which we can use one fewer macro
// blocks, with the last one consuming extra input. (But not if the
// calculation thinks that we can use zero blocks.)
if (padded_width <=
((width_overall_macro_count - 1) * consumed_width_per_macro_block + 4)) {
width_overall_macro_count -= 1;
}
width_overall_macro_count = std::max(width_overall_macro_count, 1);
// We always have to treat the final macro block along width as trailing,
// because even if it is full in terms of padded input, it will be incomplete
// in terms of output.
const int width_macro_count = width_overall_macro_count - 1;
// Micro blocks are traversed in terms of input in fill routines.
const int width_trailing_micro_repeats =
(padded_width - consumed_width_per_macro_block * width_macro_count) / 4;
const int width_overall_trailing_micro_repeats =
(padded_width - consumed_width_per_macro_block * width_macro_count + 3) /
4;
// Micro blocks are traversed in terms of output in filtering routines.
const int residual_output_micro_width =
(output_width - 1) % function_params.four_over_stride + 1;
const int output_width_trailing_micro_repeats =
residual_micro_width > (filter_size - 1)
? width_trailing_micro_repeats
: width_trailing_micro_repeats - 1;
// Check results.
TFLITE_DCHECK_GT(width_overall_trailing_micro_repeats, 0);
TFLITE_DCHECK_EQ(padded_width,
residual_micro_width +
consumed_width_per_macro_block * width_macro_count +
4 * width_trailing_micro_repeats);
TFLITE_DCHECK_LE(width_overall_macro_count, width_macro_count + 1);
TFLITE_DCHECK_GE(width_overall_macro_count, width_macro_count);
// Height repetitions and residuals.
//
int height_macro_count;
int residual_row_count;
int height_overall_macro_count;
if (stride == 1) {
TFLITE_DCHECK_EQ(row_count_per_macro, 4);
height_macro_count = row_count / 4;
residual_row_count = row_count % 4;
height_overall_macro_count = (row_count + 3) / 4;
} else {
TFLITE_DCHECK_EQ(row_count_per_macro, 2);
height_macro_count = row_count / 2;
residual_row_count = row_count % 2;
height_overall_macro_count = (row_count + 1) / 2;
}
TFLITE_DCHECK_EQ(
row_count, residual_row_count + row_count_per_macro * height_macro_count);
TFLITE_DCHECK_LE(height_overall_macro_count, height_macro_count + 1);
TFLITE_DCHECK_GE(height_overall_macro_count, height_macro_count);
// Data strides.
//
const int input_height_stride = input_width * input_depth;
const int output_height_stride = output_width * output_depth;
const int input_batch_stride = input_height_stride * input_height;
const int output_batch_stride = output_height_stride * output_height;
const int input_depth_macro_stride = has_depth_multiplication ? 0 : 64;
const int input_width_macro_stride =
input_depth * consumed_width_per_macro_block;
const int output_width_macro_stride =
output_depth * output_width_per_macro_block;
// Store parameters that do not vary across macro blocks.
//
function_params.workspace_width_micro_repeats = workspace_width_micro_repeats;
function_params.height_macro_count = height_overall_macro_count;
function_params.width_macro_count = width_overall_macro_count;
function_params.input_height_stride = input_height_stride;
function_params.output_height_stride = output_height_stride;
function_params.residual_width = residual_micro_width;
// Prefetch workspace for write, along with any necessary dummy writes.
const int max_workspace_height_stride =
16 * ((workspace_width_micro_repeats + 3) >> 2) * largest_macro_depth;
const int workspace_fill_size = std::min(
kDepthwiseConvScratchWorkspaceSize,
height_block_size * max_workspace_height_stride + kWorkspaceExtension);
WorkspacePrefetchWrite<implementation>::Run(
params.weights_offset, workspace_fill_size, macroblock_workspace);
// Main process.
//
// Most kernels are nested batch-height-width-depth. Here we proceed over
// macro blocks batch-width-depth-height.
//
// Example of handling of trailing iteration: when there is trailing depth,
// depth_overall_macro_count = depth_macro_count + 1, so we can adjust the
// dimensions for trailing macro blocks by looking for
// j_depth == depth_macro_count.
for (int b = batch_start; b < batch_end; ++b) {
for (int k_width = 0; k_width < width_overall_macro_count; ++k_width) {
// Figure out the work to be done for this macro block. If it trails in
// any dimension, the work in that dimension is adjusted.
// The work to be done across widths has 3 cases:
// (a) A full macro block,
// (b) Partial terminal macro block, with input and output ending in
// same micro block, and
// (c) Partial terminal macro block, with output corresponding to one
// fewer micro blocks, because filter extends across micro-block
// boundary.
if (k_width != width_macro_count) {
function_params.output_residual_width = 0;
function_params.input_width_micro_repeats =
workspace_width_micro_repeats;
function_params.input_width_overall_micro_repeats =
workspace_width_micro_repeats;
function_params.output_width_micro_repeats =
workspace_width_micro_repeats - 1;
} else {
function_params.output_residual_width = residual_output_micro_width;
function_params.input_width_micro_repeats =
width_trailing_micro_repeats;
function_params.input_width_overall_micro_repeats =
width_overall_trailing_micro_repeats;
function_params.output_width_micro_repeats =
output_width_trailing_micro_repeats;
}
function_params.output_width_overall_micro_repeats =
function_params.output_residual_width == 0
? function_params.output_width_micro_repeats
: function_params.output_width_micro_repeats + 1;
for (int j_depth = 0; j_depth < depth_overall_macro_count; ++j_depth) {
// Process filter and bias data.
//
function_params.depth_micro_repeats =
j_depth == depth_macro_count ? depth_trailing_micro_repeats : 8;
ProcessPerDepth<implementation>::Run(
filter_data + 64 * j_depth,
bias_data + 8 * 2 * bias_increment * j_depth,
filter_workspace[0][0][0][0], adjusted_bias_data, &function_params);
const uint8* input_data_block =
input_data + b * input_batch_stride +
j_depth * input_depth_macro_stride +
k_width * input_width_macro_stride -
function_params.padding_left * input_depth +
row_start * stride * input_height_stride -
full_padding_top * input_height_stride;
uint8* output_data_block = output_data + b * output_batch_stride +
row_start * output_height_stride +
j_depth * 64 +
k_width * output_width_macro_stride;
// Under depth multiplication the workspace_height_stride does not have
// to depend on input_width_overall_micro_repeats, but this improves the
// compactness of workspace use.
const int workspace_height_stride =
has_depth_multiplication
? 16 * ((function_params.input_width_overall_micro_repeats +
3) >>
2)
: 4 * function_params.input_width_overall_micro_repeats * 8 *
function_params.depth_micro_repeats;
TFLITE_DCHECK_EQ(workspace_height_stride % 16, 0);
function_params.workspace_height_stride = workspace_height_stride;
// For the first macro block for output rows we fill in the first few
// rows. After this we will copy them (see below in loop.)
function_params.inbound_block_height = input_height_overlap;
pack_macro_block_func(-1, k_width, input_data_block,
macroblock_workspace, &function_params);
input_data_block += input_height_stride * input_height_overlap;
for (int i_height = 0; i_height < height_overall_macro_count;
++i_height) {
if (i_height != height_macro_count) {
function_params.inbound_block_height = input_height_per_macro;
function_params.outbound_block_height = row_count_per_macro;
} else {
function_params.inbound_block_height = residual_row_count * stride;
function_params.outbound_block_height = residual_row_count;
}
TFLITE_DCHECK_LT(i_height * row_count_per_macro, row_count);
TFLITE_DCHECK_LT(i_height * input_height_per_macro, input_height);
TFLITE_DCHECK_LT(k_width * output_width_per_macro_block,
output_width);
TFLITE_DCHECK_LT(k_width * consumed_width_per_macro_block,
input_width);
// Macro blocks overlap by input_height_overlap rows, so we copy
// those instead of filling in afresh. The first macro block across
// output rows was filled in outside of the loop (above).
if (i_height > 0) {
memcpy(macroblock_workspace,
macroblock_workspace +
input_height_per_macro * workspace_height_stride,
input_height_overlap * workspace_height_stride);
}
pack_macro_block_func(
i_height, k_width, input_data_block,
macroblock_workspace +
input_height_overlap * workspace_height_stride,
&function_params);
kernel_macro_block_func(
macroblock_workspace, filter_workspace[0][0][0][0],
adjusted_bias_data, output_data_block, &function_params);
input_data_block += input_height_stride * input_height_per_macro;
output_data_block += output_height_stride * row_count_per_macro;
}
}
}
}
}