kReference,

kGenericOptimized, // Neon-free

// kMultithreadOptimized is a mixture of an Eigen-based kernel when threads

// are available and kGenericOptimized when we must use only one thread.

kMultithreadOptimized,

// The kernel uses use CBLAS interface for matrix multiplication.

// It's fast when an optimized CBLAS implementation is available (e.g. Apple

// Accelerate Framework), and it's slow when falling back to naive

// implementation.

kCblasOptimized,

// IDs are the arbitrary identifiers used by TF Lite to identify and access

// memory buffers.

int im2col_id = kTensorNotAllocated;

int hwcn_weights_id = kTensorNotAllocated;

int input_quantized_id = kTensorNotAllocated;

int scaling_factors_id = kTensorNotAllocated;

int input_offset_id = kTensorNotAllocated;

int accum_scratch_id = kTensorNotAllocated;

// Row sums are used to cache filter sums for hybrid zero-point calculations.

int row_sums_id = kTensorNotAllocated;

TfLitePaddingValues padding;

// The scaling factor from input to output (aka the 'real multiplier') can

// be represented as a fixed point multiplier plus a left shift.

int32_t output_multiplier;

int output_shift;

// Per channel output multiplier and shift.

std::vector<int32_t> per_channel_output_multiplier;

std::vector<int> per_channel_output_shift;

// The range of the fused activation layer. For example for kNone and

// uint8_t these would be 0 and 255.

int32_t output_activation_min;

int32_t output_activation_max;

// Indexes are the offset to the memory buffer in the array used to keep track

// of the allocated temporaries.

int32_t im2col_index;

int32_t hwcn_weights_index;

int32_t input_quantized_index;

int32_t scaling_factors_index;

int32_t accum_scratch_index;

int32_t input_offset_index;

int32_t row_sums_index;

bool need_hwcn_weights = false;

bool have_weights_been_transposed = false;

bool need_im2col = false;

// If it's true, it means im2col is needed but gets disabled because the

// temporary im2col tensor requires too much memory (i.e.

// >= kMaxIm2colBufferSize);

bool im2col_oversized = false;

bool supports_multithreaded_kernel = false;

bool is_hybrid_per_channel = false;

bool compute_hybrid_row_sums = true;

// Number of convolution groups.

int32_t groups = 1;

TfLiteType quantized_bias_type = kTfLiteNoType;

switch (padding) {

case TfLitePadding::kTfLitePaddingSame:

return PaddingType::kSame;

case TfLitePadding::kTfLitePaddingValid:

return PaddingType::kValid;

case TfLitePadding::kTfLitePaddingUnknown:

default:

return PaddingType::kNone;

}

#if defined(TFLITE_WITH_MULTITHREADED_EIGEN)

eigen_support::DecrementUsageCounter(context);

#endif

delete reinterpret_cast<OpData*>(buffer);

const int rows = output->dims->data[1];

const int cols = output->dims->data[0];

const float* input_data = GetTensorData<float>(input);

float* output_data = GetTensorData<float>(output);

for (int i = 0; i < rows; ++i) {

for (int j = 0; j < cols; ++j) {

const float in_value = input_data[i * cols + j];

output_data[j * rows + i] = in_value;

}

}

const TfLiteTensor* filter, OpData* data, bool is_hybrid,

KernelType kernel_type) {

// If HWCN weights are required, Im2Col not required

if (data->need_hwcn_weights) return false;

// segregate based on dilated conv & non-dialated conv

const bool need_dilated_im2col =

params->dilation_width_factor != 1 || params->dilation_height_factor != 1;

const bool need_non_dilated_im2col =

params->stride_width != 1 || params->stride_height != 1 ||

filter->dims->data[2] != 1 || filter->dims->data[1] != 1;

const bool need_im2col = need_dilated_im2col || need_non_dilated_im2col;

// Return early as basic requirement is not met

if (!need_im2col) return false;

switch (kernel_type) {

case kReference:

if (is_hybrid) {

return true;

} else {

return false;

}

case kGenericOptimized:

case kCblasOptimized:

// `need_im2col` is always satisfied.

return true;

case kMultithreadOptimized:

if (input->type == kTfLiteUInt8 || //

input->type == kTfLiteInt8 || //

input->type == kTfLiteInt16 || // quantized.

!data->supports_multithreaded_kernel) {

return true;

} else {

return false;

}

default:

return false;

}

TfLiteContext* context, TfLiteNode* node, bool is_hybrid,

bool is_per_channel, KernelType kernel_type, size_t im2col_bytes) {

auto* params = reinterpret_cast<TfLiteConvParams*>(node->builtin_data);

OpData* data = reinterpret_cast<OpData*>(node->user_data);

TF_LITE_ENSURE(context, node->inputs->size >= 2);

const TfLiteTensor* input;

TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));

const TfLiteTensor* filter;

TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &filter));

// If we're using the optimized multithreaded EigenTensor implementation of

// convolution, it expects the filter weights to be transposed compared to

// the normal TF Lite buffer format. Typical TF Lite weights are

// [filter_count, filter_height, filter_width, input_depth], but for the float

// implementation we need them as [filter_height, filter_width, input_depth,

// filter_count]. We get to that format by transposing, and create a temporary

// buffer to store the results.

// This path is only used for float processing, so only create the buffer if

// we're running with that data type.

data->need_hwcn_weights =

input->type == kTfLiteFloat32 && data->supports_multithreaded_kernel;

// We don't always need to allocate im2col. It is only used in some versions

// of the optimized Conv. This test just mimics something that happens inside

// optimized_ops.h, in order to avoid a DCHECK(!im2col_data).

data->need_im2col =

IsIm2ColRequired(input, params, filter, data, is_hybrid, kernel_type);

// If im2col_oversized is found to be true, we have to fallback to an

// execution path (like kReference in float/quantized cases) that doesn't

// require im2col operation. Therefore, we have to skip checking the hybrid

// case (but not the hybrid-per-channel one) where there's no such a fallback

// execution path.

// TODO(b/178743262): Consider making this check conditioned on the available

// memory of the system, rather than coupling to the mobile platform check.

if (IsMobilePlatform() && !(is_hybrid && !is_per_channel) &&

data->need_im2col && im2col_bytes >= kMaxIm2colBufferSizeMobile) {

data->need_im2col = false;

data->im2col_oversized = true;

}

int temporaries_count = 0;

if (data->need_im2col) {

data->im2col_index = temporaries_count;

if (data->im2col_id == kTensorNotAllocated) {

context->AddTensors(context, 1, &data->im2col_id);

}

++temporaries_count;

}

if (data->need_hwcn_weights) {

data->hwcn_weights_index = temporaries_count;

if (data->hwcn_weights_id == kTensorNotAllocated) {

context->AddTensors(context, 1, &data->hwcn_weights_id);

}

++temporaries_count;

}

if (is_hybrid) {

// Allocate tensor to store the on-the-fly quantized inputs.

data->input_quantized_index = temporaries_count;

if (data->input_quantized_id == kTensorNotAllocated) {

TF_LITE_ENSURE_OK(

context, context->AddTensors(context, 1, &data->input_quantized_id));

}

++temporaries_count;

// Allocate tensor to store the quantization params computed during

// on-the-fly input quantization.

data->scaling_factors_index = temporaries_count;

if (data->scaling_factors_id == kTensorNotAllocated) {

TF_LITE_ENSURE_OK(

context, context->AddTensors(context, 1, &data->scaling_factors_id));

}

++temporaries_count;

// Allocate tensor to store the accumulators for the matrix multiply.

data->accum_scratch_index = temporaries_count;

if (data->accum_scratch_id == kTensorNotAllocated) {

TF_LITE_ENSURE_OK(

context, context->AddTensors(context, 1, &data->accum_scratch_id));

}

++temporaries_count;

if (is_per_channel) {

data->input_offset_index = temporaries_count;

if (data->input_offset_id == kTensorNotAllocated) {

TF_LITE_ENSURE_OK(

context, context->AddTensors(context, 1, &data->input_offset_id));

}

++temporaries_count;

data->row_sums_index = temporaries_count;

if (data->row_sums_id == kTensorNotAllocated) {

TF_LITE_ENSURE_OK(context,

context->AddTensors(context, 1, &data->row_sums_id));

}

++temporaries_count;

}

}

TfLiteIntArrayFree(node->temporaries);

node->temporaries = TfLiteIntArrayCreate(temporaries_count);

return kTfLiteOk;

TfLiteNode* node) {

auto* params = reinterpret_cast<TfLiteConvParams*>(node->builtin_data);

OpData* data = reinterpret_cast<OpData*>(node->user_data);

bool has_bias = node->inputs->size == 3;

// Check number of inputs/outputs

TF_LITE_ENSURE(context, has_bias || node->inputs->size == 2);

TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);

TfLiteTensor* output;

TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));

const TfLiteTensor* input;

TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));

const TfLiteTensor* filter;

TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &filter));

// Check dimensionality of input, filter

TF_LITE_ENSURE_EQ(context, input->dims->size, 4);

TF_LITE_ENSURE_EQ(context, filter->dims->size, 4);

// Check input channels matching filter

// Filter input channel can be a factor of channels of input (grouped conv)

// or equals (normal conv).

auto input_channel = input->dims->data[3];

auto filter_input_channel = filter->dims->data[3];

TF_LITE_ENSURE(context, filter_input_channel > 0);

TF_LITE_ENSURE_EQ(context, input_channel % filter_input_channel, 0);

data->groups = input_channel / filter_input_channel;

// Check types. (We assume that UINT8 refers to quantized tensors)

TfLiteType input_type = input->type;

TF_LITE_ENSURE(context,

input_type == kTfLiteFloat32 || input_type == kTfLiteUInt8 ||

input_type == kTfLiteInt8 || input_type == kTfLiteInt16);

TF_LITE_ENSURE_TYPES_EQ(context, output->type, input_type);

// Filter must have zero zero-points in per-channel quantization.

if (input_type == kTfLiteInt16 || input_type == kTfLiteInt8) {

TF_LITE_ENSURE_EQ(context, filter->quantization.type,

kTfLiteAffineQuantization);

const auto* affine_quantization =

reinterpret_cast<TfLiteAffineQuantization*>(

filter->quantization.params);

if (affine_quantization->zero_point) {

for (int i = 0; i < affine_quantization->zero_point->size; ++i) {

TF_LITE_ENSURE_EQ(context, affine_quantization->zero_point->data[i], 0);

}

}

}

// Validate stride values

TF_LITE_ENSURE(context, params->stride_height > 0);

TF_LITE_ENSURE(context, params->stride_width > 0);

// Validate dilation values

TF_LITE_ENSURE(context, params->dilation_height_factor > 0);

TF_LITE_ENSURE(context, params->dilation_width_factor > 0);

const TfLiteTensor* bias = nullptr;

// TODO(ahentz): At this point the optimized versions require 'bias'. We can

// either change that or document that convolution requires it.

TF_LITE_ENSURE(context, has_bias);

if (has_bias) {

TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 2, &bias));

if (input_type == kTfLiteUInt8 || input_type == kTfLiteInt8) {

TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt32);

TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);

} else if (input_type == kTfLiteInt16) {

TF_LITE_ENSURE(context, (bias->type == kTfLiteInt32) ||

(bias->type == kTfLiteInt64));

TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);

} else {

TF_LITE_ENSURE_TYPES_EQ(context, bias->type, input_type);

}

TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(filter, 0));

}

if (input_type == kTfLiteInt16) {

// Quantization should be symmetric.

TF_LITE_ENSURE_EQ(context, input->params.zero_point, 0);

TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);

// Check quantized_bias_type is either kTfLiteInt64 or kTfLiteInt32.

if (params->quantized_bias_type != kTfLiteFloat32) {

TF_LITE_ENSURE(context, params->quantized_bias_type == kTfLiteInt32 ||

params->quantized_bias_type == kTfLiteInt64);

TF_LITE_ENSURE(context, (bias == nullptr) ||

bias->type == params->quantized_bias_type);

data->quantized_bias_type = params->quantized_bias_type;

}

}

const bool is_hybrid =

(input->type == kTfLiteFloat32 &&

(filter->type == kTfLiteUInt8 || filter->type == kTfLiteInt8 ||

filter->type == kTfLiteInt4));

if (is_hybrid) {

int input_num_elements = 0;

TF_LITE_ENSURE_MSG(

context, CheckedNumElements(input, input_num_elements) == kTfLiteOk,

"%s", "Conv hybrid input has too many elements.");

}

if (filter->type == kTfLiteInt4) {

int filter_num_elements = 0;

TF_LITE_ENSURE_MSG(

context, CheckedNumElements(filter, filter_num_elements) == kTfLiteOk,

"%s", "Conv int4 filter has too many elements.");

}

if (filter->quantization.type == kTfLiteAffineQuantization) {

TF_LITE_ENSURE(context, filter->quantization.params);

TF_LITE_ENSURE(context, reinterpret_cast<TfLiteAffineQuantization*>(

filter->quantization.params)

->scale);

}

if (is_hybrid &&

(filter->type == kTfLiteInt8 || filter->type == kTfLiteInt4) &&

filter->quantization.type == kTfLiteAffineQuantization &&

reinterpret_cast<TfLiteAffineQuantization*>(filter->quantization.params)

->scale->size > 1) {

const auto* affine_quantization =

reinterpret_cast<TfLiteAffineQuantization*>(

filter->quantization.params);

const float scale = affine_quantization->scale->data[0];

for (int i = 1; i < affine_quantization->scale->size; i++) {

if (affine_quantization->scale->data[i] != scale) {

data->is_hybrid_per_channel = true;

break;

}

}

}

// The multi-threaded kernel supports neither dilation nor hybrid kernels, and

// is incompatible with mutable input filters that might change between evals.

data->supports_multithreaded_kernel =

(kernel_type == kMultithreadOptimized) &&

(context->recommended_num_threads != 1) && !is_hybrid &&

(params->dilation_width_factor == 1) &&

(params->dilation_height_factor == 1) &&

(filter->allocation_type != kTfLiteArenaRw) && !IsDynamicTensor(filter);

data->need_hwcn_weights =

input->type == kTfLiteFloat32 && data->supports_multithreaded_kernel;

int channels_in = filter->dims->data[3];

int channels_out = filter->dims->data[0];

int input_width = input->dims->data[2];

int input_height = input->dims->data[1];

int filter_width = filter->dims->data[2];

int filter_height = filter->dims->data[1];

int batches = input->dims->data[0];

// Matching GetWindowedOutputSize in TensorFlow.

auto padding = params->padding;

int out_width, out_height;

data->padding = ComputePaddingHeightWidth(

params->stride_height, params->stride_width,

params->dilation_height_factor, params->dilation_width_factor,

input_height, input_width, filter_height, filter_width, padding,

&out_height, &out_width);

size_t im2col_type_size;

TF_LITE_ENSURE_STATUS(GetSizeOfType(context, input->type, &im2col_type_size));

size_t im2col_elements = 0;

size_t im2col_bytes = 0;

const bool requires_im2col =

IsIm2ColRequired(input, params, filter, data, is_hybrid, kernel_type);

if (requires_im2col) {

TF_LITE_ENSURE_OK(

context,

CheckedShapeProduct(context,

{batches, out_height, out_width, channels_in,

filter_height, filter_width},

"Conv im2col size overflowed.", im2col_elements));

TF_LITE_ENSURE_MSG(

context,

MultiplyAndCheckOverflow(im2col_elements, im2col_type_size,

&im2col_bytes) == kTfLiteOk,

"Conv im2col byte size overflowed.");

}

TF_LITE_ENSURE_STATUS(AllocateTemporaryTensorsIfRequired(

context, node, is_hybrid, data->is_hybrid_per_channel, kernel_type,

/*im2col_bytes=*/im2col_bytes));

TF_LITE_ENSURE(context, has_bias);

// Note that full fixed-point inference requires that all tensors have their

// parameters set. This is usually done during quantized training or

// calibration.

if (input_type != kTfLiteFloat32) {

TF_LITE_ENSURE_EQ(context, filter->quantization.type,

kTfLiteAffineQuantization);

const auto* affine_quantization =

reinterpret_cast<TfLiteAffineQuantization*>(

filter->quantization.params);

TF_LITE_ENSURE(context, affine_quantization);

TF_LITE_ENSURE(context, affine_quantization->scale);

TF_LITE_ENSURE(context, (affine_quantization->scale->size == 1 ||

affine_quantization->scale->size == channels_out));

data->per_channel_output_multiplier.resize(channels_out);

data->per_channel_output_shift.resize(channels_out);

TF_LITE_ENSURE_STATUS(tflite::PopulateConvolutionQuantizationParams(

context, input, filter, bias, output, params->activation,

&data->output_multiplier, &data->output_shift,

&data->output_activation_min, &data->output_activation_max,

data->per_channel_output_multiplier.data(),

data->per_channel_output_shift.data(), channels_out));

}

std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> output_size(

TfLiteIntArrayCreate(4), TfLiteIntArrayFree);

output_size->data[0] = batches;

output_size->data[1] = out_height;

output_size->data[2] = out_width;

output_size->data[3] = channels_out;

auto output_status =

context->ResizeTensor(context, output, output_size.release());

if (output_status != kTfLiteOk) return output_status;

const RuntimeShape filter_shape = GetTensorShape(filter);

if (data->need_im2col) {

// Protect downstream kernels (Im2col, TransposeIm2col) that rely

// on 32-bit signed integers for their FlatSize() and pointer arithmetic.

if (im2col_elements > std::numeric_limits<int32_t>::max()) {

TF_LITE_KERNEL_LOG(

context,

"Conv im2col elements (%zu) exceed the 32-bit integer limit.",

im2col_elements);

return kTfLiteError;

}

node->temporaries->data[data->im2col_index] = data->im2col_id;

std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> im2col_size(

TfLiteIntArrayCreate(4), TfLiteIntArrayFree);

int im2col_depth = 0;

TF_LITE_ENSURE_MSG(

context,

filter_shape.CheckedSizeFromDimension(/*start=*/1, im2col_depth), "%s",

"Conv im2col depth overflowed.");

im2col_size->data[0] = batches;

im2col_size->data[1] = out_height;

im2col_size->data[2] = out_width;

im2col_size->data[3] = im2col_depth;

TfLiteTensor* im2col =

&context->tensors[node->temporaries->data[data->im2col_index]];

im2col->type = input->type;

if (is_hybrid) {

im2col->type = filter->type == kTfLiteInt4 ? kTfLiteInt8 : filter->type;

}

im2col->allocation_type = kTfLiteArenaRw;

auto im2col_status =

context->ResizeTensor(context, im2col, im2col_size.release());

if (im2col_status != kTfLiteOk) return im2col_status;

}

if (data->need_hwcn_weights) {

node->temporaries->data[data->hwcn_weights_index] = data->hwcn_weights_id;

// Because we're treating the filter weights as a matrix when we do the

// transpose, we allocate the buffer with a two-dimensional shape, where one

// dimension is the number of elements in each filter, and the second is the

// total number of filters.

int hwcn_filter_size = 0;

TF_LITE_ENSURE_MSG(

context,

filter_shape.CheckedSizeFromDimension(/*start=*/1, hwcn_filter_size),

"%s", "Conv HWCN weights size overflowed.");

int hwcn_weights_elements = 0;

TF_LITE_ENSURE_OK(context, CheckedShapeProductToInt(

context, {hwcn_filter_size, channels_out},

"Conv HWCN weights indexing overflowed.",

hwcn_weights_elements));

std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>

hwcn_weights_size(TfLiteIntArrayCreate(2), TfLiteIntArrayFree);

hwcn_weights_size->data[0] = hwcn_filter_size;

hwcn_weights_size->data[1] = channels_out;

TfLiteTensor* hwcn_weights =

&context->tensors[node->temporaries->data[data->hwcn_weights_index]];

hwcn_weights->type = input_type;

hwcn_weights->name = "Conv_hwcn_weights";

hwcn_weights->allocation_type = kTfLiteArenaRwPersistent;

auto hwcn_weights_status = context->ResizeTensor(

context, hwcn_weights, hwcn_weights_size.release());

if (hwcn_weights_status != kTfLiteOk) return hwcn_weights_status;

// TODO(petewarden): If Resize() is called when the size hasn't actually

// changed, this will do extra redundant work.

data->have_weights_been_transposed = false;

}

if (is_hybrid) {

node->temporaries->data[data->input_quantized_index] =

data->input_quantized_id;

TfLiteTensor* input_quantized;

TF_LITE_ENSURE_OK(

context, GetTemporarySafe(context, node, data->input_quantized_index,

&input_quantized));

input_quantized->type = kTfLiteInt8;

input_quantized->allocation_type = kTfLiteArenaRw;

if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) {

TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims);

TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,

input_quantized_size));

}

node->temporaries->data[data->scaling_factors_index] =

data->scaling_factors_id;

TfLiteTensor* scaling_factors;

TF_LITE_ENSURE_OK(

context, GetTemporarySafe(context, node, data->scaling_factors_index,

&scaling_factors));

scaling_factors->type = kTfLiteFloat32;

scaling_factors->allocation_type = kTfLiteArenaRw;

// Only one scale factor per batch is typically necessary. See optimized

// implementation for why we need to allocate for the height of the inputs

// flattened to 2D.

TF_LITE_ENSURE(context, channels_in != 0);

int flattened_input_height = 0;

TF_LITE_ENSURE_OK(

context,

CheckedShapeProductToInt(

context, {batches, input_height, input_width, data->groups},

"Conv hybrid scaling factors size overflowed.",

flattened_input_height));

int scaling_dims[1] = {flattened_input_height};

if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {

std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>

scaling_factors_size(TfLiteIntArrayCreate(1), TfLiteIntArrayFree);

scaling_factors_size->data[0] = flattened_input_height;

TF_LITE_ENSURE_OK(context,

context->ResizeTensor(context, scaling_factors,

scaling_factors_size.release()));

}

node->temporaries->data[data->accum_scratch_index] = data->accum_scratch_id;

TfLiteTensor* accum_scratch;

TF_LITE_ENSURE_OK(context,

GetTemporarySafe(context, node, data->accum_scratch_index,

&accum_scratch));

accum_scratch->type = kTfLiteInt32;

accum_scratch->allocation_type = kTfLiteArenaRw;

int scratch_width = 0;

TF_LITE_ENSURE_OK(

context, CheckedShapeProductToInt(

context, {batches, out_height, out_width},

"Conv hybrid scratch size overflowed.", scratch_width));

int accum_scratch_elements = 0;

TF_LITE_ENSURE_OK(context, CheckedShapeProductToInt(

context, {channels_out, scratch_width},

"Conv hybrid scratch indexing overflowed.",

accum_scratch_elements));

int accum_scratch_dims[2] = {channels_out, scratch_width};

if (!TfLiteIntArrayEqualsArray(accum_scratch->dims, 2,

accum_scratch_dims)) {

std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>

accum_scratch_size(TfLiteIntArrayCreate(2), TfLiteIntArrayFree);

accum_scratch_size->data[0] = channels_out;

accum_scratch_size->data[1] = scratch_width;

TF_LITE_ENSURE_OK(context,

context->ResizeTensor(context, accum_scratch,

accum_scratch_size.release()));

}

if (data->is_hybrid_per_channel) {

const auto* affine_quantization =

reinterpret_cast<TfLiteAffineQuantization*>(

filter->quantization.params);

TF_LITE_ENSURE(context, affine_quantization);

TF_LITE_ENSURE(context, affine_quantization->scale);

TF_LITE_ENSURE_EQ(

context, affine_quantization->scale->size,

filter->dims->data[affine_quantization->quantized_dimension]);

node->temporaries->data[data->input_offset_index] = data->input_offset_id;

TfLiteTensor* input_offsets;

TF_LITE_ENSURE_OK(

context, GetTemporarySafe(context, node, data->input_offset_index,

&input_offsets));

input_offsets->type = kTfLiteInt32;

input_offsets->allocation_type = kTfLiteArenaRw;

// See above comment for the need to allocate for height of inputs.

TF_LITE_ENSURE(context, channels_in != 0);

const int input_offset_dims[1] = {flattened_input_height};

if (!TfLiteIntArrayEqualsArray(input_offsets->dims, 1,

input_offset_dims)) {

std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>

input_offsets_size(TfLiteIntArrayCreate(1), TfLiteIntArrayFree);

input_offsets_size->data[0] = input_offset_dims[0];

TF_LITE_ENSURE_OK(context,

context->ResizeTensor(context, input_offsets,

input_offsets_size.release()));

}

node->temporaries->data[data->row_sums_index] = data->row_sums_id;

TfLiteTensor* row_sums;

TF_LITE_ENSURE_OK(

context,

GetTemporarySafe(context, node, data->row_sums_index, &row_sums));

row_sums->type = kTfLiteInt32;

row_sums->name = "Conv_row_sums";

row_sums->allocation_type = kTfLiteArenaRwPersistent;

// See above comment for the need to allocate for height of inputs.

const int row_sums_dims[1] = {channels_out};

if (!TfLiteIntArrayEqualsArray(row_sums->dims, 1, row_sums_dims)) {

std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>

row_sums_size(TfLiteIntArrayCreate(1), TfLiteIntArrayFree);

row_sums_size->data[0] = row_sums_dims[0];

TF_LITE_ENSURE_OK(

context,

context->ResizeTensor(context, row_sums, row_sums_size.release()));

}

}

}

return kTfLiteOk;

TfLiteConvParams* params, OpData* data,

const TfLiteTensor* input, const TfLiteTensor* filter,

const TfLiteTensor* bias, TfLiteTensor* im2col,

TfLiteTensor* output) {

auto input_offset = -input->params.zero_point;

auto filter_offset = -filter->params.zero_point;

auto output_offset = output->params.zero_point;

KernelType effective_kernel_type;

if ((kernel_type == kMultithreadOptimized ||

kernel_type == kCblasOptimized) &&

(params->dilation_width_factor != 1 ||

params->dilation_height_factor != 1)) {

// kMultithreadOptimized and kCblasOptimized do not support dilation.

// Therefore, fallback to optimized.

effective_kernel_type = kGenericOptimized;

} else {

effective_kernel_type = kernel_type;

}

// We have to fallback to reference execution path when im2col is needed but

// disabled because to-be-allocated temporary im2col tensor is too large.

// See b/178743262 for the detailed motivation.

if (data->im2col_oversized) {

effective_kernel_type = kReference;

}

// Grouped convolution is right now only supported on reference kernel.

if (data->groups != 1) {

effective_kernel_type = kReference;

}

const uint8_t* filter_data = nullptr;

std::unique_ptr<int8_t[]> unpacked_filter_data = nullptr;

if (filter->type == kTfLiteInt4) {

const size_t bytes_unpacked = filter->bytes * 2;

unpacked_filter_data = std::make_unique<int8_t[]>(bytes_unpacked);

tflite::tensor_utils::UnpackPackedIntToInt8(

GetTensorData<int8_t>(filter), GetTensorShape(filter).FlatSize(),

/*bit_width=*/4, unpacked_filter_data.get());

filter_data = reinterpret_cast<const uint8_t*>(unpacked_filter_data.get());

} else {

filter_data = GetTensorData<uint8_t>(filter);

}

ConvParams op_params;

op_params.padding_type = PaddingType::kSame;

op_params.padding_values.width = data->padding.width;

op_params.padding_values.height = data->padding.height;

op_params.dilation_width_factor = params->dilation_width_factor;

op_params.dilation_height_factor = params->dilation_height_factor;

op_params.stride_width = params->stride_width;

op_params.stride_height = params->stride_height;

op_params.input_offset = input_offset;

op_params.weights_offset = filter_offset;

op_params.output_offset = output_offset;

op_params.output_multiplier = data->output_multiplier;

op_params.output_shift = -data->output_shift;

op_params.quantized_activation_min = data->output_activation_min;

op_params.quantized_activation_max = data->output_activation_max;

switch (effective_kernel_type) {

case kReference: {

reference_ops::Conv(

op_params, GetTensorShape(input), GetTensorData<uint8_t>(input),

GetTensorShape(filter), filter_data, GetTensorShape(bias),

GetTensorData<int32_t>(bias), GetTensorShape(output),

GetTensorData<uint8_t>(output), GetTensorShape(im2col),

GetTensorData<uint8_t>(im2col),

/* cpu_backend_context = */ nullptr);

break;

}

case kGenericOptimized:

case kMultithreadOptimized:

case kCblasOptimized: {

// There is only one optimized implementation for Quantized Conv.

optimized_ops::Conv(

op_params, GetTensorShape(input), GetTensorData<uint8_t>(input),

GetTensorShape(filter), filter_data, GetTensorShape(bias),

GetTensorData<int32_t>(bias), GetTensorShape(output),

GetTensorData<uint8_t>(output), GetTensorShape(im2col),

GetTensorData<uint8_t>(im2col),

CpuBackendContext::GetFromContext(context));

break;

}

}

TfLiteConvParams* params, OpData* data,

const TfLiteTensor* input,

const TfLiteTensor* filter,

const TfLiteTensor* bias, TfLiteTensor* output,

TfLiteTensor* im2col) {

ConvParams op_params;

op_params.input_offset = -input->params.zero_point;

op_params.output_offset = output->params.zero_point;

op_params.stride_height = params->stride_height;

op_params.stride_width = params->stride_width;

op_params.dilation_height_factor = params->dilation_height_factor;

op_params.dilation_width_factor = params->dilation_width_factor;

op_params.padding_values.height = data->padding.height;

op_params.padding_values.width = data->padding.width;

op_params.quantized_activation_min = data->output_activation_min;

op_params.quantized_activation_max = data->output_activation_max;

KernelType effective_kernel_type = kernel_type;

// We have to fallback to reference execution path when im2col is needed but

// disabled because to-be-allocated temporary im2col tensor is too large.

// See b/178743262 for the detailed motivation.

if (data->im2col_oversized) {

effective_kernel_type = kReference;

}

// Grouped convolution is right now only supported on reference kernel.

if (data->groups != 1) {

effective_kernel_type = kReference;

}

const int8_t* filter_data = nullptr;

std::unique_ptr<int8_t[]> unpacked_filter_data = nullptr;

if (filter->type == kTfLiteInt4) {

const size_t bytes_unpacked = filter->bytes * 2;

unpacked_filter_data = std::make_unique<int8_t[]>(bytes_unpacked);

tflite::tensor_utils::UnpackPackedIntToInt8(

GetTensorData<int8_t>(filter), GetTensorShape(filter).FlatSize(),

/*bit_width=*/4, unpacked_filter_data.get());

filter_data = unpacked_filter_data.get();

} else {

filter_data = GetTensorData<int8_t>(filter);

}

switch (effective_kernel_type) {

case kReference: {

switch (filter->type) {

case kTfLiteInt4:

case kTfLiteInt8: {

reference_integer_ops::ConvPerChannel(

op_params, data->per_channel_output_multiplier.data(),

data->per_channel_output_shift.data(), GetTensorShape(input),

GetTensorData<int8>(input), GetTensorShape(filter), filter_data,

GetTensorShape(bias), GetTensorData<int32>(bias),

GetTensorShape(output), GetTensorData<int8>(output));

break;

}

default: {

TF_LITE_KERNEL_LOG(context,

"Weight type %s (%d) not supported for filter.",

TfLiteTypeGetName(filter->type), filter->type);

break;

}

}

break;

}

case kGenericOptimized:

case kMultithreadOptimized:

case kCblasOptimized:

switch (filter->type) {

case kTfLiteInt4:

case kTfLiteInt8: {

optimized_integer_ops::ConvPerChannel(

op_params, data->per_channel_output_multiplier.data(),

data->per_channel_output_shift.data(), GetTensorShape(input),

GetTensorData<int8>(input), GetTensorShape(filter), filter_data,

GetTensorShape(bias), GetTensorData<int32>(bias),

GetTensorShape(output), GetTensorData<int8>(output),

GetTensorShape(im2col), GetTensorData<int8>(im2col),

CpuBackendContext::GetFromContext(context));

break;

}

default: {

TF_LITE_KERNEL_LOG(context,

"Weight type %s (%d) not supported for filter.",

TfLiteTypeGetName(filter->type), filter->type);

break;

}

}

}

TfLiteConvParams* params, OpData* data,

const TfLiteTensor* input,

const TfLiteTensor* filter,

const TfLiteTensor* bias, TfLiteTensor* output,

TfLiteTensor* im2col) {

ConvParams op_params;

op_params.input_offset = -input->params.zero_point;

op_params.output_offset = output->params.zero_point;

op_params.stride_height = params->stride_height;

op_params.stride_width = params->stride_width;

op_params.dilation_height_factor = params->dilation_height_factor;

op_params.dilation_width_factor = params->dilation_width_factor;

op_params.padding_values.height = data->padding.height;

op_params.padding_values.width = data->padding.width;

op_params.quantized_activation_min = data->output_activation_min;

op_params.quantized_activation_max = data->output_activation_max;

KernelType effective_kernel_type = kernel_type;

// We have to fallback to reference execution path when im2col is needed but

// disabled because to-be-allocated temporary im2col tensor is too large.

// See b/178743262 for the detailed motivation.

if (data->im2col_oversized) {

effective_kernel_type = kReference;

}

// Grouped convolution is right now only supported on reference kernel.

if (data->groups != 1) {

effective_kernel_type = kReference;

}

// To prevent 32bit accum overflow for 16x8 quantization, it enables the

// optimized path only when zero_point is 0.

bool has_non_zero_point = input->params.zero_point ||

filter->params.zero_point ||

output->params.zero_point;

const int8_t* filter_data = nullptr;

std::unique_ptr<int8_t[]> unpacked_filter_data = nullptr;

if (filter->type == kTfLiteInt4) {

const size_t bytes_unpacked = filter->bytes * 2;

unpacked_filter_data = std::make_unique<int8_t[]>(bytes_unpacked);

tflite::tensor_utils::UnpackPackedIntToInt8(

GetTensorData<int8_t>(filter), GetTensorShape(filter).FlatSize(),

/*bit_width=*/4, unpacked_filter_data.get());

filter_data = unpacked_filter_data.get();

} else {

filter_data = GetTensorData<int8_t>(filter);

}

if (data->quantized_bias_type == kTfLiteInt32) {

if (effective_kernel_type == kReference || has_non_zero_point) {

reference_integer_ops::ConvPerChannel(

op_params, data->per_channel_output_multiplier.data(),

data->per_channel_output_shift.data(), GetTensorShape(input),

GetTensorData<int16>(input), GetTensorShape(filter), filter_data,

GetTensorShape(bias), GetTensorData<int32_t>(bias),

GetTensorShape(output), GetTensorData<int16>(output));