wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_COMMON_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_COMMON_H_
#include <cstdint>
namespace tflite {
namespace optimized_4bit {
// Since we need to convert int4 to int8 with shifts, it is faster if we
// can use unsigned int4, so just subtract zero_point_4bit from all values.
// Fold input * zero_point into quantization since we need to quantize
// each input and multiply by zero_point_4bit to convert back to signed int.
constexpr int zero_point_4bit = -7;
inline int8_t upper(int8_t value) { return value >> 4; }
inline int8_t lower(int8_t value) {
uint8_t sign_y = UINT8_C(256) - (value & UINT8_C(8));
return (value & UINT8_C(7)) | sign_y;
}
inline int8_t merge(int8_t upper, int8_t lower) {
const auto to_int4 = [](int8_t v) -> uint8_t {
int32_t x = v + 7;
return static_cast<uint8_t>(x);
};
return (to_int4(upper) << 4) | to_int4(lower);
}
} // namespace optimized_4bit
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_COMMON_H_
tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference.cc
+365
View File
@@ -0,0 +1,365 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include <stdint.h>
#include <algorithm>
#include <cstdlib>
#include <cstring>
#include <vector>
#include "tensorflow/lite/kernels/internal/cppmath.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference_impl.h"
namespace tflite {
namespace optimized_4bit {
void ReferencePackInner(const int8_t* src, uint8_t* box, int src_rows,
int src_cols, int outer_row, int outer_col,
int outer_rows, int outer_cols, int inner_rows,
int inner_cols) {
// Create a kernel-specific layout for a packed unit.
const int width = inner_rows;
const int depth = inner_cols;
const int real_depth = depth / 2;
const int real_src_cols = src_cols / 2;
// Determine row and column of source tensor.
const int row = outer_row * inner_rows;
const int col = outer_col * inner_cols;
// The source width (rows) and depth (columns).
int src_width = std::min(width, src_rows - row);
int src_depth = std::min(depth, src_cols - col);
int real_col = col / 2;
const int8_t* src_data = src + row * real_src_cols + real_col;
int real_src_depth = src_depth / 2;
// Src is [rows / src_rows, cols / src_depth, src_rows, src_cols]
// Reshape and pad to [outer_rows, outer_cols, width, depth]
// Interleave values [u1,u2,...,u_depth] to
// [u1,u_{depth/2+1},u2,u_{depth/2+2},.., u_{depth/2},u_depth]
// So that after shifting, we get [u1,u2...u_{depth/2}] and
// [u_{depth/2} + 1, ... u_{depth}].
for (int m = 0; m < src_width; ++m) {
int i = 0;
int k = 0;
int half_depth = depth / 2;
int half_half_depth = half_depth / 2;
for (; i < (real_src_depth & (~(half_depth - 1))); i += half_depth) {
for (int j = 0; j < half_half_depth; ++j) {
const int8_t v1 = (int8_t)src_data[i + j];
int8_t uv1 = upper(v1);
int8_t lv1 = lower(v1);
const int8_t v2 = (int8_t)src_data[i + j + half_half_depth];
int8_t uv2 = upper(v2);
int8_t lv2 = lower(v2);
box[k] = merge(lv1, lv2);
box[k + 1] = merge(uv1, uv2);
k += 2;
}
}
// Handle remaining values -- if greater than or equal to
// 16 values remaining, do the shuffle.
if (i < real_src_depth) {
const int remaining = half_half_depth < (real_src_depth - i)
? half_half_depth
: real_src_depth - i;
for (int j = 0; j < remaining; ++j) {
const int8_t v1 = (int8_t)src_data[i + j];
int8_t uv1 = upper(v1);
int8_t lv1 = lower(v1);
int8_t uv2 = 0;
int8_t lv2 = 0;
if ((i + j + half_half_depth) < real_src_depth) {
const int8_t v2 = (int8_t)src_data[i + j + half_half_depth];
uv2 = upper(v2);
lv2 = lower(v2);
}
box[k] = merge(lv1, lv2);
box[k + 1] = merge(uv1, uv2);
k += 2;
}
}
box += real_depth;
src_data += real_src_cols;
}
}
void ReferencePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth) {
size_t size = layout_rows * layout_cols / 2;
memset(dest, static_cast<uint8_t>(0x77), sizeof(uint8_t) * size);
int outer_cols = layout_cols / depth;
int outer_rows = layout_rows / width;
int inner_cols = depth;
int inner_rows = width;
for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
// Each outer row x outer col contains width x depth, copied
// from tensor at the cluster_index.
const int cluster_index = outer_row * outer_cols + outer_col;
const int real_depth = inner_cols / 2;
uint8_t* box = dest + cluster_index * real_depth * inner_rows;
ReferencePackInner(tensor, box, src_rows, src_cols, outer_row, outer_col,
outer_rows, outer_cols, inner_rows, inner_cols);
}
}
}
void ReferenceBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width,
int depth, int32_t* input_offsets) {
const int rows = n_batch;
const int cols = n_data;
// depth is always cols
const int layout_rows = (rows + (width - 1)) & ~(width - 1);
const int layout_cols = (cols + (depth - 1)) & ~(depth - 1);
const int size = layout_rows * layout_cols;
int8_t* data = quantized_data_ptr;
memset(data, 0, sizeof(int8_t) * size);
memset(input_offsets, 0, sizeof(int32_t) * layout_rows);
const float* tensor_data = float_data_ptr;
// basically, we need to make a new 4D matrix
// [rows / width, cols / depth, width, depth] in depth-first
const int outer_cols = layout_cols / depth;
const int outer_rows = layout_rows / width;
float* scaling_factors_ptr = scaling_factors;
for (int outer_row = 0; outer_row < outer_rows; outer_row++) {
std::vector<float> scale(width);
const int row = width * outer_row;
scaling_factors_ptr = scaling_factors + row;
for (int w = 0; w < width; ++w) {
if ((row + w) >= rows) {
continue;
}
const float* start = tensor_data + (row + w) * cols;
float scale_denom = 0;
for (int c = 0; c < cols; ++c) {
scale_denom = std::max(scale_denom, std::abs(*(start++)));
}
if (scale_denom == 0) {
scale_denom = 127.0;
}
scale[w] = 127.0 / scale_denom;
scaling_factors_ptr[w] = scale_denom / 127.0;
}
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const int col = depth * outer_col;
const int src_width = std::min(width, rows - row);
const int src_depth = std::min(depth, cols - col);
const int cluster_index = outer_row * outer_cols + outer_col;
int8_t* box = data + cluster_index * depth * width;
for (int w = 0; w < src_width; ++w) {
const float* float_data = tensor_data + (row + w) * cols + col;
for (int d = 0; d < src_depth; ++d) {
int8_t q = static_cast<int8_t>(TfLiteRound(float_data[d] * scale[w]));
box[w * depth + d] = q;
input_offsets[row + w] += q;
}
}
}
}
for (int r = 0; r < layout_rows; ++r) {
// Multiply the input by zero-point so that we don't have to calculate
// later.
input_offsets[r] = input_offsets[r] * zero_point_4bit;
}
}
void ReferenceAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
const float* filter_scales,
const float* bias_ptr,
float* output_ptr, int output_depth,
int batch_size) {
if (bias_ptr) {
for (int b = 0; b < batch_size; ++b) {
const float val = *input_offsets++ * *batch_scales++;
const float* filter_scales_ptr = filter_scales;
const float* bias_ptr_tmp = bias_ptr;
for (int i = 0; i < output_depth; i++) {
*output_ptr++ = (val * *filter_scales_ptr++) + *bias_ptr_tmp++;
}
}
return;
}
for (int b = 0; b < batch_size; ++b) {
const float val = *input_offsets++ * *batch_scales++;
const float* filter_scales_ptr = filter_scales;
for (int i = 0; i < output_depth; i++) {
*output_ptr++ = (val * *filter_scales_ptr++);
}
}
}
/* Unpack the accumulated scratch buffer by transposing and multiplying
* by input and filter scales.
* Before, dst contains integer accumulated values with layout:
* [rhs_layout_rows // rhs_width, lhs_layout_rows // lhs_width,
* rhs_width, lhs_width]
* Transpose and dequantize to [batch_size, num_units].
*/
template <int Depth, int Width>
void ReferenceUnpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols) {
// Width == 1 is when batch size == 1, the most frequent case.
// No need to iterate over outer rows.
if (Width == 1) {
const int outer_rows = dst_layout_rows / Width;
const int outer_cols = dst_layout_cols / Depth;
const int32_t* dst_ptr = dst;
int unit = 0;
for (int outer_col = 0; outer_col < outer_cols;
++outer_col, unit += Depth) {
float* tmp_output_ptr = output_ptr + unit;
int len = num_units - unit < Depth ? num_units - unit : Depth;
const float* scaling_factors_ptr = scaling_factors;
for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
const float scale = *scaling_factors_ptr;
const float* filter_scales_ptr = filter_scales + unit;
for (int i = 0; i < len; ++i) {
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
}
dst_ptr += (Depth - len);
scaling_factors_ptr += Width;
tmp_output_ptr += (num_units - len);
}
}
return;
}
const int outer_rows = dst_layout_rows / Width;
const int outer_cols = dst_layout_cols / Depth;
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const int unit = outer_col * Depth;
const int remaining_units = std::min(num_units - unit, Depth);
const int depth_offset = Depth - remaining_units;
const int width_offset = num_units - remaining_units;
int outer_row = 0;
for (; outer_row < outer_rows; ++outer_row) {
const int batch = outer_row * Width;
const int remaining_width = std::min(batch_size - batch, Width);
const int cluster_index = outer_col * outer_rows + outer_row;
const int32_t* dst_ptr = dst + cluster_index * Depth * Width;
float* tmp_output_ptr = output_ptr + batch * num_units + unit;
const float* scale = scaling_factors + batch;
int w = remaining_width;
for (; w > 0; --w, scale++) {
int d = remaining_units;
const float* filter_scales_ptr = filter_scales + unit;
for (; d > 0; --d) {
*tmp_output_ptr++ += *dst_ptr++ * (*scale) * (*filter_scales_ptr++);
}
dst_ptr += depth_offset;
tmp_output_ptr += width_offset;
}
}
}
}
template <int RowsLeft, int RowsRight, int Cols>
void ReferenceRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols) {
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
int32_t* elementPtr = dst;
const int outer_rows = (clamped_end_row + RowsLeft - 1) / RowsLeft;
const int outer_cols = (clamped_end_col + RowsRight - 1) / RowsRight;
const int depth = std::min(lhs_layout_cols / Cols, rhs_layout_cols / Cols);
for (int i = start_row; i < outer_rows; ++i) {
int left_index = i * RowsLeft * lhs_layout_cols / 2;
const uint8_t* lhs_val_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_val = lhs_val_data;
int right_index = j * RowsRight * rhs_layout_cols;
const int8_t* rhs_val = rhs + right_index;
int32_t accum[RowsLeft * RowsRight];
memset(accum, 0, sizeof(int32_t) * RowsLeft * RowsRight);
for (int k = 0; k < depth; ++k) {
uint8_t lhs_[RowsLeft][Cols];
for (int m = 0; m < RowsLeft; ++m) {
for (int n = 0; n < Cols / 2; ++n) {
uint8_t val = *(lhs_val++);
lhs_[m][n] = (val >> 4 & 15);
lhs_[m][n + (Cols / 2)] = (val & 15);
}
}
int8_t rhs_[RowsRight][Cols];
for (int m = 0; m < RowsRight; ++m) {
for (int n = 0; n < Cols; ++n) {
rhs_[m][n] = *(rhs_val++);
}
}
for (int r = 0; r < RowsRight; ++r) {
for (int l = 0; l < RowsLeft; ++l) {
for (int i = 0; i < Cols; ++i) {
accum[r * RowsLeft + l] += lhs_[l][i] * rhs_[r][i];
}
}
}
} // end depth
for (int r = 0; r < RowsRight; ++r) {
for (int l = 0; l < RowsLeft; ++l) {
int32_t q = accum[r * RowsLeft + l];
*(elementPtr++) = q;
}
}
}
}
}
template void ReferenceUnpack<4, 1>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales,
int dst_layout_rows, int dst_layout_cols);
template void ReferenceUnpack<4, 2>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales,
int dst_layout_rows, int dst_layout_cols);
template void ReferenceUnpack<4, 4>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales,
int dst_layout_rows, int dst_layout_cols);
template void ReferenceRunKernel<4, 1, 32>(
const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
template void ReferenceRunKernel<4, 2, 32>(
const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
template void ReferenceRunKernel<4, 4, 32>(
const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
} // namespace optimized_4bit
} // namespace tflite
tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference.h
+140
View File
@@ -0,0 +1,140 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_H_
#include <cstdint>
#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference_impl.h"
namespace tflite {
namespace optimized_4bit {
/* Returns the maximum number of rhs rows supported/compiled.
*
* In general 4 is the most we can do without running out of registers on
* aarch64. For x86/aarch64, 4x32 4bit = 64 bytes can be held in cache line
* size. This is required to set the packing layout for the rhs.
*
* For reference, return 1.
*/
inline int GetMaxSupportedRows() { return 1; }
/* Pack a 4bit inner_rows x inner_cols array from src.
* This is called as an inner function for Prepack.
*/
inline void PackInner(const int8_t* src, uint8_t* box, int src_rows,
int src_cols, int outer_row, int outer_col,
int outer_rows, int outer_cols, int inner_rows,
int inner_cols) {
ReferencePackInner(src, box, src_rows, src_cols, outer_row, outer_col,
outer_rows, outer_cols, inner_rows, inner_cols);
}
/* Prepack lhs matrix into dest.
* Transform tensor from (src_rows, src_cols) to
* (layout_rows / width, layout_cols / depth, width, depth) with possibly
* padding, and interleaving values along depth / 2 dimensions.
* dest should be aligned and allocated before prepack.
*/
inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth) {
ReferencePrepack(dest, tensor, layout_rows, layout_cols, src_rows, src_cols,
width, depth);
}
/* Quantize input floats to 8bit and calculate sum of each column.
* Data in float_data_ptr of shape (n_batch x n_data), is quantized and
* packed into (n_batch / width, n_data / depth, width, data) into
* quantized_data_ptr and input_offsets will contain the product of filter
* zero_point and input.
*/
inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width,
int depth, int32_t* input_offsets) {
ReferenceBatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
quantized_data_ptr, scaling_factors, width,
depth, input_offsets);
}
/* Write bias + input offset * filter_scale to output_ptr.
* output_ptr of size (batch_size, output_depth) will have
* output_ptr[output_depth * b + o] =
* bias_ptr[o] + input_offsets[b] * batch_scales[b] * filter_scale[o]
*/
inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr,
float* output_ptr, int output_depth,
int batch_size) {
ReferenceAssignBiasAndComputeOffsets(input_offsets, batch_scales,
filter_scales, bias_ptr, output_ptr,
output_depth, batch_size);
}
/* Add accumulated integer sums in dst to float output.
* output_ptr of size (batch_size, output_depth) will have
* output_ptr[b * output_depth + o] = \
* dst[b / dst_layout_rows, o / dst_layout_cols,
* b % dst_layout_rows, o % dst_layout_cols] * scaling_filters[b] *
* filter_scales[o]
*/
template <int Depth, int Width>
void Unpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols) {
ReferenceUnpack<Depth, Width>(output_ptr, dst, batch_size, num_units,
scaling_factors, filter_scales, dst_layout_rows,
dst_layout_cols);
}
/* Computes dst = (lchd,rchd->lr, lhs, rhs)
* Where l = lhs_layout_rows, r = rhs_layout_rows,
* c = rhs_layout_cols = lhs_layout_cols.
*/
template <int RowsLeft, int RowsRight, int Cols>
void RunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows, int dst_layout_cols) {
ReferenceRunKernel<RowsLeft, RowsRight, Cols>(
lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
rhs_layout_cols, dst_layout_rows, dst_layout_cols);
}
// Compute sum of lhs * rhs columnwise and write output to output_ptr.
inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int output_depth, int batch_size,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols,
float* output_ptr, const float* scaling_factors,
const float* filter_scales) {
ReferenceRunKernel<4, 1, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
rhs_layout_rows, rhs_layout_cols,
dst_layout_rows, dst_layout_cols);
ReferenceUnpack<4, 1>(output_ptr, dst, batch_size, output_depth,
scaling_factors, filter_scales, dst_layout_rows,
dst_layout_cols);
}
} // namespace optimized_4bit
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_H_
tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference_impl.h
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_IMPL_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_IMPL_H_
#include <stdint.h>
namespace tflite {
namespace optimized_4bit {
void ReferencePackInner(const int8_t* src, uint8_t* box, int src_rows,
int src_cols, int outer_row, int outer_col,
int outer_rows, int outer_cols, int inner_rows,
int inner_cols);
void ReferencePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth);
void ReferenceBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width,
int depth, int32_t* input_offsets);
void ReferenceAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
const float* filter_scales,
const float* bias_ptr,
float* output_ptr, int output_depth,
int batch_size);
template <int Depth, int Width>
extern void ReferenceUnpack(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template <int RowsLeft, int RowsRight, int Cols>
extern void ReferenceRunKernel(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
} // namespace optimized_4bit
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_IMPL_H_
tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected.cc
+609
View File
@@ -0,0 +1,609 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include <arm_neon.h>
#include <stdint.h>
#include <algorithm>
#include <cstring>
#include <vector>
#include "include/cpuinfo.h"
#include "tensorflow/lite/kernels/internal/cppmath.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
namespace tflite {
namespace optimized_4bit {
#ifndef __aarch64__
inline int16x8_t vqmovn_high_s32(int16x4_t a_s16x4, int32x4_t b_s32x4) {
return vcombine_s16(a_s16x4, vqmovn_s32(b_s32x4));
}
inline int8x16_t vqmovn_high_s16(int8x8_t a_s8x8, int16x8_t b_s16x8) {
return vcombine_s8(a_s8x8, vqmovn_s16(b_s16x8));
}
inline int32x4_t vpaddq_s32(int32x4_t a, int32x4_t b) {
int32x2_t a0 = vpadd_s32(vget_low_s32(a), vget_high_s32(a));
int32x2_t b0 = vpadd_s32(vget_low_s32(b), vget_high_s32(b));
return vcombine_s32(a0, b0);
}
inline float vmaxvq_f32(float32x4_t max_f32x4) {
float32x2_t max_f32x2 =
vmax_f32(vget_low_f32(max_f32x4), vget_high_f32(max_f32x4));
max_f32x2 = vpmax_f32(max_f32x2, max_f32x2);
return vget_lane_f32(max_f32x2, 0);
}
inline int32x4_t vcvtaq_s32_f32(float32x4_t a_f32x4) {
float32x4_t half = vdupq_n_f32(.5);
float32x4_t sign =
vcvtq_f32_u32(vshrq_n_u32(vreinterpretq_u32_f32(a_f32x4), 31));
float32x4_t add_half = vaddq_f32(a_f32x4, half);
float32x4_t round = vsubq_f32(add_half, sign);
return vcvtq_s32_f32(round);
}
void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr, float* output_ptr,
int output_depth, int batch_size) {
if (bias_ptr) {
for (int b = 0; b < batch_size; ++b) {
const float* filter_scales_ptr = filter_scales;
const float* tmp_bias_ptr = bias_ptr;
float val = input_offsets[b] * batch_scales[b];
int o = output_depth;
const float32x4_t v4_f32x4 = vdupq_n_f32(val);
for (; o >= 4; o -= 4) {
float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
filter_scales_ptr += 4;
float32x4_t v5_f32x4 = vld1q_f32(tmp_bias_ptr);
tmp_bias_ptr += 4;
v5_f32x4 = vmlaq_f32(v5_f32x4, v0_f32x4, v4_f32x4);
vst1q_f32(output_ptr, v5_f32x4);
output_ptr += 4;
}
for (; o > 0; --o) {
*output_ptr++ = val * (*filter_scales_ptr++) + (*tmp_bias_ptr++);
}
}
return;
}
for (int b = 0; b < batch_size; ++b) {
const float* filter_scales_ptr = filter_scales;
float val = input_offsets[b] * batch_scales[b];
int o = output_depth;
const float32x4_t v4_f32x4 = vdupq_n_f32(val);
for (; o >= 4; o -= 4) {
float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
filter_scales_ptr += 4;
float32x4_t v13_f32x4 = vmulq_f32(v0_f32x4, v4_f32x4);
vst1q_f32(output_ptr, v13_f32x4);
output_ptr += 4;
}
for (; o > 0; --o) {
*output_ptr++ = val * (*filter_scales_ptr++);
}
}
}
#else
void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr, float* output_ptr,
int output_depth, int batch_size) {
if (bias_ptr) {
for (int b = 0; b < batch_size; ++b) {
const float* filter_scales_ptr = filter_scales;
const float* tmp_bias_ptr = bias_ptr;
float val = input_offsets[b] * batch_scales[b];
int o = output_depth;
const float32x4_t v4_f32x4 = vdupq_n_f32(val);
for (; o >= 16; o -= 16) {
float32x4x4_t v0_to_v3_f32x4x4 = vld1q_f32_x4(filter_scales_ptr);
filter_scales_ptr += 16;
float32x4x4_t v5_to_v8_f32x4x4 = vld1q_f32_x4(tmp_bias_ptr);
tmp_bias_ptr += 16;
v5_to_v8_f32x4x4.val[0] = vfmaq_f32(v5_to_v8_f32x4x4.val[0],
v0_to_v3_f32x4x4.val[0], v4_f32x4);
v5_to_v8_f32x4x4.val[1] = vfmaq_f32(v5_to_v8_f32x4x4.val[1],
v0_to_v3_f32x4x4.val[1], v4_f32x4);
v5_to_v8_f32x4x4.val[2] = vfmaq_f32(v5_to_v8_f32x4x4.val[2],
v0_to_v3_f32x4x4.val[2], v4_f32x4);
v5_to_v8_f32x4x4.val[3] = vfmaq_f32(v5_to_v8_f32x4x4.val[3],
v0_to_v3_f32x4x4.val[3], v4_f32x4);
vst1q_f32_x4(output_ptr, v5_to_v8_f32x4x4);
output_ptr += 16;
}
if (o >= 8) {
float32x4x2_t v0_to_v1_f32x4x2 = vld1q_f32_x2(filter_scales_ptr);
filter_scales_ptr += 8;
float32x4x2_t v5_to_v6_f32x4x2 = vld1q_f32_x2(tmp_bias_ptr);
tmp_bias_ptr += 8;
v5_to_v6_f32x4x2.val[0] = vfmaq_f32(v5_to_v6_f32x4x2.val[0],
v0_to_v1_f32x4x2.val[0], v4_f32x4);
v5_to_v6_f32x4x2.val[1] = vfmaq_f32(v5_to_v6_f32x4x2.val[1],
v0_to_v1_f32x4x2.val[1], v4_f32x4);
vst1q_f32_x2(output_ptr, v5_to_v6_f32x4x2);
output_ptr += 8;
o -= 8;
}
if (o >= 4) {
float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
filter_scales_ptr += 4;
float32x4_t v5_f32x4 = vld1q_f32(tmp_bias_ptr);
tmp_bias_ptr += 4;
v5_f32x4 = vfmaq_f32(v5_f32x4, v0_f32x4, v4_f32x4);
vst1q_f32(output_ptr, v5_f32x4);
output_ptr += 4;
o -= 4;
}
for (; o > 0; --o) {
*output_ptr++ = val * (*filter_scales_ptr++) + (*tmp_bias_ptr++);
}
}
return;
}
for (int b = 0; b < batch_size; ++b) {
const float* filter_scales_ptr = filter_scales;
float val = input_offsets[b] * batch_scales[b];
int o = output_depth;
const float32x4_t v4_f32x4 = vdupq_n_f32(val);
for (; o >= 16; o -= 16) {
float32x4x4_t v0_to_v3_f32x4x4 = vld1q_f32_x4(filter_scales_ptr);
filter_scales_ptr += 16;
float32x4x4_t v13_to_v16_f32x4x4;
v13_to_v16_f32x4x4.val[0] = vmulq_f32(v0_to_v3_f32x4x4.val[0], v4_f32x4);
v13_to_v16_f32x4x4.val[1] = vmulq_f32(v0_to_v3_f32x4x4.val[1], v4_f32x4);
v13_to_v16_f32x4x4.val[2] = vmulq_f32(v0_to_v3_f32x4x4.val[2], v4_f32x4);
v13_to_v16_f32x4x4.val[3] = vmulq_f32(v0_to_v3_f32x4x4.val[3], v4_f32x4);
vst1q_f32_x4(output_ptr, v13_to_v16_f32x4x4);
output_ptr += 16;
}
if (o >= 8) {
float32x4x2_t v0_to_v1_f32x4x2 = vld1q_f32_x2(filter_scales_ptr);
filter_scales_ptr += 8;
float32x4x2_t v11_to_v12_f32x4x2;
v11_to_v12_f32x4x2.val[0] = vmulq_f32(v0_to_v1_f32x4x2.val[0], v4_f32x4);
v11_to_v12_f32x4x2.val[1] = vmulq_f32(v0_to_v1_f32x4x2.val[1], v4_f32x4);
vst1q_f32_x2(output_ptr, v11_to_v12_f32x4x2);
output_ptr += 8;
o -= 8;
}
if (o >= 4) {
float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
filter_scales_ptr += 4;
float32x4_t v13_f32x4 = vmulq_f32(v0_f32x4, v4_f32x4);
vst1q_f32(output_ptr, v13_f32x4);
output_ptr += 4;
o -= 4;
}
for (; o > 0; --o) {
*output_ptr++ = val * (*filter_scales_ptr++);
}
}
}
#endif
void NeonPackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
int outer_row, int outer_col, int outer_rows, int outer_cols,
int inner_rows, int inner_cols) {
// create a kernel-specific layout and store it into packed
const int width = inner_rows;
const int depth = inner_cols;
const int real_depth = depth / 2;
const int real_src_cols = src_cols / 2;
// which virtual row
const int row = outer_row * inner_rows;
const int col = outer_col * inner_cols;
int src_width = std::min(width, src_rows - row);
int src_depth = std::min(depth, src_cols - col);
int real_col = col / 2;
const int8_t* src_data = src + row * real_src_cols + real_col;
int real_src_depth = src_depth / 2;
const int8x16_t seven = vdupq_n_s8(7);
const int8x8_t seven8 = vdup_n_s8(7);
for (int m = 0; m < src_width; ++m) {
int i = 0;
int k = 0;
for (; i < (real_src_depth & (~15)); i += 16) {
int8x16_t values_16x8 = vld1q_s8(src_data + i);
int8x16_t uv1 = vshrq_n_s8(values_16x8, 4);
int8x16_t lv1 = vshlq_n_s8(values_16x8, 4);
uv1 = vaddq_s8(uv1, seven);
lv1 = vshrq_n_s8(lv1, 4);
lv1 = vaddq_s8(lv1, seven);
int8x8_t iuvl = vget_low_s8(uv1);
int8x8_t iuvh = vget_high_s8(uv1);
int8x8_t ilvl = vget_low_s8(lv1);
int8x8_t ilvh = vget_high_s8(lv1);
uint8x8_t uvl = vshl_n_u8(vreinterpret_u8_s8(iuvl), 4);
uint8x8_t lvl = vshl_n_u8(vreinterpret_u8_s8(ilvl), 4);
uint8x8_t uv = vorr_u8(uvl, vreinterpret_u8_s8(iuvh));
uint8x8_t lv = vorr_u8(lvl, vreinterpret_u8_s8(ilvh));
uint8x8x2_t zipped = vzip_u8(lv, uv);
uint8x16_t combined = vcombine_u8(zipped.val[0], zipped.val[1]);
vst1q_u8(box + k, combined);
k += 16;
}
// If exactly 16 values remaining, use fast path
if (real_src_depth == (real_src_depth & (~7))) {
for (; i < (real_src_depth & (~7)); i += 8) {
int8x8_t values_8x8 = vld1_s8(src_data + i);
int8x8_t uv1 = vshr_n_s8(values_8x8, 4);
int8x8_t lv1 = vshl_n_s8(values_8x8, 4);
uv1 = vadd_s8(uv1, seven8);
lv1 = vshr_n_s8(lv1, 4);
lv1 = vadd_s8(lv1, seven8);
uint8x8_t uvl = vshl_n_u8(vreinterpret_u8_s8(uv1), 4);
uint8x8_t lvl = vshl_n_u8(vreinterpret_u8_s8(lv1), 4);
uint8x8x2_t zipped = vzip_u8(lvl, uvl);
uint8x16_t combined = vcombine_u8(zipped.val[0], zipped.val[1]);
vst1q_u8(box + k, combined);
k += 16;
}
}
// Handle remaining values -- if greater than 16 values,
// shuffle.
if (i < real_src_depth) {
int remaining = 8;
remaining =
remaining < (real_src_depth - i) ? remaining : real_src_depth - i;
for (int j = 0; j < remaining; ++j) {
const int8_t v1 = (int8_t)src_data[i + j];
int8_t uv1 = upper(v1);
int8_t lv1 = lower(v1);
int8_t uv2 = 0;
int8_t lv2 = 0;
if ((i + j + 8) < real_src_depth) {
const int8_t v2 = (int8_t)src_data[i + j + 8];
uv2 = upper(v2);
lv2 = lower(v2);
}
box[k] = merge(lv1, lv2);
box[k + 1] = merge(uv1, uv2);
k += 2;
}
}
box += real_depth;
src_data += real_src_cols;
}
}
void NeonPrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth) {
// depth is always cols
size_t size = layout_rows * layout_cols / 2;
memset(dest, static_cast<uint8_t>(119), sizeof(uint8_t) * size);
// basically, we need to make a new 4D matrix
// [rows / width, cols / depth, width, depth] in depth-first
int outer_cols = layout_cols / depth;
int outer_rows = layout_rows / width;
int inner_cols = depth;
int inner_rows = width;
for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const int cluster_index = outer_row * outer_cols + outer_col;
const int real_depth = inner_cols / 2;
uint8_t* box = dest + cluster_index * real_depth * inner_rows;
NeonPackInner(tensor, box, src_rows, src_cols, outer_row, outer_col,
outer_rows, outer_cols, inner_rows, inner_cols);
}
}
}
void NeonBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width, int depth,
int32_t* input_offsets) {
const int rows = n_batch;
const int cols = n_data;
// depth is alpways cols
const int layout_rows = (rows + (width - 1)) & ~(width - 1);
const int layout_cols = (cols + (depth - 1)) & ~(depth - 1);
const int size = layout_rows * layout_cols;
int8_t* data = quantized_data_ptr;
memset(data, 0, sizeof(int8_t) * size);
memset(input_offsets, 0, sizeof(int32_t) * layout_rows);
const float* tensor_data = float_data_ptr;
// basically, we need to make a new 4D matrix
// [rows / width, cols / depth, width, depth] in depth-first
const int outer_cols = layout_cols / depth;
const int outer_rows = layout_rows / width;
float* scaling_factors_ptr = scaling_factors;
for (int outer_row = 0; outer_row < outer_rows; outer_row++) {
std::vector<float> scale(width);
const int row = width * outer_row;
scaling_factors_ptr = scaling_factors + row;
for (int w = 0; w < width; ++w) {
if ((row + w) >= rows) {
continue;
}
int c = 0;
const float* start = tensor_data + (row + w) * cols;
float32x4_t v1_f32x4 = vdupq_n_f32(0);
for (; c < (cols & ~3); c += 4) {
float32x4_t v0_f32x4 = vld1q_f32(start);
v0_f32x4 = vabsq_f32(v0_f32x4);
start += 4;
v1_f32x4 = vmaxq_f32(v0_f32x4, v1_f32x4);
}
float scale_denom = vmaxvq_f32(v1_f32x4);
for (; c < cols; ++c) {
scale_denom = std::max(scale_denom, std::abs(*(start++)));
}
if (scale_denom == 0) {
scale_denom = 127.0;
}
scale[w] = 127.0 / scale_denom;
scaling_factors_ptr[w] = scale_denom / 127.0;
}
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const int col = depth * outer_col;
const int src_width = std::min(width, rows - row);
const int src_depth = std::min(depth, cols - col);
const int cluster_index = outer_row * outer_cols + outer_col;
int8_t* box = data + cluster_index * depth * width;
__builtin_prefetch(box, 1, 3);
for (int w = 0; w < src_width; ++w) {
const float scale_w = scale[w];
const float* float_data = tensor_data + (row + w) * cols + col;
__builtin_prefetch(float_data, 0, 3);
int32_t* input_offsets_ptr = input_offsets + (row + w);
__builtin_prefetch(input_offsets_ptr, 1, 3);
int8_t* x0 = box + w * depth;
const float* x1 = float_data;
int16x8_t v12_s16x8 = vdupq_n_s16(0);
int32x4_t v13_s32x4 = vdupq_n_s32(0);
size_t run_depth = 0;
float32x4_t v0_f32x4 = vdupq_n_f32(scale_w);
#ifdef __aarch64__
for (; run_depth < (src_depth & ~15); run_depth += 16) {
const float32x4x4_t v1_f32x4x4 = vld1q_f32_x4(x1);
x1 += 16;
const float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4x4.val[0], v0_f32x4);
const float32x4_t v6_f32x4 = vmulq_f32(v1_f32x4x4.val[1], v0_f32x4);
const float32x4_t v7_f32x4 = vmulq_f32(v1_f32x4x4.val[2], v0_f32x4);
const float32x4_t v8_f32x4 = vmulq_f32(v1_f32x4x4.val[3], v0_f32x4);
const int32x4_t v5_s32x4 = vcvtaq_s32_f32(v5_f32x4);
const int32x4_t v6_s32x4 = vcvtaq_s32_f32(v6_f32x4);
const int32x4_t v7_s32x4 = vcvtaq_s32_f32(v7_f32x4);
const int32x4_t v8_s32x4 = vcvtaq_s32_f32(v8_f32x4);
const int16x4_t v9_low_s16x4 = vqmovn_s32(v5_s32x4);
const int16x8_t v9_s16x8 = vqmovn_high_s32(v9_low_s16x4, v6_s32x4);
const int16x4_t v10_low_s16x4 = vqmovn_s32(v7_s32x4);
const int16x8_t v10_s16x8 = vqmovn_high_s32(v10_low_s16x4, v8_s32x4);
const int8x8_t v11_low_s8x8 = vqmovn_s16(v9_s16x8);
const int8x16_t v11_s8x16 = vqmovn_high_s16(v11_low_s8x8, v10_s16x8);
v12_s16x8 = vaddq_s16(v12_s16x8, v9_s16x8);
v12_s16x8 = vaddq_s16(v12_s16x8, v10_s16x8);
vst1q_s8(x0, v11_s8x16);
x0 += 16;
}
for (; run_depth < (src_depth & ~7); run_depth += 8) {
const float32x4x2_t v1_f32x4x2 = vld1q_f32_x2(x1);
x1 += 8;
const float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4x2.val[0], v0_f32x4);
const float32x4_t v6_f32x4 = vmulq_f32(v1_f32x4x2.val[1], v0_f32x4);
const int32x4_t v5_s32x4 = vcvtaq_s32_f32(v5_f32x4);
const int32x4_t v6_s32x4 = vcvtaq_s32_f32(v6_f32x4);
const int16x4_t v9_low_s16x4 = vqmovn_s32(v5_s32x4);
const int16x8_t v9_s16x8 = vqmovn_high_s32(v9_low_s16x4, v6_s32x4);
const int8x8_t v11_low_s8x8 = vqmovn_s16(v9_s16x8);
vst1_s8(x0, v11_low_s8x8);
x0 += 8;
v12_s16x8 = vaddq_s16(v12_s16x8, v9_s16x8);
}
#else
for (; run_depth < (src_depth & ~7); run_depth += 8) {
const float32x4_t v1_f32x4_0 = vld1q_f32(x1);
x1 += 4;
const float32x4_t v1_f32x4_1 = vld1q_f32(x1);
x1 += 4;
const float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4_0, v0_f32x4);
const float32x4_t v6_f32x4 = vmulq_f32(v1_f32x4_1, v0_f32x4);
const int32x4_t v5_s32x4 = vcvtaq_s32_f32(v5_f32x4);
const int32x4_t v6_s32x4 = vcvtaq_s32_f32(v6_f32x4);
const int16x4_t v9_low_s16x4 = vqmovn_s32(v5_s32x4);
const int16x8_t v9_s16x8 = vqmovn_high_s32(v9_low_s16x4, v6_s32x4);
const int8x8_t v11_low_s8x8 = vqmovn_s16(v9_s16x8);
vst1_s8(x0, v11_low_s8x8);
x0 += 8;
v12_s16x8 = vaddq_s16(v12_s16x8, v9_s16x8);
}
#endif
int32_t row_sum = 0;
if (run_depth > 0) {
v13_s32x4 = vpadalq_s16(v13_s32x4, v12_s16x8);
v13_s32x4 = vpaddq_s32(v13_s32x4, v13_s32x4);
v13_s32x4 = vpaddq_s32(v13_s32x4, v13_s32x4);
row_sum += vgetq_lane_s32(v13_s32x4, 0);
}
for (; run_depth < src_depth; run_depth++) {
const float f = *x1++;
const int8_t q =
static_cast<int8_t>(::tflite::TfLiteRound(f * scale_w));
*x0++ = q;
row_sum += q;
}
input_offsets[row + w] += row_sum;
}
}
}
for (int r = 0; r < layout_rows; ++r) {
input_offsets[r] = input_offsets[r] * zero_point_4bit;
}
}
void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr, float* output_ptr,
int output_depth, int batch_size);
template <int Depth, int Width>
void NeonUnpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols) {
if (Width == 1) {
const int outer_rows = dst_layout_rows / Width;
const int outer_cols = dst_layout_cols / Depth;
const int32_t* dst_ptr = dst;
int unit = 0;
for (int outer_col = 0; outer_col < outer_cols;
++outer_col, unit += Depth) {
float* tmp_output_ptr = output_ptr + unit;
int len = num_units - unit < Depth ? num_units - unit : Depth;
int cond = len & ~3;
const float* scaling_factors_ptr = scaling_factors;
for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
const float scale = *scaling_factors_ptr;
const float* filter_scales_ptr = filter_scales + unit;
int i = 0;
for (; i < cond; i += 4) {
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
}
for (; i < len; ++i) {
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
}
dst_ptr += (Depth - len);
scaling_factors_ptr += Width;
tmp_output_ptr += (num_units - len);
}
}
return;
}
const int outer_rows = dst_layout_rows / Width;
const int outer_cols = dst_layout_cols / Depth;
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const int unit = outer_col * Depth;
const int remaining_units = std::min(num_units - unit, Depth);
const int depth_offset = Depth - remaining_units;
const int width_offset = num_units - remaining_units;
int outer_row = 0;
for (; outer_row < outer_rows; ++outer_row) {
const int batch = outer_row * Width;
const int remaining_width = std::min(batch_size - batch, Width);
const int cluster_index = outer_col * outer_rows + outer_row;
const int32_t* dst_ptr = dst + cluster_index * Depth * Width;
float* tmp_output_ptr = output_ptr + batch * num_units + unit;
const float* scale = scaling_factors + batch;
int w = remaining_width;
for (; w > 0; --w, scale++) {
int d = remaining_units;
const float* filter_scales_ptr = filter_scales + unit;
float32x4_t v3_f32x4 = vld1q_dup_f32(scale);
for (; d > 3; d -= 4) {
int32x4_t v0_s32x4 = vld1q_s32(dst_ptr);
dst_ptr += 4;
float32x4_t v1_f32x4 = vcvtq_f32_s32(v0_s32x4);
float32x4_t v2_f32x4 = vld1q_f32(filter_scales_ptr);
filter_scales_ptr += 4;
float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4, v3_f32x4);
float32x4_t v6_f32x4 = vmulq_f32(v5_f32x4, v2_f32x4);
float32x4_t v7_f32x4 = vld1q_f32(tmp_output_ptr);
float32x4_t v8_f32x4 = vaddq_f32(v6_f32x4, v7_f32x4);
vst1q_f32(tmp_output_ptr, v8_f32x4);
tmp_output_ptr += 4;
}
for (; d > 0; --d) {
*tmp_output_ptr++ += *dst_ptr++ * (*scale) * (*filter_scales_ptr++);
}
dst_ptr += depth_offset;
tmp_output_ptr += width_offset;
}
}
}
}
inline bool HasSDot() {
// CPUInfo already guards against double init
if (!cpuinfo_initialize()) {
// If we failed to init CPUInfo, assume ARM v8.2a-dotprod is not supported.
return false;
};
return cpuinfo_has_arm_neon_dot();
}
template <int RowsLeft, int RowsRight, int Cols>
void NeonRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols) {
if (HasSDot()) {
NeonRunKernelSDot<RowsLeft, RowsRight, Cols>(
lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
rhs_layout_cols, dst_layout_rows, dst_layout_cols);
return;
}
NeonRunKernelNoSDot<RowsLeft, RowsRight, Cols>(
lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
rhs_layout_cols, dst_layout_rows, dst_layout_cols);
}
template void NeonUnpack<4, 1>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template void NeonRunKernel<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
#ifdef __aarch64__
template void NeonUnpack<4, 2>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template void NeonRunKernel<4, 2, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
template void NeonUnpack<4, 4>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template void NeonRunKernel<4, 4, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
#endif
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_NEON)...
tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected.h
+130
View File
@@ -0,0 +1,130 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_H_
#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
namespace tflite {
namespace optimized_4bit {
// Maximum RowsRight compiled RunKernel implementations.
inline int GetMaxSupportedRows() {
#ifdef __aarch64__
return 4;
#else
return 1;
#endif
}
// Pack a 4bit inner_rows x inner_cols array from src.
inline void PackInner(const int8_t* src, uint8_t* box, int src_rows,
int src_cols, int outer_row, int outer_col,
int outer_rows, int outer_cols, int inner_rows,
int inner_cols) {
NeonPackInner(src, box, src_rows, src_cols, outer_row, outer_col, outer_rows,
outer_cols, inner_rows, inner_cols);
}
// Prepack lhs matrix into dest.
inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth) {
NeonPrepack(dest, tensor, layout_rows, layout_cols, src_rows, src_cols, width,
depth);
}
// Quantize input floats to 8bit and calculate sum of each column.
inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width,
int depth, int32_t* input_offsets) {
NeonBatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
quantized_data_ptr, scaling_factors, width, depth,
input_offsets);
}
// Write bias + input offset * filter_scale to output_ptr.
inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr,
float* output_ptr, int output_depth,
int batch_size) {
NeonAssignBiasAndComputeOffsets(input_offsets, batch_scales, filter_scales,
bias_ptr, output_ptr, output_depth,
batch_size);
}
// Add accumulated integer sums in dst to float output.
template <int Depth, int Width>
void Unpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols) {
NeonUnpack<Depth, Width>(output_ptr, dst, batch_size, num_units,
scaling_factors, filter_scales, dst_layout_rows,
dst_layout_cols);
}
// Compute sum of lhs * rhs columnwise.
template <int RowsLeft, int RowsRight, int Cols>
void RunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows, int dst_layout_cols) {
NeonRunKernel<RowsLeft, RowsRight, Cols>(
lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
rhs_layout_cols, dst_layout_rows, dst_layout_cols);
}
// Compute sum of lhs * rhs columnwise and write output to output_ptr.
inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int output_depth, int batch_size,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols,
float* output_ptr, const float* scaling_factors,
const float* filter_scales) {
#ifdef __aarch64__
if (rhs_width >= 4) {
NeonRunKernel<4, 4, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
dst_layout_cols);
NeonUnpack<4, 4>(output_ptr, dst, batch_size, output_depth, scaling_factors,
filter_scales, dst_layout_rows, dst_layout_cols);
return;
}
if (rhs_width >= 2) {
NeonRunKernel<4, 2, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
dst_layout_cols);
NeonUnpack<4, 2>(output_ptr, dst, batch_size, output_depth, scaling_factors,
filter_scales, dst_layout_rows, dst_layout_cols);
return;
}
#endif
NeonRunKernel<4, 1, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
dst_layout_cols);
NeonUnpack<4, 1>(output_ptr, dst, batch_size, output_depth, scaling_factors,
filter_scales, dst_layout_rows, dst_layout_cols);
}
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_NEON)...
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_H_
tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_aarch64_nosdot.cc
+607
View File
@@ -0,0 +1,607 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include <arm_neon.h>
#include <stdint.h>
#include <algorithm>
#include <vector>
#include "tensorflow/lite/kernels/internal/cppmath.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
namespace tflite {
namespace optimized_4bit {
template <int RowsLeft, int RowsRight, int Cols>
void NeonRunKernelNoSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols) {}
template <>
void NeonRunKernelNoSDot<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 1;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
movi v24.16b, #15
ld1 {v4.16b}, [%[lhs_ptr]], #16
movi v16.4s, #0
movi v17.4s, #0
ld1 {v5.16b}, [%[lhs_ptr]], #16
movi v18.4s, #0
movi v19.4s, #0
ld1 {v6.16b}, [%[lhs_ptr]], #16
and v8.16b, v4.16b, v24.16b
and v9.16b, v5.16b, v24.16b
ld1 {v7.16b}, [%[lhs_ptr]], #16
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
ld1 {v0.16b}, [%[rhs_ptr]], #16
and v10.16b, v6.16b, v24.16b
and v11.16b, v7.16b, v24.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.ls 1f /* skip loop */
0: /* loop start */
ld1 {v4.16b}, [%[lhs_ptr]], #16
smull v20.8h, v12.8b, v0.8b
smull v21.8h, v13.8b, v0.8b
smull v22.8h, v14.8b, v0.8b
ld1 {v5.16b}, [%[lhs_ptr]], #16
smull v23.8h, v15.8b, v0.8b
smlal v20.8h, v8.8b, v1.8b
smlal v21.8h, v9.8b, v1.8b
ld1 {v6.16b}, [%[lhs_ptr]], #16
smlal v22.8h, v10.8b, v1.8b
smlal v23.8h, v11.8b, v1.8b
smlal2 v20.8h, v12.16b, v0.16b
ld1 {v7.16b}, [%[lhs_ptr]], #16
smlal2 v21.8h, v13.16b, v0.16b
smlal2 v22.8h, v14.16b, v0.16b
smlal2 v23.8h, v15.16b, v0.16b
smlal2 v20.8h, v8.16b, v1.16b
smlal2 v21.8h, v9.16b, v1.16b
smlal2 v22.8h, v10.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
smlal2 v23.8h, v11.16b, v1.16b
sadalp v16.4s, v20.8h
sadalp v17.4s, v21.8h
sadalp v18.4s, v22.8h
sadalp v19.4s, v23.8h
ld1 {v1.16b}, [%[rhs_ptr]], #16
and v8.16b, v4.16b, v24.16b
and v9.16b, v5.16b, v24.16b
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
and v10.16b, v6.16b, v24.16b
and v11.16b, v7.16b, v24.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.hi 0b /* loop branch */
1: /* loop end */
smull v20.8h, v12.8b, v0.8b
smull v21.8h, v13.8b, v0.8b
smull v22.8h, v14.8b, v0.8b
smull v23.8h, v15.8b, v0.8b
smlal v20.8h, v8.8b, v1.8b
smlal v21.8h, v9.8b, v1.8b
smlal v22.8h, v10.8b, v1.8b
smlal v23.8h, v11.8b, v1.8b
smlal2 v20.8h, v12.16b, v0.16b
smlal2 v21.8h, v13.16b, v0.16b
smlal2 v22.8h, v14.16b, v0.16b
smlal2 v23.8h, v15.16b, v0.16b
smlal2 v20.8h, v8.16b, v1.16b
smlal2 v21.8h, v9.16b, v1.16b
smlal2 v22.8h, v10.16b, v1.16b
smlal2 v23.8h, v11.16b, v1.16b
sadalp v16.4s, v20.8h
sadalp v17.4s, v21.8h
sadalp v18.4s, v22.8h
sadalp v19.4s, v23.8h
addp v4.4s, v16.4s, v17.4s
addp v5.4s, v18.4s, v19.4s
addp v6.4s, v4.4s, v5.4s
st1 {v6.4s}, [%[dst]], #16
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "v0", "v1", "v4", "v5", "v6", "v7", "v8", "v9",
"v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18",
"v19", "v20", "v21", "v22", "v23", "v24");
}
}
}
template <>
void NeonRunKernelNoSDot<4, 2, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 2;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
ld1 {v4.16b}, [%[lhs_ptr]], #16
movi v31.16b, #15
movi v16.4s, #0
movi v17.4s, #0
ld1 {v5.16b}, [%[lhs_ptr]], #16
movi v18.4s, #0
movi v19.4s, #0
ld1 {v6.16b}, [%[lhs_ptr]], #16
and v8.16b, v4.16b, v31.16b
and v9.16b, v5.16b, v31.16b
ld1 {v7.16b}, [%[lhs_ptr]], #16
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
movi v24.4s, #0
ld1 {v0.16b}, [%[rhs_ptr]], #16
movi v25.4s, #0
movi v26.4s, #0
ld1 {v1.16b}, [%[rhs_ptr]], #16
movi v27.4s, #0
and v10.16b, v6.16b, v31.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
and v11.16b, v7.16b, v31.16b
ushr v14.16b, v6.16b, #4
ld1 {v3.16b}, [%[rhs_ptr]], #16
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.ls 1f /* skip loop */
0: /* loop start */
smull v20.8h, v12.8b, v0.8b
smull v21.8h, v13.8b, v0.8b
smull v22.8h, v14.8b, v0.8b
ld1 {v4.16b}, [%[lhs_ptr]], #16
smull v23.8h, v15.8b, v0.8b
smlal v20.8h, v8.8b, v1.8b
smlal v21.8h, v9.8b, v1.8b
ld1 {v5.16b}, [%[lhs_ptr]], #16
smlal v22.8h, v10.8b, v1.8b
smlal v23.8h, v11.8b, v1.8b
smlal2 v20.8h, v12.16b, v0.16b
ld1 {v6.16b}, [%[lhs_ptr]], #16
smlal2 v21.8h, v13.16b, v0.16b
smlal2 v22.8h, v14.16b, v0.16b
smlal2 v23.8h, v15.16b, v0.16b
ld1 {v7.16b}, [%[lhs_ptr]], #16
smlal2 v20.8h, v8.16b, v1.16b
smlal2 v21.8h, v9.16b, v1.16b
smlal2 v22.8h, v10.16b, v1.16b
smlal2 v23.8h, v11.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
sadalp v16.4s, v20.8h
sadalp v17.4s, v21.8h
sadalp v18.4s, v22.8h
sadalp v19.4s, v23.8h
ld1 {v1.16b}, [%[rhs_ptr]], #16
smull v28.8h, v12.8b, v2.8b
smull v29.8h, v13.8b, v2.8b
smull v30.8h, v14.8b, v2.8b
smull v20.8h, v15.8b, v2.8b
smlal v28.8h, v8.8b, v3.8b
smlal v29.8h, v9.8b, v3.8b
smlal v30.8h, v10.8b, v3.8b
smlal v20.8h, v11.8b, v3.8b
smlal2 v28.8h, v12.16b, v2.16b
smlal2 v29.8h, v13.16b, v2.16b
smlal2 v30.8h, v14.16b, v2.16b
smlal2 v20.8h, v15.16b, v2.16b
smlal2 v28.8h, v8.16b, v3.16b
smlal2 v29.8h, v9.16b, v3.16b
smlal2 v30.8h, v10.16b, v3.16b
smlal2 v20.8h, v11.16b, v3.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
sadalp v24.4s, v28.8h
sadalp v25.4s, v29.8h
sadalp v26.4s, v30.8h
sadalp v27.4s, v20.8h
ld1 {v3.16b}, [%[rhs_ptr]], #16
and v8.16b, v4.16b, v31.16b
and v9.16b, v5.16b, v31.16b
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
and v10.16b, v6.16b, v31.16b
and v11.16b, v7.16b, v31.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.hi 0b /* loop branch */
1: /* loop end */
smull v20.8h, v12.8b, v0.8b
smull v21.8h, v13.8b, v0.8b
smull v22.8h, v14.8b, v0.8b
smull v23.8h, v15.8b, v0.8b
smlal v20.8h, v8.8b, v1.8b
smlal v21.8h, v9.8b, v1.8b
smlal v22.8h, v10.8b, v1.8b
smlal v23.8h, v11.8b, v1.8b
smlal2 v20.8h, v12.16b, v0.16b
smlal2 v21.8h, v13.16b, v0.16b
smlal2 v22.8h, v14.16b, v0.16b
smlal2 v23.8h, v15.16b, v0.16b
smlal2 v20.8h, v8.16b, v1.16b
smlal2 v21.8h, v9.16b, v1.16b
smlal2 v22.8h, v10.16b, v1.16b
smlal2 v23.8h, v11.16b, v1.16b
smull v28.8h, v12.8b, v2.8b
smull v29.8h, v13.8b, v2.8b
smull v30.8h, v14.8b, v2.8b
smull v31.8h, v15.8b, v2.8b
smlal v28.8h, v8.8b, v3.8b
smlal v29.8h, v9.8b, v3.8b
smlal v30.8h, v10.8b, v3.8b
smlal v31.8h, v11.8b, v3.8b
smlal2 v28.8h, v12.16b, v2.16b
smlal2 v29.8h, v13.16b, v2.16b
smlal2 v30.8h, v14.16b, v2.16b
smlal2 v31.8h, v15.16b, v2.16b
smlal2 v28.8h, v8.16b, v3.16b
smlal2 v29.8h, v9.16b, v3.16b
smlal2 v30.8h, v10.16b, v3.16b
smlal2 v31.8h, v11.16b, v3.16b
sadalp v16.4s, v20.8h
sadalp v17.4s, v21.8h
sadalp v18.4s, v22.8h
sadalp v19.4s, v23.8h
sadalp v24.4s, v28.8h
sadalp v25.4s, v29.8h
sadalp v26.4s, v30.8h
sadalp v27.4s, v31.8h
addp v4.4s, v16.4s, v17.4s
addp v5.4s, v18.4s, v19.4s
addp v8.4s, v24.4s, v25.4s
addp v9.4s, v26.4s, v27.4s
addp v6.4s, v4.4s, v5.4s
addp v7.4s, v8.4s, v9.4s
st1 {v6.4s, v7.4s}, [%[dst]], #32
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
"v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
"v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26",
"v27", "v28", "v29", "v30", "v31");
}
}
}
template <>
void NeonRunKernelNoSDot<4, 4, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 4;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
movi v3.16b, #15
ld1 {v4.16b}, [%[lhs_ptr]], #16
movi v16.4s, #0
movi v17.4s, #0
movi v18.4s, #0
movi v19.4s, #0
ld1 {v5.16b}, [%[lhs_ptr]], #16
movi v20.4s, #0
movi v21.4s, #0
movi v22.4s, #0
ld1 {v6.16b}, [%[lhs_ptr]], #16
movi v23.4s, #0
movi v24.4s, #0
movi v25.4s, #0
ld1 {v7.16b}, [%[lhs_ptr]], #16
movi v26.4s, #0
movi v27.4s, #0
movi v28.4s, #0
movi v29.4s, #0
ld1 {v0.16b}, [%[rhs_ptr]], #16
movi v30.4s, #0
movi v31.4s, #0
ld1 {v1.16b}, [%[rhs_ptr]], #16
and v8.16b, v4.16b, v3.16b
and v9.16b, v5.16b, v3.16b
and v10.16b, v6.16b, v3.16b
and v11.16b, v7.16b, v3.16b
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.ls 1f /* skip loop */
0: /* loop start */
smull v4.8h, v12.8b, v0.8b
smull v5.8h, v13.8b, v0.8b
smull v6.8h, v14.8b, v0.8b
smull v7.8h, v15.8b, v0.8b
smlal2 v4.8h, v12.16b, v0.16b
smlal2 v5.8h, v13.16b, v0.16b
smlal2 v6.8h, v14.16b, v0.16b
smlal2 v7.8h, v15.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
smlal v4.8h, v8.8b, v1.8b
smlal v5.8h, v9.8b, v1.8b
smlal v6.8h, v10.8b, v1.8b
smlal v7.8h, v11.8b, v1.8b
smlal2 v4.8h, v8.16b, v1.16b
smlal2 v5.8h, v9.16b, v1.16b
smlal2 v6.8h, v10.16b, v1.16b
smlal2 v7.8h, v11.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
sadalp v16.4s, v4.8h
sadalp v17.4s, v5.8h
sadalp v18.4s, v6.8h
sadalp v19.4s, v7.8h
smull v4.8h, v12.8b, v2.8b
smull v5.8h, v13.8b, v2.8b
smull v6.8h, v14.8b, v2.8b
smull v7.8h, v15.8b, v2.8b
smlal2 v4.8h, v12.16b, v2.16b
smlal2 v5.8h, v13.16b, v2.16b
smlal2 v6.8h, v14.16b, v2.16b
smlal2 v7.8h, v15.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
smlal v4.8h, v8.8b, v0.8b
smlal v5.8h, v9.8b, v0.8b
smlal v6.8h, v10.8b, v0.8b
smlal v7.8h, v11.8b, v0.8b
smlal2 v4.8h, v8.16b, v0.16b
smlal2 v5.8h, v9.16b, v0.16b
smlal2 v6.8h, v10.16b, v0.16b
smlal2 v7.8h, v11.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
sadalp v20.4s, v4.8h
sadalp v21.4s, v5.8h
sadalp v22.4s, v6.8h
sadalp v23.4s, v7.8h
smull v4.8h, v12.8b, v1.8b
smull v5.8h, v13.8b, v1.8b
smull v6.8h, v14.8b, v1.8b
smull v7.8h, v15.8b, v1.8b
smlal2 v4.8h, v12.16b, v1.16b
smlal2 v5.8h, v13.16b, v1.16b
smlal2 v6.8h, v14.16b, v1.16b
smlal2 v7.8h, v15.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
smlal v4.8h, v8.8b, v2.8b
smlal v5.8h, v9.8b, v2.8b
smlal v6.8h, v10.8b, v2.8b
smlal v7.8h, v11.8b, v2.8b
smlal2 v4.8h, v8.16b, v2.16b
smlal2 v5.8h, v9.16b, v2.16b
smlal2 v6.8h, v10.16b, v2.16b
smlal2 v7.8h, v11.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
sadalp v24.4s, v4.8h
sadalp v25.4s, v5.8h
sadalp v26.4s, v6.8h
sadalp v27.4s, v7.8h
smull v4.8h, v12.8b, v0.8b
smull v5.8h, v13.8b, v0.8b
smull v6.8h, v14.8b, v0.8b
smull v7.8h, v15.8b, v0.8b
smlal2 v4.8h, v12.16b, v0.16b
smlal2 v5.8h, v13.16b, v0.16b
smlal2 v6.8h, v14.16b, v0.16b
smlal2 v7.8h, v15.16b, v0.16b
ld1 {v12.16b}, [%[lhs_ptr]], #16
smlal v4.8h, v8.8b, v1.8b
smlal v5.8h, v9.8b, v1.8b
smlal v6.8h, v10.8b, v1.8b
smlal v7.8h, v11.8b, v1.8b
ld1 {v13.16b}, [%[lhs_ptr]], #16
smlal2 v4.8h, v8.16b, v1.16b
smlal2 v5.8h, v9.16b, v1.16b
smlal2 v6.8h, v10.16b, v1.16b
smlal2 v7.8h, v11.16b, v1.16b
ld1 {v14.16b}, [%[lhs_ptr]], #16
sadalp v28.4s, v4.8h
sadalp v29.4s, v5.8h
sadalp v30.4s, v6.8h
sadalp v31.4s, v7.8h
ld1 {v15.16b}, [%[lhs_ptr]], #16
and v8.16b, v12.16b, v3.16b
and v9.16b, v13.16b, v3.16b
and v10.16b, v14.16b, v3.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
and v11.16b, v15.16b, v3.16b
ushr v12.16b, v12.16b, #4
ushr v13.16b, v13.16b, #4
ld1 {v1.16b}, [%[rhs_ptr]], #16
ushr v14.16b, v14.16b, #4
ushr v15.16b, v15.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.hi 0b /* loop branch */
1: /* loop end */
smull v4.8h, v12.8b, v0.8b
smull v5.8h, v13.8b, v0.8b
smull v6.8h, v14.8b, v0.8b
smull v7.8h, v15.8b, v0.8b
smlal2 v4.8h, v12.16b, v0.16b
smlal2 v5.8h, v13.16b, v0.16b
smlal2 v6.8h, v14.16b, v0.16b
smlal2 v7.8h, v15.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
smlal v4.8h, v8.8b, v1.8b
smlal v5.8h, v9.8b, v1.8b
smlal v6.8h, v10.8b, v1.8b
smlal v7.8h, v11.8b, v1.8b
smlal2 v4.8h, v8.16b, v1.16b
smlal2 v5.8h, v9.16b, v1.16b
smlal2 v6.8h, v10.16b, v1.16b
smlal2 v7.8h, v11.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
sadalp v16.4s, v4.8h
sadalp v17.4s, v5.8h
sadalp v18.4s, v6.8h
sadalp v19.4s, v7.8h
smull v4.8h, v12.8b, v2.8b
smull v5.8h, v13.8b, v2.8b
smull v6.8h, v14.8b, v2.8b
smull v7.8h, v15.8b, v2.8b
smlal2 v4.8h, v12.16b, v2.16b
smlal2 v5.8h, v13.16b, v2.16b
smlal2 v6.8h, v14.16b, v2.16b
smlal2 v7.8h, v15.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
smlal v4.8h, v8.8b, v0.8b
smlal v5.8h, v9.8b, v0.8b
smlal v6.8h, v10.8b, v0.8b
smlal v7.8h, v11.8b, v0.8b
smlal2 v4.8h, v8.16b, v0.16b
smlal2 v5.8h, v9.16b, v0.16b
smlal2 v6.8h, v10.16b, v0.16b
smlal2 v7.8h, v11.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
sadalp v20.4s, v4.8h
sadalp v21.4s, v5.8h
sadalp v22.4s, v6.8h
sadalp v23.4s, v7.8h
smull v4.8h, v12.8b, v1.8b
smull v5.8h, v13.8b, v1.8b
smull v6.8h, v14.8b, v1.8b
smull v7.8h, v15.8b, v1.8b
smlal2 v4.8h, v12.16b, v1.16b
smlal2 v5.8h, v13.16b, v1.16b
smlal2 v6.8h, v14.16b, v1.16b
smlal2 v7.8h, v15.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
smlal v4.8h, v8.8b, v2.8b
smlal v5.8h, v9.8b, v2.8b
smlal v6.8h, v10.8b, v2.8b
smlal v7.8h, v11.8b, v2.8b
smlal2 v4.8h, v8.16b, v2.16b
smlal2 v5.8h, v9.16b, v2.16b
smlal2 v6.8h, v10.16b, v2.16b
smlal2 v7.8h, v11.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
sadalp v24.4s, v4.8h
sadalp v25.4s, v5.8h
sadalp v26.4s, v6.8h
sadalp v27.4s, v7.8h
smull v4.8h, v12.8b, v0.8b
smull v5.8h, v13.8b, v0.8b
smull v6.8h, v14.8b, v0.8b
smull v7.8h, v15.8b, v0.8b
smlal2 v4.8h, v12.16b, v0.16b
smlal2 v5.8h, v13.16b, v0.16b
smlal2 v6.8h, v14.16b, v0.16b
smlal2 v7.8h, v15.16b, v0.16b
smlal v4.8h, v8.8b, v1.8b
smlal v5.8h, v9.8b, v1.8b
smlal v6.8h, v10.8b, v1.8b
smlal v7.8h, v11.8b, v1.8b
smlal2 v4.8h, v8.16b, v1.16b
smlal2 v5.8h, v9.16b, v1.16b
smlal2 v6.8h, v10.16b, v1.16b
smlal2 v7.8h, v11.16b, v1.16b
sadalp v28.4s, v4.8h
sadalp v29.4s, v5.8h
sadalp v30.4s, v6.8h
sadalp v31.4s, v7.8h
addp v14.4s, v16.4s, v17.4s
addp v15.4s, v18.4s, v19.4s
addp v12.4s, v20.4s, v21.4s
addp v13.4s, v22.4s, v23.4s
addp v10.4s, v24.4s, v25.4s
addp v11.4s, v26.4s, v27.4s
addp v8.4s, v28.4s, v29.4s
addp v9.4s, v30.4s, v31.4s
addp v4.4s, v14.4s, v15.4s
addp v5.4s, v12.4s, v13.4s
addp v6.4s, v10.4s, v11.4s
addp v7.4s, v8.4s, v9.4s
st1 {v4.4s, v5.4s, v6.4s, v7.4s}, [%[dst]], #64
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
"v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
"v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26",
"v27", "v28", "v29", "v30", "v31");
}
}
}
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
// defined(__ARM_NEON))
tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_aarch64_sdot.cc
+418
View File
@@ -0,0 +1,418 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include <stdint.h>
#include <algorithm>
#include <vector>
#include "tensorflow/lite/kernels/internal/cppmath.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
#define DOTPROD_ATTRIBUTE __attribute__((target("dotprod")))
namespace tflite {
namespace optimized_4bit {
template <int RowsLeft, int RowsRight, int Cols>
void NeonRunKernelSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
template <>
DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 1, 32>(
const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 1;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
movi v24.16b, #15
ld1 {v4.16b}, [%[lhs_ptr]], #16
movi v16.4s, #0
movi v17.4s, #0
ld1 {v5.16b}, [%[lhs_ptr]], #16
movi v18.4s, #0
movi v19.4s, #0
ld1 {v6.16b, v7.16b}, [%[lhs_ptr]], #32
and v8.16b, v4.16b, v24.16b
and v9.16b, v5.16b, v24.16b
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
ld1 {v0.16b, v1.16b}, [%[rhs_ptr]], #32
and v10.16b, v6.16b, v24.16b
and v11.16b, v7.16b, v24.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.ls 1f /* skip loop */
0: /* loop start */
ld1 {v4.16b, v5.16b, v6.16b, v7.16b}, [%[lhs_ptr]], #64
sdot v16.4s, v8.16b, v1.16b
sdot v17.4s, v9.16b, v1.16b
sdot v18.4s, v10.16b, v1.16b
sdot v19.4s, v11.16b, v1.16b
sdot v16.4s, v12.16b, v0.16b
sdot v17.4s, v13.16b, v0.16b
sdot v18.4s, v14.16b, v0.16b
sdot v19.4s, v15.16b, v0.16b
and v8.16b, v4.16b, v24.16b
and v9.16b, v5.16b, v24.16b
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
ld1 {v0.16b, v1.16b}, [%[rhs_ptr]], #32
and v10.16b, v6.16b, v24.16b
and v11.16b, v7.16b, v24.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.hi 0b /* loop branch */
1: /* loop end */
sdot v16.4s, v8.16b, v1.16b
sdot v17.4s, v9.16b, v1.16b
sdot v18.4s, v10.16b, v1.16b
sdot v19.4s, v11.16b, v1.16b
sdot v16.4s, v12.16b, v0.16b
sdot v17.4s, v13.16b, v0.16b
sdot v18.4s, v14.16b, v0.16b
sdot v19.4s, v15.16b, v0.16b
addp v4.4s, v16.4s, v17.4s
addp v5.4s, v18.4s, v19.4s
addp v6.4s, v4.4s, v5.4s
st1 {v6.4s}, [%[dst]], #16
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "v0", "v1", "v4", "v5", "v6", "v7", "v8", "v9",
"v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18",
"v19", "v24");
}
}
}
// Note: NeonRunKernelSDot<4, 2, 32> does not mutate registers v25-v31 in its
// inline assembly block, so they are intentionally omitted from the clobber
// list to avoid redundant register preservation.
template <>
DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 2, 32>(
const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 2;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
movi v24.16b, #15
ld1 {v4.16b}, [%[lhs_ptr]], #16
movi v16.4s, #0
movi v17.4s, #0
ld1 {v5.16b}, [%[lhs_ptr]], #16
movi v18.4s, #0
movi v19.4s, #0
ld1 {v6.16b, v7.16b}, [%[lhs_ptr]], #32
movi v20.4s, #0
movi v21.4s, #0
and v8.16b, v4.16b, v24.16b
and v9.16b, v5.16b, v24.16b
movi v22.4s, #0
movi v23.4s, #0
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
ld1 {v0.16b, v1.16b, v2.16b, v3.16b}, [%[rhs_ptr]], #64
and v10.16b, v6.16b, v24.16b
and v11.16b, v7.16b, v24.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.ls 1f /* skip loop */
0: /* loop start */
ld1 {v4.16b, v5.16b, v6.16b, v7.16b}, [%[lhs_ptr]], #64
sdot v16.4s, v12.16b, v0.16b
sdot v17.4s, v13.16b, v0.16b
sdot v18.4s, v14.16b, v0.16b
sdot v19.4s, v15.16b, v0.16b
sdot v16.4s, v8.16b, v1.16b
sdot v17.4s, v9.16b, v1.16b
sdot v18.4s, v10.16b, v1.16b
sdot v19.4s, v11.16b, v1.16b
sdot v20.4s, v12.16b, v2.16b
sdot v21.4s, v13.16b, v2.16b
sdot v22.4s, v14.16b, v2.16b
sdot v23.4s, v15.16b, v2.16b
sdot v20.4s, v8.16b, v3.16b
sdot v21.4s, v9.16b, v3.16b
sdot v22.4s, v10.16b, v3.16b
sdot v23.4s, v11.16b, v3.16b
and v8.16b, v4.16b, v24.16b
and v9.16b, v5.16b, v24.16b
ushr v12.16b, v4.16b, #4
ld1 {v0.16b, v1.16b, v2.16b, v3.16b}, [%[rhs_ptr]], #64
ushr v13.16b, v5.16b, #4
and v10.16b, v6.16b, v24.16b
and v11.16b, v7.16b, v24.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.hi 0b /* loop branch */
1: /* loop end */
sdot v16.4s, v12.16b, v0.16b
sdot v17.4s, v13.16b, v0.16b
sdot v18.4s, v14.16b, v0.16b
sdot v19.4s, v15.16b, v0.16b
sdot v16.4s, v8.16b, v1.16b
sdot v17.4s, v9.16b, v1.16b
sdot v18.4s, v10.16b, v1.16b
sdot v19.4s, v11.16b, v1.16b
sdot v20.4s, v12.16b, v2.16b
sdot v21.4s, v13.16b, v2.16b
sdot v22.4s, v14.16b, v2.16b
sdot v23.4s, v15.16b, v2.16b
sdot v20.4s, v8.16b, v3.16b
sdot v21.4s, v9.16b, v3.16b
sdot v22.4s, v10.16b, v3.16b
sdot v23.4s, v11.16b, v3.16b
addp v4.4s, v16.4s, v17.4s
addp v5.4s, v18.4s, v19.4s
addp v8.4s, v20.4s, v21.4s
addp v9.4s, v22.4s, v23.4s
addp v6.4s, v4.4s, v5.4s
addp v7.4s, v8.4s, v9.4s
st1 {v6.4s, v7.4s}, [%[dst]], #32
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
"v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
"v18", "v19", "v20", "v21", "v22", "v23", "v24");
}
}
}
template <>
DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 4, 32>(
const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 4;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
movi v3.16b, #15
ld1 {v4.16b}, [%[lhs_ptr]], #16
movi v16.4s, #0
movi v17.4s, #0
ld1 {v5.16b}, [%[lhs_ptr]], #16
movi v18.4s, #0
movi v19.4s, #0
ld1 {v6.16b, v7.16b}, [%[lhs_ptr]], #32
and v8.16b, v4.16b, v3.16b
and v9.16b, v5.16b, v3.16b
movi v20.4s, #0
movi v21.4s, #0
movi v22.4s, #0
movi v23.4s, #0
movi v24.4s, #0
movi v25.4s, #0
movi v26.4s, #0
movi v27.4s, #0
movi v28.4s, #0
movi v29.4s, #0
movi v30.4s, #0
movi v31.4s, #0
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
ld1 {v0.16b, v1.16b}, [%[rhs_ptr]], #32
and v10.16b, v6.16b, v3.16b
and v11.16b, v7.16b, v3.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.ls 1f /* skip loop */
0: /* loop start */
ld1 {v4.16b, v5.16b, v6.16b, v7.16b}, [%[lhs_ptr]], #64
sdot v16.4s, v12.16b, v0.16b
sdot v17.4s, v13.16b, v0.16b
sdot v18.4s, v14.16b, v0.16b
sdot v19.4s, v15.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
sdot v16.4s, v8.16b, v1.16b
sdot v17.4s, v9.16b, v1.16b
sdot v18.4s, v10.16b, v1.16b
sdot v19.4s, v11.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
sdot v20.4s, v12.16b, v2.16b
sdot v21.4s, v13.16b, v2.16b
sdot v22.4s, v14.16b, v2.16b
sdot v23.4s, v15.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
sdot v20.4s, v8.16b, v0.16b
sdot v21.4s, v9.16b, v0.16b
sdot v22.4s, v10.16b, v0.16b
sdot v23.4s, v11.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
sdot v24.4s, v12.16b, v1.16b
sdot v25.4s, v13.16b, v1.16b
sdot v26.4s, v14.16b, v1.16b
sdot v27.4s, v15.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
sdot v24.4s, v8.16b, v2.16b
sdot v25.4s, v9.16b, v2.16b
sdot v26.4s, v10.16b, v2.16b
sdot v27.4s, v11.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
sdot v28.4s, v12.16b, v0.16b
sdot v29.4s, v13.16b, v0.16b
sdot v30.4s, v14.16b, v0.16b
sdot v31.4s, v15.16b, v0.16b
sdot v28.4s, v8.16b, v1.16b
sdot v29.4s, v9.16b, v1.16b
sdot v30.4s, v10.16b, v1.16b
sdot v31.4s, v11.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
and v8.16b, v4.16b, v3.16b
and v9.16b, v5.16b, v3.16b
ushr v12.16b, v4.16b, #4
ushr v13.16b, v5.16b, #4
ld1 {v1.16b}, [%[rhs_ptr]], #16
and v10.16b, v6.16b, v3.16b
and v11.16b, v7.16b, v3.16b
ushr v14.16b, v6.16b, #4
ushr v15.16b, v7.16b, #4
subs %w[run_depth], %w[run_depth], #1
b.hi 0b /* loop branch */
1: /* loop end */
sdot v16.4s, v12.16b, v0.16b
sdot v17.4s, v13.16b, v0.16b
sdot v18.4s, v14.16b, v0.16b
sdot v19.4s, v15.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
sdot v16.4s, v8.16b, v1.16b
sdot v17.4s, v9.16b, v1.16b
sdot v18.4s, v10.16b, v1.16b
sdot v19.4s, v11.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
sdot v20.4s, v12.16b, v2.16b
sdot v21.4s, v13.16b, v2.16b
sdot v22.4s, v14.16b, v2.16b
sdot v23.4s, v15.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
sdot v20.4s, v8.16b, v0.16b
sdot v21.4s, v9.16b, v0.16b
sdot v22.4s, v10.16b, v0.16b
sdot v23.4s, v11.16b, v0.16b
ld1 {v2.16b}, [%[rhs_ptr]], #16
sdot v24.4s, v12.16b, v1.16b
sdot v25.4s, v13.16b, v1.16b
sdot v26.4s, v14.16b, v1.16b
sdot v27.4s, v15.16b, v1.16b
ld1 {v0.16b}, [%[rhs_ptr]], #16
sdot v24.4s, v8.16b, v2.16b
sdot v25.4s, v9.16b, v2.16b
sdot v26.4s, v10.16b, v2.16b
sdot v27.4s, v11.16b, v2.16b
ld1 {v1.16b}, [%[rhs_ptr]], #16
sdot v28.4s, v12.16b, v0.16b
sdot v29.4s, v13.16b, v0.16b
sdot v30.4s, v14.16b, v0.16b
sdot v31.4s, v15.16b, v0.16b
sdot v28.4s, v8.16b, v1.16b
sdot v29.4s, v9.16b, v1.16b
sdot v30.4s, v10.16b, v1.16b
sdot v31.4s, v11.16b, v1.16b
addp v14.4s, v16.4s, v17.4s
addp v15.4s, v18.4s, v19.4s
addp v12.4s, v20.4s, v21.4s
addp v13.4s, v22.4s, v23.4s
addp v10.4s, v24.4s, v25.4s
addp v11.4s, v26.4s, v27.4s
addp v8.4s, v28.4s, v29.4s
addp v9.4s, v30.4s, v31.4s
addp v4.4s, v14.4s, v15.4s
addp v5.4s, v12.4s, v13.4s
addp v6.4s, v10.4s, v11.4s
addp v7.4s, v8.4s, v9.4s
st1 {v4.4s, v5.4s, v6.4s, v7.4s}, [%[dst]], #64
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
"v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
"v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26",
"v27", "v28", "v29", "v30", "v31");
}
}
}
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
// defined(__ARM_NEON))
tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_arm32_nosdot.cc
+159
View File
@@ -0,0 +1,159 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include <stdint.h>
#include <algorithm>
#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
namespace tflite {
namespace optimized_4bit {
template <int RowsLeft, int RowsRight, int Cols>
void NeonRunKernelNoSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
template <>
void NeonRunKernelNoSDot<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 1;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
const int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
const int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
vmov.i8 q14, #15
vld1.8 {q4}, [%[lhs_ptr]]!
vmov.i32 q0, #0
vmov.i32 q1, #0
vld1.8 {q5}, [%[lhs_ptr]]!
vmov.i32 q2, #0
vmov.i32 q3, #0
vld1.8 {q6}, [%[lhs_ptr]]!
vand q8, q4, q14
vand q9, q5, q14
vld1.8 {q7}, [%[lhs_ptr]]!
vshr.u8 q4, q4, #4
vshr.u8 q5, q5, #4
vld1.8 {d24, d25}, [%[rhs_ptr]]!
vand q10, q6, q14
vand q11, q7, q14
vld1.8 {d26, d27}, [%[rhs_ptr]]!
vshr.u8 q6, q6, #4
vshr.u8 q7, q7, #4
subs %[run_depth], %[run_depth], #1
bls 1f /* skip loop */
0: /* loop start */
vmull.s8 q15, d8, d24
vmlal.s8 q15, d9, d25
vmlal.s8 q15, d16, d26
vmlal.s8 q15, d17, d27
vpadal.s16 q0, q15
vmull.s8 q15, d10, d24
vmlal.s8 q15, d11, d25
vmlal.s8 q15, d18, d26
vmlal.s8 q15, d19, d27
vpadal.s16 q1, q15
vld1.8 {q4, q5}, [%[lhs_ptr]]!
vmull.s8 q15, d12, d24
vmlal.s8 q15, d13, d25
vmlal.s8 q15, d20, d26
vmlal.s8 q15, d21, d27
vpadal.s16 q2, q15
vmull.s8 q15, d14, d24
vmlal.s8 q15, d15, d25
vmlal.s8 q15, d22, d26
vmlal.s8 q15, d23, d27
vpadal.s16 q3, q15
vld1.8 {q6, q7}, [%[lhs_ptr]]!
vld1.8 {d24, d25, d26, d27}, [%[rhs_ptr]]!
vand q8, q4, q14
vand q9, q5, q14
vand q10, q6, q14
vand q11, q7, q14
vshr.u8 q4, q4, #4
vshr.u8 q5, q5, #4
vshr.u8 q6, q6, #4
vshr.u8 q7, q7, #4
subs %[run_depth], %[run_depth], #1
bhi 0b /* loop branch */
1: /* loop end */
vmull.s8 q15, d8, d24
vmlal.s8 q15, d9, d25
vmlal.s8 q15, d16, d26
vmlal.s8 q15, d17, d27
vpadal.s16 q0, q15
vmull.s8 q15, d10, d24
vmlal.s8 q15, d11, d25
vmlal.s8 q15, d18, d26
vmlal.s8 q15, d19, d27
vpadal.s16 q1, q15
vmull.s8 q15, d12, d24
vmlal.s8 q15, d13, d25
vmlal.s8 q15, d20, d26
vmlal.s8 q15, d21, d27
vpadal.s16 q2, q15
vmull.s8 q15, d14, d24
vmlal.s8 q15, d15, d25
vmlal.s8 q15, d22, d26
vmlal.s8 q15, d23, d27
vpadal.s16 q3, q15
vpadd.i32 d0, d0, d1
vpadd.i32 d1, d2, d3
vpadd.i32 d2, d4, d5
vpadd.i32 d3, d6, d7
vpadd.i32 d4, d0, d1
vpadd.i32 d5, d2, d3
vst1.32 {d4, d5}, [%[dst]]!
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "d0", "d1", "d2", "d3", "d4", "d5", "d6", "d7",
"d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", "d16", "d17",
"d18", "d19", "d20", "d21", "d22", "d23", "d24", "d25", "d26",
"d27", "q0", "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9",
"q10", "q11", "q14", "q15");
}
}
}
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
// defined(__ARM_NEON))
tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_arm32_sdot.cc
+134
View File
@@ -0,0 +1,134 @@
/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include <stdint.h>
#include <algorithm>
#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
#define DOTPROD_ATTRIBUTE __attribute__((target("dotprod")))
namespace tflite {
namespace optimized_4bit {
template <int RowsLeft, int RowsRight, int Cols>
void NeonRunKernelSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
template <>
DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 1, 32>(
const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols) {
const int rows_left = 4;
const int rows_right = 1;
const int cols = 32;
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
for (int i = start_row; i < outer_rows; ++i) {
const int left_index = i * rows_left * lhs_layout_cols / 2;
const uint8_t* lhs_ptr_data = lhs + left_index;
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_ptr = lhs_ptr_data;
const int right_index = j * rows_right * rhs_layout_cols;
const int8_t* rhs_ptr = rhs + right_index;
int run_depth = depth;
asm(R"asm(
vmov.i8 q14, #15
vld1.8 {q4}, [%[lhs_ptr]]!
vmov.i32 q0, #0
vmov.i32 q1, #0
vld1.8 {q5}, [%[lhs_ptr]]!
vmov.i32 q2, #0
vmov.i32 q3, #0
vld1.8 {q6}, [%[lhs_ptr]]!
vand q8, q4, q14
vand q9, q5, q14
vld1.8 {q7}, [%[lhs_ptr]]!
vshr.u8 q4, q4, #4
vshr.u8 q5, q5, #4
vld1.8 {q12}, [%[rhs_ptr]]!
vand q10, q6, q14
vand q11, q7, q14
vld1.8 {q13}, [%[rhs_ptr]]!
vshr.u8 q6, q6, #4
vshr.u8 q7, q7, #4
subs %[run_depth], %[run_depth], #1
bls 1f /* skip loop */
0: /* loop start */
vsdot.s8 q0, q4, q12
vsdot.s8 q0, q8, q13
vsdot.s8 q1, q5, q12
vsdot.s8 q1, q9, q13
vld1.8 {q4, q5}, [%[lhs_ptr]]!
vsdot.s8 q2, q6, q12
vsdot.s8 q2, q10, q13
vsdot.s8 q3, q7, q12
vsdot.s8 q3, q11, q13
vld1.8 {q6, q7}, [%[lhs_ptr]]!
vld1.8 {q12, q13}, [%[rhs_ptr]]!
vand q8, q4, q14
vand q9, q5, q14
vand q10, q6, q14
vand q11, q7, q14
vshr.u8 q4, q4, #4
vshr.u8 q5, q5, #4
vshr.u8 q6, q6, #4
vshr.u8 q7, q7, #4
subs %[run_depth], %[run_depth], #1
bhi 0b /* loop branch */
1: /* loop end */
vsdot.s8 q0, q4, q12
vsdot.s8 q0, q8, q13
vsdot.s8 q1, q5, q12
vsdot.s8 q1, q9, q13
vsdot.s8 q2, q6, q12
vsdot.s8 q2, q10, q13
vsdot.s8 q3, q7, q12
vsdot.s8 q3, q11, q13
vpadd.i32 d0, d0, d1
vpadd.i32 d1, d2, d3
vpadd.i32 d2, d4, d5
vpadd.i32 d3, d6, d7
vpadd.i32 d4, d0, d1
vpadd.i32 d5, d2, d3
vst1.32 {d4, d5}, [%[dst]]!
)asm"
: [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
[run_depth] "+r"(run_depth)
:
: "cc", "memory", "d0", "d1", "d2", "d3", "d4", "d5", "d6", "d7",
"q0", "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9", "q10",
"q11", "q12", "q13", "q14");
}
}
}
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
// defined(__ARM_NEON))
tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h
+78
View File
@@ -0,0 +1,78 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_IMPL_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_IMPL_H_
#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include <stdint.h>
#if !defined(EIGEN_MAX_ALIGN_BYTES) && !defined(__aarch64__)
#define EIGEN_MAX_ALIGN_BYTES 32
#elif !defined(EIGEN_MAX_ALIGN_BYTES)
#define EIGEN_MAX_ALIGN_BYTES 64
#endif
namespace tflite {
namespace optimized_4bit {
void NeonPackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
int outer_row, int outer_col, int outer_rows, int outer_cols,
int inner_rows, int inner_cols);
void NeonPrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth);
void NeonBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width, int depth,
int32_t* input_offsets);
void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr, float* output_ptr,
int output_depth, int batch_size);
template <int Depth, int Width>
extern void NeonUnpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template <int RowsLeft, int RowsRight, int Cols>
extern void NeonRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
template <int RowsLeft, int RowsRight, int Cols>
extern void NeonRunKernelNoSDot(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
template <int RowsLeft, int RowsRight, int Cols>
extern void NeonRunKernelSDot(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_NEON)...
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_IMPL_H_
tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected.cc
+439
View File
@@ -0,0 +1,439 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
#include <stdint.h>
#include <stdlib.h>
// NOLINTBEGIN
#include <tmmintrin.h>
#include <algorithm>
#include <cstdlib>
#include <cstring>
#include <vector>
#include "tensorflow/lite/kernels/internal/cppmath.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h"
#include "tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected_impl.h"
namespace tflite {
namespace optimized_4bit {
#define is_aligned(ptr, bytes) ((((size_t)(ptr)) & (bytes - 1)) == 0)
void SsePackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
int outer_row, int outer_col, int outer_rows, int outer_cols,
int inner_rows, int inner_cols) {
const int width = inner_rows;
const int depth = inner_cols;
const int real_depth = depth / 2;
const int real_src_cols = src_cols / 2;
const int row = outer_row * inner_rows;
const int col = outer_col * inner_cols;
int src_width = std::min(width, src_rows - row);
int src_depth = std::min(depth, src_cols - col);
int real_col = col / 2;
const int8_t* src_data =
src + static_cast<size_t>(row) * real_src_cols + real_col;
int real_src_depth = src_depth / 2;
const __m128i bitmask_upper = _mm_set1_epi16(255U << 8);
const __m128i bitmask_lower = _mm_set1_epi16(255U);
const __m128i seven = _mm_set1_epi8(7);
for (int m = 0; m < src_width; ++m) {
int i = 0;
int k = 0;
for (; i < (real_src_depth & (~15)); i += 16) {
const __m128i values_128i = _mm_loadu_si128((__m128i*)(src_data + i));
// sign extend uv1
__m128i uv1 = _mm_srai_epi16(values_128i, 4);
uv1 = _mm_add_epi8(uv1, seven);
uv1 = _mm_and_si128(uv1, bitmask_upper);
__m128i uv2 = _mm_slli_epi16(values_128i, 8);
uv2 = _mm_srai_epi16(uv2, 12);
uv2 = _mm_add_epi8(uv2, seven);
uv2 = _mm_and_si128(uv2, bitmask_lower);
uv1 = _mm_or_si128(uv1, uv2);
__m128i lv1 = _mm_slli_epi16(values_128i, 4);
lv1 = _mm_srai_epi16(lv1, 4);
lv1 = _mm_add_epi8(lv1, seven);
lv1 = _mm_and_si128(lv1, bitmask_upper);
__m128i lv2 = _mm_slli_epi16(values_128i, 12);
lv2 = _mm_srai_epi16(lv2, 12);
lv2 = _mm_add_epi8(lv2, seven);
lv2 = _mm_and_si128(lv2, bitmask_lower);
lv1 = _mm_or_si128(lv1, lv2);
__m128i u = _mm_or_si128(_mm_slli_epi16(uv1, 4),
_mm_unpackhi_epi64(uv1, _mm_setzero_si128()));
__m128i l = _mm_or_si128(_mm_slli_epi16(lv1, 4),
_mm_unpackhi_epi64(lv1, _mm_setzero_si128()));
__m128i v = _mm_unpacklo_epi8(l, u);
_mm_store_si128((__m128i*)(box + k), v);
k += 16;
}
// Handle remaining values -- if greater than or equal to
// 16 values remaining, do the shuffle.
if (i < real_src_depth) {
int remaining = 8;
remaining =
remaining < (real_src_depth - i) ? remaining : real_src_depth - i;
for (int j = 0; j < remaining; j++) {
const int8_t v1 = (int8_t)src_data[i + j];
int8_t uv1 = upper(v1);
int8_t lv1 = lower(v1);
int8_t uv2 = 0;
int8_t lv2 = 0;
if ((i + j + 8) < real_src_depth) {
const int8_t v2 = (int8_t)src_data[i + j + 8];
uv2 = upper(v2);
lv2 = lower(v2);
}
box[k] = merge(lv1, lv2);
box[k + 1] = merge(uv1, uv2);
k += 2;
}
}
box += real_depth;
src_data += real_src_cols;
}
}
void SsePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth) {
size_t size = static_cast<size_t>(layout_rows) * layout_cols / 2;
memset(dest, static_cast<uint8_t>(119), sizeof(uint8_t) * size);
int outer_cols = layout_cols / depth;
int outer_rows = layout_rows / width;
int inner_cols = depth;
int inner_rows = width;
for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const size_t cluster_index =
static_cast<size_t>(outer_row) * outer_cols + outer_col;
const int real_depth = inner_cols / 2;
uint8_t* box = dest + cluster_index * real_depth * inner_rows;
SsePackInner(tensor, box, src_rows, src_cols, outer_row, outer_col,
outer_rows, outer_cols, inner_rows, inner_cols);
}
}
}
void SseBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width, int depth,
int32_t* input_offsets) {
const int rows = n_batch;
const int cols = n_data;
// depth is always cols
const int layout_rows = (rows + (width - 1)) & ~(width - 1);
const int layout_cols = (cols + (depth - 1)) & ~(depth - 1);
const size_t size = static_cast<size_t>(layout_rows) * layout_cols;
int8_t* data = quantized_data_ptr;
memset(data, 0, sizeof(int8_t) * size);
memset(input_offsets, 0, sizeof(int32_t) * layout_rows);
const float* tensor_data = float_data_ptr;
// basically, we need to make a new 4D matrix
// [rows / width, cols / depth, width, depth] in depth-first
const int outer_cols = layout_cols / depth;
const int outer_rows = layout_rows / width;
float* scaling_factors_ptr = scaling_factors;
for (int outer_row = 0; outer_row < outer_rows; outer_row++) {
std::vector<float> scale(width);
const int row = width * outer_row;
scaling_factors_ptr = scaling_factors + row;
for (int w = 0; w < width; ++w) {
if ((row + w) >= rows) {
continue;
}
const float* start = tensor_data + static_cast<size_t>(row + w) * cols;
int c = 0;
float scale_denom = 0;
for (; c < cols; ++c) {
scale_denom = std::max(scale_denom, std::abs(*(start++)));
}
if (scale_denom == 0) {
scale_denom = 127.0;
}
scale[w] = 127.0 / scale_denom;
scaling_factors_ptr[w] = scale_denom / 127.0;
}
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const int col = depth * outer_col;
const int src_width = std::min(width, rows - row);
const int src_depth = std::min(depth, cols - col);
const size_t cluster_index =
static_cast<size_t>(outer_row) * outer_cols + outer_col;
int8_t* box = data + cluster_index * depth * width;
for (int w = 0; w < src_width; ++w) {
const float* float_data =
tensor_data + static_cast<size_t>(row + w) * cols + col;
int d = 0;
for (; d < src_depth; ++d) {
int8_t q = static_cast<int8_t>(TfLiteRound(float_data[d] * scale[w]));
box[w * depth + d] = q;
input_offsets[row + w] += q;
}
}
}
}
for (int r = 0; r < layout_rows; ++r) {
input_offsets[r] = input_offsets[r] * zero_point_4bit;
}
}
void SseAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
const float* filter_scales,
const float* bias_ptr, float* output_ptr,
int output_depth, int batch_size) {
if (bias_ptr) {
for (int b = 0; b < batch_size; ++b) {
const float val = *input_offsets++ * *batch_scales++;
const float* filter_scales_ptr = filter_scales;
const float* bias_ptr_tmp = bias_ptr;
int i = 0;
for (; i < output_depth; i++) {
*output_ptr++ = (val * *filter_scales_ptr++) + *bias_ptr_tmp++;
}
}
return;
}
for (int b = 0; b < batch_size; ++b) {
const float val = *input_offsets++ * *batch_scales++;
const float* filter_scales_ptr = filter_scales;
int i = 0;
for (; i < output_depth; i++) {
*output_ptr++ = (val * *filter_scales_ptr++);
}
}
}
template <int Depth, int Width>
void SseUnpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols) {
if (Width == 1) {
const int outer_rows = dst_layout_rows / Width;
const int outer_cols = dst_layout_cols / Depth;
const int32_t* dst_ptr = dst;
int unit = 0;
for (int outer_col = 0; outer_col < outer_cols;
++outer_col, unit += Depth) {
float* tmp_output_ptr = output_ptr + unit;
int len = num_units - unit < Depth ? num_units - unit : Depth;
int cond = len & ~3;
const float* scaling_factors_ptr = scaling_factors;
for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
const float scale = *scaling_factors_ptr;
const float* filter_scales_ptr = filter_scales + unit;
int i = 0;
for (; i < cond; i += 4) {
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
}
for (; i < len; ++i) {
*(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
}
dst_ptr += (Depth - len);
scaling_factors_ptr += Width;
tmp_output_ptr += (num_units - len);
}
}
return;
}
const int outer_rows = dst_layout_rows / Width;
const int outer_cols = dst_layout_cols / Depth;
for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
const int unit = outer_col * Depth;
const int remaining_units = std::min(num_units - unit, Depth);
const int depth_offset = Depth - remaining_units;
const int width_offset = num_units - remaining_units;
int outer_row = 0;
for (; outer_row < outer_rows; ++outer_row) {
const int batch = outer_row * Width;
const int remaining_width = std::min(batch_size - batch, Width);
const size_t cluster_index =
static_cast<size_t>(outer_col) * outer_rows + outer_row;
const int32_t* dst_ptr = dst + cluster_index * Depth * Width;
float* tmp_output_ptr =
output_ptr + static_cast<size_t>(batch) * num_units + unit;
const float* scale = scaling_factors + batch;
int w = remaining_width;
for (; w > 0; --w, scale++) {
int d = remaining_units;
const float* filter_scales_ptr = filter_scales + unit;
for (; d > 0; --d) {
*tmp_output_ptr++ += *dst_ptr++ * (*scale) * (*filter_scales_ptr++);
}
dst_ptr += depth_offset;
tmp_output_ptr += width_offset;
}
}
}
}
inline __m128i DotProdInt8x4x4(__m128i acc_32x4, __m128i a_8x16,
__m128i b_8x16) {
b_8x16 = _mm_sign_epi8(b_8x16, a_8x16);
a_8x16 = _mm_abs_epi8(a_8x16);
__m128i sumprod_16x8 = _mm_maddubs_epi16(a_8x16, b_8x16);
return _mm_add_epi32(acc_32x4,
_mm_madd_epi16(sumprod_16x8, _mm_set1_epi16(1)));
}
inline __m128i ReduceInt32x4x4(__m128i a, __m128i b, __m128i c, __m128i d) {
// Assuming x = [x0, x1, x2, x3]
const __m128i a_b_lo_half = _mm_unpacklo_epi32(a, b); // [a0, b0, a1, b1]
const __m128i a_b_hi_half = _mm_unpackhi_epi32(a, b); // [a2, b2, a3, b3]
const __m128i a_plus_b =
_mm_add_epi32(a_b_lo_half, a_b_hi_half); // [a0+a2, b0+b2, a1+a3, b1+b3]
const __m128i c_d_lo_half = _mm_unpacklo_epi32(c, d); // [c0, d0, c1, d1]
const __m128i c_d_hi_half = _mm_unpackhi_epi32(c, d); // [c2, d2, c3, d3]
const __m128i c_plus_d =
_mm_add_epi32(c_d_lo_half, c_d_hi_half); // [c0+c2, d0+d2, c1+c3, d1+d3]
const __m128i all_evns =
_mm_unpacklo_epi64(a_plus_b, c_plus_d); // [a02, b02, c02, d02]
const __m128i all_odds =
_mm_unpackhi_epi64(a_plus_b, c_plus_d); // [a13, b13, c13, d13]
return _mm_add_epi32(all_evns, all_odds); // [a0123, b0123, c0123, d0123]
}
template <int RowsLeft, int RowsRight, int Cols>
void SseRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols) {
const int start_row = 0;
const int start_col = 0;
const int end_row = lhs_layout_rows;
const int end_col = rhs_layout_rows;
const int clamped_end_row = std::min(end_row, dst_layout_cols);
const int clamped_end_col = std::min(end_col, dst_layout_rows);
int32_t* elementPtr = dst;
const int outer_rows = (clamped_end_row + RowsLeft - 1) / RowsLeft;
const int outer_cols = (clamped_end_col + RowsRight - 1) / RowsRight;
const int depth = std::min(lhs_layout_cols / Cols, rhs_layout_cols / Cols);
const __m128i bitmask = _mm_set1_epi8(15);
const uintptr_t padding = 15;
std::vector<uint8_t> lhs_vec_data(
(static_cast<size_t>(RowsLeft) * lhs_layout_cols / 2) + padding);
uint8_t* lhs_vec = lhs_vec_data.data();
for (int i = start_row; i < outer_rows; ++i) {
size_t left_index = static_cast<size_t>(i) * RowsLeft * lhs_layout_cols / 2;
const uint8_t* lhs_val_data = lhs + left_index;
if (!is_aligned(lhs_val_data, 16)) {
size_t size = static_cast<size_t>(RowsLeft) * lhs_layout_cols / 2;
uintptr_t aligned =
(reinterpret_cast<uintptr_t>(lhs_vec) + padding) & ~(padding);
lhs_vec = reinterpret_cast<uint8_t*>(aligned);
memcpy(lhs_vec, lhs_val_data, size);
lhs_val_data = lhs_vec;
}
for (int j = start_col; j < outer_cols; ++j) {
const uint8_t* lhs_val = lhs_val_data;
size_t right_index = static_cast<size_t>(j) * RowsRight * rhs_layout_cols;
const int8_t* rhs_val = rhs + right_index;
__m128i accum[RowsRight * RowsLeft];
for (int m = 0; m < (RowsLeft * RowsRight); ++m) {
accum[m] = _mm_set1_epi8(0);
}
for (int k = 0; k < depth; ++k) {
__m128i lhs_row[RowsLeft];
for (int m = 0; m < RowsLeft; ++m) {
lhs_row[m] = _mm_load_si128((__m128i*)(lhs_val));
lhs_val += 16;
}
__m128i rhs[RowsRight][2];
for (int m = 0; m < RowsRight; ++m) {
for (int n = 0; n < 2; ++n) {
rhs[m][n] = _mm_loadu_si128((__m128i*)(rhs_val));
rhs_val += 16;
}
}
__m128i lhs_row_8[RowsLeft][2];
for (int m = 0; m < RowsLeft; ++m) {
lhs_row_8[m][0] = _mm_srli_epi16(lhs_row[m], 4);
lhs_row_8[m][1] = _mm_and_si128(lhs_row[m], bitmask);
}
for (int m = 0; m < RowsLeft; ++m) {
lhs_row_8[m][0] = _mm_and_si128(lhs_row_8[m][0], bitmask);
}
for (int i = 0; i < 2; ++i) {
for (int r = 0; r < RowsRight; ++r) {
for (int l = 0; l < RowsLeft; ++l) {
accum[r * RowsLeft + l] = DotProdInt8x4x4(
accum[r * RowsLeft + l], lhs_row_8[l][i], rhs[r][i]);
}
}
}
}
for (int r = 0; r < RowsRight; ++r) {
__m128i sum =
ReduceInt32x4x4(accum[r * RowsLeft], accum[r * RowsLeft + 1],
accum[r * RowsLeft + 2], accum[r * RowsLeft + 3]);
_mm_storeu_si128((__m128i*)elementPtr, sum);
elementPtr += 4;
}
}
}
}
// NOLINTEND
template void SseUnpack<4, 1>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template void SseUnpack<4, 2>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template void SseUnpack<4, 4>(float* output_ptr, const int32_t* dst,
int batch_size, int num_units,
const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template void SseRunKernel<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
template void SseRunKernel<4, 2, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
template void SseRunKernel<4, 4, 32>(const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int lhs_layout_rows,
int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows,
int dst_layout_cols);
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_SSE) && defined(__SSSE3__)
tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected.h
+125
View File
@@ -0,0 +1,125 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_H_
#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
#include <stdint.h>
#include "tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected_impl.h"
namespace tflite {
namespace optimized_4bit {
// Maximum RowsRight compiled RunKernel implementations.
inline int GetMaxSupportedRows() { return 4; }
// Pack a 4bit inner_rows x inner_cols array from src.
inline void PackInner(const int8_t* src, uint8_t* box, int src_rows,
int src_cols, int outer_row, int outer_col,
int outer_rows, int outer_cols, int inner_rows,
int inner_cols) {
SsePackInner(src, box, src_rows, src_cols, outer_row, outer_col, outer_rows,
outer_cols, inner_rows, inner_cols);
}
// Prepack lhs matrix into dest.
inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth) {
SsePrepack(dest, tensor, layout_rows, layout_cols, src_rows, src_cols, width,
depth);
}
// Quantize input floats to 8bit and calculate sum of each column.
inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width,
int depth, int32_t* input_offsets) {
SseBatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
quantized_data_ptr, scaling_factors, width, depth,
input_offsets);
}
// Write bias + input offset * filter_scale to output_ptr.
inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr,
float* output_ptr, int output_depth,
int batch_size) {
SseAssignBiasAndComputeOffsets(input_offsets, batch_scales, filter_scales,
bias_ptr, output_ptr, output_depth,
batch_size);
}
// Add accumulated integer sums in dst to float output.
template <int Depth, int Width>
void Unpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols) {
SseUnpack<Depth, Width>(output_ptr, dst, batch_size, num_units,
scaling_factors, filter_scales, dst_layout_rows,
dst_layout_cols);
}
// Compute sum of lhs * rhs columnwise.
template <int RowsLeft, int RowsRight, int Cols>
void RunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
int rhs_layout_cols, int dst_layout_rows, int dst_layout_cols) {
SseRunKernel<RowsLeft, RowsRight, Cols>(
lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
rhs_layout_cols, dst_layout_rows, dst_layout_cols);
}
// Compute sum of lhs * rhs columnwise and write output to output_ptr.
inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int output_depth, int batch_size,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols,
float* output_ptr, const float* scaling_factors,
const float* filter_scales) {
if (rhs_width >= 4) {
SseRunKernel<4, 4, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
dst_layout_cols);
SseUnpack<4, 4>(output_ptr, dst, batch_size, output_depth, scaling_factors,
filter_scales, dst_layout_rows, dst_layout_cols);
return;
}
if (rhs_width >= 2) {
SseRunKernel<4, 2, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
dst_layout_cols);
SseUnpack<4, 2>(output_ptr, dst, batch_size, output_depth, scaling_factors,
filter_scales, dst_layout_rows, dst_layout_cols);
return;
}
SseRunKernel<4, 1, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
dst_layout_cols);
SseUnpack<4, 1>(output_ptr, dst, batch_size, output_depth, scaling_factors,
filter_scales, dst_layout_rows, dst_layout_cols);
}
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_SSE) && defined(__SSSE3__)
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_H_
tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected_impl.h
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_IMPL_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_IMPL_H_
#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
#include <stdint.h>
#ifndef EIGEN_MAX_ALIGN_BYTES
#define EIGEN_MAX_ALIGN_BYTES 64
#endif
namespace tflite {
namespace optimized_4bit {
void SsePackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
int outer_row, int outer_col, int outer_rows, int outer_cols,
int inner_rows, int inner_cols);
void SsePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth);
void SseBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width, int depth,
int32_t* input_offsets);
void SseAssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
const float* filter_scales,
const float* bias_ptr, float* output_ptr,
int output_depth, int batch_size);
template <int Depth, int Width>
extern void SseUnpack(float* output_ptr, const int32_t* dst, int batch_size,
int num_units, const float* scaling_factors,
const float* filter_scales, int dst_layout_rows,
int dst_layout_cols);
template <int RowsLeft, int RowsRight, int Cols>
extern void SseRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols);
} // namespace optimized_4bit
} // namespace tflite
#endif // defined(FC_4BIT_SSE) && defined(__SSSE3__)
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_IMPL_H_
tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h
+160
View File
@@ -0,0 +1,160 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_AVX2_QUANTIZATION_UTILS_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_AVX2_QUANTIZATION_UTILS_H_
#ifdef __AVX2__
#include <immintrin.h>
#include <limits>
#include "tensorflow/lite/kernels/internal/compatibility.h"
namespace tflite {
namespace avx2_utils {
static inline void mm_storeu_si64(void *dst, __m128i v) {
#if (defined __clang__) || (defined _MSC_VER)
_mm_storeu_si64(dst, v);
#else
// GCC 9 lacks support for _mm_storeu_si64.
*static_cast<std::int64_t *>(dst) = _mm_extract_epi64(v, 0);
#endif
}
static inline __m256i mm256_blendv_epi32(const __m256i &a, const __m256i &b,
const __m256i &mask) {
__m256 result =
_mm256_blendv_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b),
_mm256_castsi256_ps(mask));
return _mm256_castps_si256(result);
}
static inline __m256i rounding_right_shift(const __m256i &value,
int32_t right_shift) {
TFLITE_DCHECK_GT(right_shift, 0);
const int32_t one_shift_exp_minus1 = 1 << (right_shift - 1);
__m256i nudge = _mm256_set1_epi32(one_shift_exp_minus1);
const __m256i r_plus_nudge = _mm256_add_epi32(value, nudge);
const __m256i shifted_sum =
_mm256_srav_epi32(r_plus_nudge, _mm256_set1_epi32(right_shift));
// Identify overflow in each lane and create mask.
const __m256i mask_num_plus_nudge_overflow = _mm256_cmpgt_epi32(
value, _mm256_set1_epi32(0x7fffffff - one_shift_exp_minus1));
// Fill results with either (value + nudge) >> exponent or
// std::numeric_limits<std::int32_t>::max() in the case of overflow.
return mm256_blendv_epi32(
shifted_sum, _mm256_set1_epi32(std::numeric_limits<std::int32_t>::max()),
mask_num_plus_nudge_overflow);
}
static inline __m256i rounding_right_shift(const __m256i &value,
const __m256i right_shift) {
const __m256i zeros = _mm256_setzero_si256();
const __m256i mask_rightshift_gtz = _mm256_cmpgt_epi32(right_shift, zeros);
const __m256i one_shift_exp_minus1 =
_mm256_sllv_epi32(_mm256_set1_epi32(1),
_mm256_sub_epi32(right_shift, _mm256_set1_epi32(1)));
__m256i nudge =
mm256_blendv_epi32(zeros, one_shift_exp_minus1, mask_rightshift_gtz);
const __m256i r_plus_nudge = _mm256_add_epi32(value, nudge);
const __m256i shifted_sum = _mm256_srav_epi32(r_plus_nudge, right_shift);
// Identify overflow in each lane and create mask.
const __m256i mask_num_plus_nudge_overflow = _mm256_cmpgt_epi32(
value, _mm256_sub_epi32(_mm256_set1_epi32(0x7fffffff), nudge));
// Fill results with either (value + nudge) >> exponent or
// std::numeric_limits<std::int32_t>::max() in the case of overflow.
return mm256_blendv_epi32(
shifted_sum, _mm256_set1_epi32(std::numeric_limits<std::int32_t>::max()),
mask_num_plus_nudge_overflow);
}
inline void CastInt32ToInt16AndStore(int16 *dst, const __m256i v) {
// As _mm256_cvtepi32_epi16 is not supported in AVX2, use the below repack.
// Select bytes 0, 1, 4, 5, 8, 9, 12, 13 within each lane, effectively
// truncating each 16-bit integer.
const __m256i repack_perm = _mm256_set1_epi64x(0x0d0c090805040100);
const __m256i shuffled_v = _mm256_shuffle_epi8(v, repack_perm);
mm_storeu_si64(dst, _mm256_extracti128_si256(shuffled_v, 0));
mm_storeu_si64(dst + 4, _mm256_extracti128_si256(shuffled_v, 1));
}
inline __m256i MultiplyByQuantizedMultiplier(const __m256i &value,
const int32_t multiplier,
const int32_t left_shift) {
const __m256i repack_perm = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
const __m256i shifted_value =
left_shift > 0 ? _mm256_sllv_epi32(value, _mm256_set1_epi32(left_shift))
: value;
__m256i scaled_v_low = _mm256_mul_epi32(
_mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 0)),
_mm256_set1_epi64x(multiplier));
__m256i scaled_v_high = _mm256_mul_epi32(
_mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 1)),
_mm256_set1_epi64x(multiplier));
scaled_v_low = _mm256_srlv_epi64(scaled_v_low, _mm256_set1_epi64x(31));
scaled_v_high = _mm256_srlv_epi64(scaled_v_high, _mm256_set1_epi64x(31));
// As _mm256_cvtepi64_epi32 is not supported in AVX2, use the below permute.
scaled_v_high = _mm256_slli_epi64(scaled_v_high, 32);
__m256i result = _mm256_blend_epi32(scaled_v_low, scaled_v_high, 0xaa);
result = _mm256_permutevar8x32_epi32(result, repack_perm);
if (left_shift >= 0) {
return result;
}
return rounding_right_shift(result, -left_shift);
}
inline __m256i MultiplyByQuantizedMultiplier(const __m256i &value,
const __m256i multiplier,
const __m256i left_shift) {
const __m256i zero_vector = _mm256_setzero_si256();
const __m256i positive_left_shift = _mm256_max_epi32(left_shift, zero_vector);
const __m256i positive_right_shift =
_mm256_max_epi32(_mm256_sub_epi32(zero_vector, left_shift), zero_vector);
const __m256i repack_perm = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
const __m256i shifted_value = _mm256_sllv_epi32(value, positive_left_shift);
const __m256i multiplier_low =
_mm256_cvtepi32_epi64(_mm256_extracti128_si256(multiplier, 0));
const __m256i multiplier_high =
_mm256_cvtepi32_epi64(_mm256_extracti128_si256(multiplier, 1));
__m256i scaled_v_low = _mm256_mul_epi32(
_mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 0)),
multiplier_low);
__m256i scaled_v_high = _mm256_mul_epi32(
_mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 1)),
multiplier_high);
scaled_v_low = _mm256_srlv_epi64(scaled_v_low, _mm256_set1_epi64x(31));
scaled_v_high = _mm256_srlv_epi64(scaled_v_high, _mm256_set1_epi64x(31));
// As _mm256_cvtepi64_epi32 is not supported in AVX2, use the below permute.
scaled_v_high = _mm256_slli_epi64(scaled_v_high, 32);
__m256i result = _mm256_blend_epi32(scaled_v_low, scaled_v_high, 0xaa);
result = _mm256_permutevar8x32_epi32(result, repack_perm);
return rounding_right_shift(result, positive_right_shift);
}
} // namespace avx2_utils
} // namespace tflite
#endif // __AVX2__
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_AVX2_QUANTIZATION_UTILS_H_
tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils_test.cc
+138
View File
@@ -0,0 +1,138 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
#include <gmock/gmock.h>
#include "tensorflow/lite/kernels/internal/common.h"
#ifdef __AVX2__
namespace tflite {
namespace avx2_utils {
namespace {
using ::testing::ElementsAreArray;
__m256i FillVectorWithInt32(const std::vector<int32_t>& src) {
return _mm256_set_epi32(src[7], src[6], src[5], src[4], src[3], src[2],
src[1], src[0]);
}
void CompareWithReferenceValue(std::vector<int32_t>& reference_values,
const __m256i& result) {
// As _mm256_extract_epi32 only supports const int, which should be known
// at the comile time, it puts down 8 comparison instead of for-loop.
EXPECT_NEAR(reference_values[0], _mm256_extract_epi32(result, 0), 1);
EXPECT_NEAR(reference_values[1], _mm256_extract_epi32(result, 1), 1);
EXPECT_NEAR(reference_values[2], _mm256_extract_epi32(result, 2), 1);
EXPECT_NEAR(reference_values[3], _mm256_extract_epi32(result, 3), 1);
EXPECT_NEAR(reference_values[4], _mm256_extract_epi32(result, 4), 1);
EXPECT_NEAR(reference_values[5], _mm256_extract_epi32(result, 5), 1);
EXPECT_NEAR(reference_values[6], _mm256_extract_epi32(result, 6), 1);
EXPECT_NEAR(reference_values[7], _mm256_extract_epi32(result, 7), 1);
}
TEST(CastInt32ToInt16AndStoreTest, CastInt32ToInt16AndStoreTest) {
const std::vector<int16_t> src = {1, 2, 3, 4, 5, 6, 7, 8};
int16_t dst[8];
const __m256i src_vector = _mm256_set_epi32(src[7], src[6], src[5], src[4],
src[3], src[2], src[1], src[0]);
CastInt32ToInt16AndStore(dst, src_vector);
EXPECT_THAT(src, ElementsAreArray(dst));
}
TEST(MultiplyByQuantizedMultiplierTest, PositiveLeftShiftTest) {
std::vector<int32_t> values = {100, 200, 300, 400, 500, 600, 700, 800};
const __m256i src_vector = FillVectorWithInt32(values);
const int32_t left_shift = 20;
const int32_t multiplier = 12345;
const __m256i result =
MultiplyByQuantizedMultiplier(src_vector, multiplier, left_shift);
// Get the reference values.
for (int i = 0; i < values.size(); i++) {
values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multiplier,
left_shift);
}
CompareWithReferenceValue(values, result);
}
TEST(MultiplyByQuantizedMultiplierTest, NegativeLeftShiftTest) {
std::vector<int32_t> values = {1000, 2000, 3000, 4000,
5000, 6000, 7000, 8000};
const __m256i src_vector = FillVectorWithInt32(values);
const int32_t left_shift = -3;
const int32_t multiplier = 1234567890;
const __m256i result =
MultiplyByQuantizedMultiplier(src_vector, multiplier, left_shift);
// Get the reference values.
for (int i = 0; i < values.size(); i++) {
values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multiplier,
left_shift);
}
CompareWithReferenceValue(values, result);
}
TEST(MultiplyByQuantizedMultiplierTest, VectorPositiveLeftShiftTest) {
std::vector<int32_t> values = {100, 200, 300, 400, 500, 600, 700, 800};
const std::vector<int32_t> left_shifts = {20, 19, 18, 17, 16, 15, 14, 13};
const std::vector<int32_t> multipliers = {10000, 20000, 30000, 40000,
50000, 60000, 70000, 80000};
const __m256i src_vector = FillVectorWithInt32(values);
const __m256i left_shifts_vector = FillVectorWithInt32(left_shifts);
const __m256i multipliers_vector = FillVectorWithInt32(multipliers);
const __m256i result = MultiplyByQuantizedMultiplier(
src_vector, multipliers_vector, left_shifts_vector);
// Get the reference values.
for (int i = 0; i < values.size(); i++) {
values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multipliers[i],
left_shifts[i]);
}
CompareWithReferenceValue(values, result);
}
TEST(MultiplyByQuantizedMultiplierTest, VectorNegativeLeftShiftTest) {
std::vector<int32_t> values = {1000, 2000, 3000, 4000,
5000, 6000, 7000, 8000};
const std::vector<int32_t> left_shifts = {-3, -4, -5, -6, -7, -8, -9, -10};
const std::vector<int32_t> multipliers = {1000000000, 1100000000, 1200000000,
1300000000, 1400000000, 1500000000,
1600000000, 1700000000};
const __m256i src_vector = FillVectorWithInt32(values);
const __m256i left_shifts_vector = FillVectorWithInt32(left_shifts);
const __m256i multipliers_vector = FillVectorWithInt32(multipliers);
const __m256i result = MultiplyByQuantizedMultiplier(
src_vector, multipliers_vector, left_shifts_vector);
// Get the reference values.
for (int i = 0; i < values.size(); i++) {
values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multipliers[i],
left_shifts[i]);
}
CompareWithReferenceValue(values, result);
}
} // namespace
} // namespace avx2_utils
} // namespace tflite
#endif // __AVX2__
tensorflow/lite/kernels/internal/optimized/cpu_check.cc
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#if defined __linux__ && defined __aarch64__
#include <sys/auxv.h>
#endif
namespace tflite {
namespace {
// The implementation of dotprod detection is copied from ruy's internal
// function DetectDotprod().
// At the moment it's only implemented on Linux ARM64. Consider syncing again
// with ruy in the future to share improvements.
#if defined __linux__ && defined __aarch64__
bool DetectDotprodByLinuxAuxvMethod() {
// This is the value of HWCAP_ASIMDDP in sufficiently recent Linux headers,
// however we need to support building against older headers for the time
// being.
const int kLocalHwcapAsimddp = 1 << 20;
return getauxval(AT_HWCAP) & kLocalHwcapAsimddp;
}
#endif
} // namespace
bool DetectArmNeonDotprod() {
#if defined __linux__ && defined __aarch64__
return DetectDotprodByLinuxAuxvMethod();
#else
return false;
#endif
}
} // namespace tflite
tensorflow/lite/kernels/internal/optimized/cpu_check.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_CPU_CHECK_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_CPU_CHECK_H_
// This include is superfluous. However, it's been here for a while, and a
// number of files have been relying on it to include neon_check.h for them.
// This should be removed, but with a global run of presubmits to catch
// any such issues. This requires running more than just TFLite presubmits.
#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
namespace tflite {
// On A64, returns true if the dotprod extension is present.
// On other architectures, returns false unconditionally.
bool DetectArmNeonDotprod();
struct CpuFlags {
bool neon_dotprod = false;
};
inline void GetCpuFlags(CpuFlags* cpu_flags) {
cpu_flags->neon_dotprod = DetectArmNeonDotprod();
}
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_CPU_CHECK_H_
tensorflow/lite/kernels/internal/optimized/depthwiseconv_3x3_filter_common.h
+590
View File
@@ -0,0 +1,590 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_3X3_FILTER_COMMON_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_3X3_FILTER_COMMON_H_
#include <algorithm>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_ops {
namespace depthwise_conv {
constexpr int kDepthwiseConvScratchWorkspaceSize = 10 * 10 * 64;
constexpr int kDepthwiseConvAdjustedBiasLimit = 64;
// In cases such as depth multiplication, we want to be able to load data from
// the workspace that is beyond the valid range. Macro-block sizes are adjusted
// to allow for this.
constexpr int kWorkspaceExtension = 16;
#ifdef USE_NEON
#ifndef __aarch64__
inline int8x16_t vqtbl4q_s8(int8x16x4_t a, int8x16_t b) {
const uint8x16_t mask = vtstq_s8(b, vdupq_n_s8(8));
// Delete bit 3 from the indices.
const int8x16_t high_bits = vshrq_n_s8(b, 4);
int8x16_t deleted_bit_3 = b;
deleted_bit_3 = vsliq_n_s8(deleted_bit_3, high_bits, 3);
int8x8x4_t repacked_data;
// Calculate for lower indices.
repacked_data.val[0] = vget_low_s8(a.val[0]);
repacked_data.val[1] = vget_low_s8(a.val[1]);
repacked_data.val[2] = vget_low_s8(a.val[2]);
repacked_data.val[3] = vget_low_s8(a.val[3]);
const int8x16_t output_for_lower =
vcombine_s8(vtbl4_s8(repacked_data, vget_low_s8(deleted_bit_3)),
vtbl4_s8(repacked_data, vget_high_s8(deleted_bit_3)));
// Calculate for high indices.
repacked_data.val[0] = vget_high_s8(a.val[0]);
repacked_data.val[1] = vget_high_s8(a.val[1]);
repacked_data.val[2] = vget_high_s8(a.val[2]);
repacked_data.val[3] = vget_high_s8(a.val[3]);
const int8x16_t output_for_higher =
vcombine_s8(vtbl4_s8(repacked_data, vget_low_s8(deleted_bit_3)),
vtbl4_s8(repacked_data, vget_high_s8(deleted_bit_3)));
// Merge.
int8x16_t output = vbslq_s8(mask, output_for_higher, output_for_lower);
return output;
}
#endif // !__aarch64__
// Convenience-compatibility functions.
// Compatibility: Intrinsics reflect a mixture of older and newer ARM
// instructions. This actually results in ZIP1 / ZIP2 asm instructions, but
// one intrinsic is provided. Also older instructions operated in place,
// and it seems more defensive to assume that some versions of intrinsics
// might reflect this
// Convenience: Callers in these kernels want both ZIP1 and ZIP2, and we do not
// want the calling code to get cluttered with unpacking int8x16x2_t.
inline void vzipq_s8_in_place(int8x16_t* a, int8x16_t* b) {
int8x16x2_t r8x16;
r8x16 = vzipq_s8(*a, *b);
*a = r8x16.val[0];
*b = r8x16.val[1];
}
inline void vzipq_s8x2_in_place(int8x16_t* a, int8x16_t* b) {
int16x8x2_t r16x8;
r16x8 = vzipq_s16(vreinterpretq_s16_s8(*a), vreinterpretq_s16_s8(*b));
*a = vreinterpretq_s8_s16(r16x8.val[0]);
*b = vreinterpretq_s8_s16(r16x8.val[1]);
}
// Similar rationale to the zip-in_place functions, but callers only actually
// need the TRN1 asm instruction result.
inline void vtrn1_s8x2_in_place(int8x16_t* a, int8x16_t* b) {
int16x8x2_t r16x8;
r16x8 = vtrnq_s16(vreinterpretq_s16_s8(*a), vreinterpretq_s16_s8(*b));
*a = vreinterpretq_s8_s16(r16x8.val[0]);
}
// Similar rationale to the zip-in_place functions, but callers only actually
// need the ZIP1 or ZIP2 asm instruction results.
inline int8x16_t vzip1q_s8(int8x16_t a, int8x16_t b) {
return vzipq_s8(a, b).val[0];
}
inline int8x16_t vzip2q_s8(int8x16_t a, int8x16_t b) {
return vzipq_s8(a, b).val[1];
}
inline void biregister_rotate_8(int8x16_t* left, int8x16_t* right) {
*left = vreinterpretq_s8_u32(vshrq_n_u32(vreinterpretq_u32_s8(*left), 8));
*left = vreinterpretq_s8_u32(vsliq_n_u32(vreinterpretq_u32_s8(*left),
vreinterpretq_u32_s8(*right), 24));
*right = vreinterpretq_s8_u32(vshrq_n_u32(vreinterpretq_u32_s8(*right), 8));
}
#ifndef __aarch64__
inline int32x4_t vpaddq_s32(int32x4_t a, int32x4_t b) {
int32x2_t a0 = vpadd_s32(vget_low_s32(a), vget_high_s32(a));
int32x2_t b0 = vpadd_s32(vget_low_s32(b), vget_high_s32(b));
return vcombine_s32(a0, b0);
}
#endif // !__aarch64__
#ifdef __ARM_FEATURE_DOTPROD
// The vdotq_lane_s32 takes int8x8t for the rhs parameter, whereas the actual
// instruction selects from between 4 32-bit (4x8-bit packed) sub-registers, an
// unusual interpretation of "lane".
inline int32x4_t vdotq_four_lane_s32(int32x4_t acc, int8x16_t lhs,
int8x16_t rhs, const int lane) {
switch (lane) {
case 0:
return vdotq_lane_s32(acc, lhs, vget_low_s8(rhs), 0);
case 1:
return vdotq_lane_s32(acc, lhs, vget_low_s8(rhs), 1);
case 2:
return vdotq_lane_s32(acc, lhs, vget_high_s8(rhs), 0);
case 3:
default:
return vdotq_lane_s32(acc, lhs, vget_high_s8(rhs), 1);
}
}
#else
inline int32x4_t vdotq_s32(int32x4_t acc, int8x16_t lhs, int8x16_t rhs) {
int32x4_t sum0 = vpaddlq_s16(vmull_s8(vget_low_s8(lhs), vget_low_s8(rhs)));
int32x4_t sum1 = vpaddlq_s16(vmull_s8(vget_high_s8(lhs), vget_high_s8(rhs)));
int32x4_t sum = vpaddq_s32(sum0, sum1);
return vaddq_s32(acc, sum);
}
inline int32x4_t vdotq_four_lane_s32(int32x4_t acc, int8x16_t lhs,
int8x16_t rhs, int lane) {
int8x8_t lane_rhs;
if (lane == 0) {
lane_rhs = vreinterpret_s8_s32(
vdup_lane_s32(vreinterpret_s32_s8(vget_low_s8(rhs)), 0));
} else if (lane == 1) {
lane_rhs = vreinterpret_s8_s32(
vdup_lane_s32(vreinterpret_s32_s8(vget_low_s8(rhs)), 1));
} else if (lane == 2) {
lane_rhs = vreinterpret_s8_s32(
vdup_lane_s32(vreinterpret_s32_s8(vget_high_s8(rhs)), 0));
} else {
lane_rhs = vreinterpret_s8_s32(
vdup_lane_s32(vreinterpret_s32_s8(vget_high_s8(rhs)), 1));
}
int32x4_t sum0 = vpaddlq_s16(vmull_s8(vget_low_s8(lhs), lane_rhs));
int32x4_t sum1 = vpaddlq_s16(vmull_s8(vget_high_s8(lhs), lane_rhs));
int32x4_t sum = vpaddq_s32(sum0, sum1);
return vaddq_s32(acc, sum);
}
#endif // !__ARM_FEATURE_DOTPROD
#endif // ARM NEON
// This structure is typically used for reducing the magnitude of outputs, and
// the historical name reflects that.
template <DepthwiseConvOutputRounding output_rounding>
struct DivideByPOT {};
template <>
struct DivideByPOT<DepthwiseConvOutputRounding::kAwayFromZero> {
template <typename IntegerType>
static inline IntegerType Run(IntegerType x, int exponent) {
return RoundingDivideByPOT(x, exponent);
}
// Mult versions use the exponents directly, rather than negated.
template <typename IntegerType>
static inline IntegerType RunMult(IntegerType x, int exponent) {
return RoundingDivideByPOT(x, -exponent);
}
};
#ifdef USE_NEON
template <>
struct DivideByPOT<DepthwiseConvOutputRounding::kUpward> {
template <typename IntegerType>
static inline IntegerType Run(IntegerType x, int exponent) {
return vqrshlq_s32(x, vdupq_n_s32(static_cast<int32_t>(-exponent)));
}
template <typename IntegerType>
static inline IntegerType RunMult(IntegerType x, IntegerType exponent) {
return vqrshlq_s32(x, exponent);
}
template <typename IntegerType>
static inline IntegerType RunMult(IntegerType x, int exponent) {
return vqrshlq_s32(x, vdupq_n_s32(static_cast<int32_t>(exponent)));
}
};
#endif // ARM NEON
// See CategorizeDotProductKernel for definitive taxonomy.
enum class DotProduct3x3KernelType {
kNone = 0, // Parameter combination is not supported for dot product kernels.
kPlain,
kWithDepthMultiplicationStride1,
kWithDepthMultiplicationStride2,
kStride2,
};
enum class QuantizationType {
kNonPerChannelUint8 = 0,
kPerChannelInt8 = 1,
};
template <QuantizationType quantization_type>
struct QuantizationTypeImpl {};
template <>
struct QuantizationTypeImpl<QuantizationType::kNonPerChannelUint8> {
typedef uint8_t ExternalType;
static constexpr int kIntSymmetricZeroPoint = 128;
static constexpr uint8_t kUint8SignBit = 0x80;
};
template <>
struct QuantizationTypeImpl<QuantizationType::kPerChannelInt8> {
typedef int8_t ExternalType;
static constexpr int kIntSymmetricZeroPoint = 0;
static constexpr uint8_t kUint8SignBit = 0x0;
};
template <
QuantizationType quantization_type = QuantizationType::kNonPerChannelUint8>
inline DotProduct3x3KernelType CategorizeDotProductKernel(
const RuntimeShape& input_shape, const RuntimeShape& filter_shape,
const RuntimeShape& output_shape, const DepthwiseParams& params,
const int32_t* output_shift_ptr = nullptr) {
constexpr int kSymmetricZeroPoint =
QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
const int padding =
std::max(params.padding_values.width, params.padding_values.height);
const int stride = params.stride_width;
const int32_t input_depth = input_shape.Dims(3);
const int32_t depth_multiplier = params.depth_multiplier;
const int32_t filter_height = filter_shape.Dims(1);
const int32_t filter_width = filter_shape.Dims(2);
bool supported = stride == params.stride_height && stride <= 2 &&
padding <= 1 && filter_width == 3 && filter_height == 3 &&
params.dilation_width_factor == 1 &&
params.dilation_height_factor == 1 &&
(((input_depth % 8) == 0 && depth_multiplier == 1) ||
(input_depth == 1 && depth_multiplier > 1));
if (!supported) {
return DotProduct3x3KernelType::kNone;
}
if (params.weights_offset != -kSymmetricZeroPoint) {
return DotProduct3x3KernelType::kNone;
}
if (quantization_type == QuantizationType::kPerChannelInt8) {
if (output_shift_ptr == nullptr) {
return DotProduct3x3KernelType::kNone;
}
} else if (params.output_shift > 0) {
return DotProduct3x3KernelType::kNone;
}
if (params.depth_multiplier == 1) {
if (stride == 1) {
return DotProduct3x3KernelType::kPlain;
} else if (stride == 2) {
return DotProduct3x3KernelType::kStride2;
} else {
return DotProduct3x3KernelType::kNone;
}
} else {
if (stride == 1) {
return DotProduct3x3KernelType::kWithDepthMultiplicationStride1;
} else if (stride == 2) {
return DotProduct3x3KernelType::kWithDepthMultiplicationStride2;
} else {
return DotProduct3x3KernelType::kNone;
}
}
}
// Encapsulates constant parameters used in DepthwiseConv.
// 64-bit is used for types that will be added to 64-bit addresses in asm.
struct DepthwiseConvParams {
int64_t input_depth;
int64_t input_row_size;
int64_t output_depth;
int64_t output_row_size;
int64_t filter_row_size;
int32_t input_offset;
int32_t output_offset;
int32_t filter_offset;
int32_t output_multiplier;
int32_t output_activation_min;
int32_t output_activation_max;
int32_t output_right_shift;
int32_t input_width;
int32_t input_height;
int32_t stride_width;
int32_t stride_height;
int32_t output_width;
int32_t output_height;
float float_output_activation_min;
float float_output_activation_max;
};
// Encapsulates constant parameters used in DepthwiseConv using dot-product ops.
// 64-bit is used for types that will be added to 64-bit addresses in asm.
//
// This structure is specifically designed for use in asm.
struct DepthwiseConvDotProdParams {
int64_t input_depth;
int64_t output_depth;
int32_t stride;
int32_t bias_increment;
//
int32_t input_offset;
int32_t output_offset;
int32_t output_multiplier;
int32_t output_shift;
int32_t quantized_activation_min;
int32_t quantized_activation_max;
//
int32_t padding_left;
int32_t padding_right;
int32_t padding_top;
int32_t padding_bottom;
//
int32_t depth_micro_repeats;
//
int32_t width_macro_count;
int32_t input_width_overall_micro_repeats;
int32_t input_width_micro_repeats;
int32_t residual_width;
int32_t output_width_overall_micro_repeats;
int32_t output_width_micro_repeats;
int32_t output_residual_width;
int32_t workspace_width_micro_repeats;
//
int32_t height_macro_count;
int32_t inbound_block_height;
int32_t outbound_block_height;
int32_t input_height_stride;
int32_t output_height_stride;
int32_t workspace_height_stride;
//
int32_t four_over_stride;
//
const int32_t* output_multiplier_per_channel;
const int32_t* output_shift_per_channel;
};
template <DepthwiseConvOutputRounding output_rounding, int32_t kDepth,
int32_t kStrideWidth, int32_t kStrideHeight>
struct DepthwiseConvWindow {};
template <DepthwiseConvOutputRounding output_rounding, int32_t kDepth,
int32_t kStrideWidth, int32_t kStrideHeight>
struct DepthwiseConvWindowPerChannel {};
enum class EdgeType { kCorner, kHorizontal, kVertical, kCenter };
template <DepthwiseConvOutputRounding output_rounding, EdgeType kEdgeType,
int kPadWidth, int kPadHeight>
struct DepthwiseConvPartial {};
template <DepthwiseConvOutputRounding output_rounding, EdgeType kEdgeType,
int kPadWidth, int kPadHeight>
struct DepthwiseConvPartialPerChannel {};
// Copies a subset of the input designated by |input_ptr| into |output_ptr|
// with the specified output dimensions. Supports output depths of 64 only as
// this is the cache line size.
template <typename T>
inline void ShuffleInput(const T* input_ptr, int64_t input_depth,
int32_t input_width, int32_t input_height,
int64_t output_depth, int32_t output_width,
int32_t output_height, T* output_ptr) {
const int64_t input_row_size = input_depth * input_width;
for (int32_t y = 0; y < output_height; y++) {
const T* ptr = input_ptr;
for (int32_t x = 0; x < output_width; x++) {
memcpy(output_ptr, ptr, output_depth);
output_ptr += output_depth;
ptr += input_depth;
}
input_ptr += input_row_size;
}
}
// Calculates the input size depending on stride and output.
inline int32_t get_shuffle_input_size(int32_t stride, int32_t output) {
return stride * (output - 1) + 3;
}
// Indicates the input and output dimensions used when shuffling input
// activations.
struct ShuffleParams {
int32_t output_width;
int32_t output_height;
int32_t input_width;
int32_t input_height;
ShuffleParams() = default;
ShuffleParams(int32_t output_width, int32_t output_height,
int32_t stride_width, int32_t stride_height)
: output_width(output_width),
output_height(output_height),
input_width(get_shuffle_input_size(stride_width, output_width)),
input_height(get_shuffle_input_size(stride_height, output_height)) {}
};
template <
QuantizationType quantization_type = QuantizationType::kNonPerChannelUint8>
inline bool Fast3x3FilterKernelSupported(
const RuntimeShape& input_shape, const RuntimeShape& filter_shape,
int32_t stride_width, int32_t stride_height, int32_t dilation_width_factor,
int32_t dilation_height_factor, int32_t pad_width, int32_t pad_height,
int32_t depth_multiplier, const RuntimeShape& output_shape,
int32_t output_shift, const int32_t* output_shift_ptr = nullptr) {
const int32_t input_height = input_shape.Dims(1);
const int32_t input_width = input_shape.Dims(2);
const int32_t input_depth = input_shape.Dims(3);
const int32_t filter_height = filter_shape.Dims(1);
const int32_t filter_width = filter_shape.Dims(2);
const int32_t output_height = output_shape.Dims(1);
const int32_t output_width = output_shape.Dims(2);
bool supported =
filter_width == 3 && filter_height == 3 && depth_multiplier == 1 &&
(stride_width == 1 || stride_width == 2) &&
(stride_height == 1 || stride_height == 2) &&
(stride_width == stride_height) && (pad_width == 0 || pad_width == 1) &&
(pad_height == 0 || pad_height == 1) && (pad_width == pad_height) &&
(input_depth % 8) == 0 && (output_shift <= 0) &&
dilation_width_factor == 1 && dilation_height_factor == 1;
if (!supported) {
return false;
}
// Handle case where padding is zero but padding type is not kValid.
// This would require special boundary case handling that is not supported.
const int32_t out_x = output_width - 1;
const int32_t out_y = output_height - 1;
const int32_t in_x_origin = (out_x * stride_width) - pad_width;
const int32_t in_y_origin = (out_y * stride_height) - pad_height;
const int32_t in_x_end = in_x_origin + filter_width;
const int32_t in_y_end = in_y_origin + filter_height;
// Supported only if filter on the right and bottom boundary lies completely
// within the input if padding is zero.
if (pad_width == 0 && pad_height == 0) {
return in_x_end <= input_width && in_y_end <= input_height;
}
// Else if padding is 1, supported if bottom right filter lies +1 past input
// width and height.
supported = in_x_end <= (input_width + 1) && in_y_end <= (input_height + 1);
if (!supported) {
return false;
}
// Shapes with width 1 and height > 1, and vice versa are not supported yet.
if (input_width == 1) {
supported = (input_width == input_height);
} else if (input_height == 1) {
supported = (input_width == input_height);
}
return supported;
}
// Permute filter data, and adjust bias data to account for symmetric input
// offset. Details are provided in the implementation of the
// kUseCModel3x3DotProduct version.
//
// See the comments preceding DepthwiseConvDotProduct3x3() for further notes.
template <DepthwiseConvImplementation implementation,
QuantizationType quantization_type>
struct ProcessPerDepth {
// Routine is contained in a static Run() method. No default template version
// is supplied, so that all implementations are deliberate choices of template
// specialization.
//
// Note that the signature of the Run() method will be designed for the asm
// implementation rather than conforming to style.
};
// Copy a macro block of data from the input buffer into the workspace,
// permuting data within each micro block.
//
// (a) Copy a macro block of data, padding as required along the width and
// height.
// (b) Transpose the data within each micro block.
//
// See the comments preceding DepthwiseConvDotProduct3x3() for further notes.
template <DepthwiseConvImplementation implementation,
QuantizationType quantization_type,
DepthwiseConvDepthMultiplication depth_multiplication,
int32_t max_padding>
struct PackMacroBlock {
// Routine is contained in a static Run() method. No default template version
// is supplied, so that all implementations are deliberate choices of template
// specialization.
//
// Note that the signature of the Run() method will be designed for the asm
// implementation rather than conforming to style.
};
// Apply filter to macro block of input data and store results. Details are
// provided in the implementation of the kUseCModel3x3DotProduct version.
//
// Parameters for repeats and residual sizes are in terms of outputs.
//
// See the comments preceding DepthwiseConvDotProduct3x3() for further notes.
template <DepthwiseConvImplementation implementation,
QuantizationType quantization_type,
DepthwiseConvDepthMultiplication depth_multiplication, int32_t stride>
struct KernelMacroBlock {
// Routine is contained in a static Run() method. No default template version
// is supplied, so that all implementations are deliberate choices of template
// specialization.
//
// Note that the signature of the Run() method will be designed for the asm
// implementation rather than conforming to style.
};
#if defined(__aarch64__)
// Experiments suggest that a modest performance improvement is seen, at least
// on 855 chipset big cores, with cache hints.
template <typename T>
inline void PreloadInputBlock(
const T* input_block_data,
const DepthwiseConvDotProdParams* function_params) {
// Preload.
const int input_width_micro_repeats =
function_params->input_width_micro_repeats;
const int block_height = function_params->inbound_block_height;
const int residual_width = function_params->residual_width;
const int input_height_stride = function_params->input_height_stride;
const int input_depth = function_params->input_depth;
const int total_width = 4 * input_width_micro_repeats + residual_width;
const T* row_ptr = input_block_data;
for (int k_height = 0; k_height < block_height; ++k_height) {
const T* ptr = row_ptr;
for (int j = 0; j < total_width; ++j) {
// Input data is loaded once.
optimized_ops_preload_l1_keep(ptr);
ptr += input_depth;
}
row_ptr += input_height_stride;
}
}
#endif // __aarch64__
} // namespace depthwise_conv
} // namespace optimized_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_3X3_FILTER_COMMON_H_
tensorflow/lite/kernels/internal/optimized/depthwiseconv_float.h
+1117
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/depthwiseconv_multithread.h
+189
View File
@@ -0,0 +1,189 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_MULTITHREAD_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_MULTITHREAD_H_
#include <algorithm>
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_float.h"
#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8.h"
namespace tflite {
namespace optimized_ops {
// TODO(luwa): add multithread to per-channel depthwise_conv
// DepthwiseConv can run with multi threads on the dim specified by thread_dim.
// Each thread processes output elements on dim, thread_dim, in the range of
// [thread_start, thread_end).
// For example, assume thread_start = 2, thread_end = 6, and thread_dim = 1, it
// means that it will calculate DepthwiseConv for output_data[:, 2:5, :, :].
template <typename T, typename TS>
struct DepthwiseConvWorkerTask : cpu_backend_threadpool::Task {
DepthwiseConvWorkerTask(const DepthwiseParams& params,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& filter_shape,
const T* filter_data, const RuntimeShape& bias_shape,
const TS* bias_data, const RuntimeShape& output_shape,
T* output_data, const CpuFlags& cpu_flags,
int thread_start, int thread_end, int thread_dim)
: params_(params),
input_shape_(input_shape),
input_data_(input_data),
filter_shape_(filter_shape),
filter_data_(filter_data),
bias_shape_(bias_shape),
bias_data_(bias_data),
output_shape_(output_shape),
output_data_(output_data),
cpu_flags_(cpu_flags),
thread_start_(thread_start),
thread_end_(thread_end),
thread_dim_(thread_dim) {}
void Run() override {
DepthwiseConvImpl(params_, input_shape_, input_data_, filter_shape_,
filter_data_, bias_shape_, bias_data_, output_shape_,
output_data_, cpu_flags_, thread_start_, thread_end_,
thread_dim_);
}
private:
const DepthwiseParams& params_;
const RuntimeShape& input_shape_;
const T* input_data_;
const RuntimeShape& filter_shape_;
const T* filter_data_;
const RuntimeShape& bias_shape_;
const TS* bias_data_;
const RuntimeShape& output_shape_;
T* output_data_;
const CpuFlags& cpu_flags_;
int thread_start_;
int thread_end_;
int thread_dim_;
};
inline int HowManyConvThreads(const RuntimeShape& output_shape,
const RuntimeShape& filter_shape) {
// How many scalar multiplications are needed to make it worth using one
// more thread
static constexpr int kMinMulPerThread = 1 << 13; // 8k
const int filter_height = filter_shape.Dims(1);
const int filter_width = filter_shape.Dims(2);
const int num_muls = output_shape.FlatSize() * filter_height * filter_width;
// Try to avoid real runtime divisions if possible by dividing by a
// compile-time constant.
int thread_count = std::max(1, num_muls / kMinMulPerThread);
return thread_count;
}
inline bool MultithreadAlongBatches(int thread_count, int batches) {
TFLITE_DCHECK_GE(thread_count, 2);
// If there are fewer batch entries than the number of threads we want to use,
// then better do intra-batch-entry multithreading.
if (batches < thread_count) {
return false;
}
// If there are at least 2 batch entries to be handed to each thread, then
// it's safe to proceed with batch-wise multithreading: each thread will have
// approximately equal number of batch entries to handle, so the load
// balancing will be reasonable, and the amount to which the load is not
// perfectly balanced will be offset by the inherent advantages of
// batch-wise multithreading (each thread is more efficient thanks to working
// on larger buffers with less boundary-handling overhead).
if (batches >= 2 * thread_count) {
return true;
}
// In the limit case were there are at least 1 but not much more than 1
// batch entries per thread, it may be a good idea to do per-batch
// multithreading if the number of batch entries is a multiple of the number
// of threads, so that each thread will have the same number of batch entries
// to process.
return ((batches % thread_count) == 0);
}
template <typename T, typename TS>
inline void DepthwiseConv(const DepthwiseParams& params,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& filter_shape,
const T* filter_data, const RuntimeShape& bias_shape,
const TS* bias_data, const RuntimeShape& output_shape,
T* output_data,
CpuBackendContext* cpu_backend_context) {
ruy::profiler::ScopeLabel label("DepthwiseConv");
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
int thread_count = HowManyConvThreads(output_shape, filter_shape);
const int max_threads = cpu_backend_context->max_num_threads();
thread_count = std::max(1, std::min(thread_count, max_threads));
#ifndef TFLITE_WITH_RUY
// Cap the number of threads to 2 for float path to avoid regression in
// performance (b/132294857).
if (std::is_floating_point<T>::value) {
thread_count = std::min(thread_count, 2);
}
#endif
const int output_batches = output_shape.Dims(0);
const int output_height = output_shape.Dims(1);
CpuFlags cpu_flags;
GetCpuFlags(&cpu_flags);
if (thread_count == 1) {
DepthwiseConvImpl(params, input_shape, input_data, filter_shape,
filter_data, bias_shape, bias_data, output_shape,
output_data, cpu_flags, /*thread_start=*/0,
/*thread_end=*/output_height, /*thread_dim=*/1);
return;
}
int thread_dim, thread_dim_size;
if (MultithreadAlongBatches(thread_count, output_batches)) {
thread_dim = 0;
thread_dim_size = output_batches;
} else {
thread_dim = 1;
thread_dim_size = output_height;
}
std::vector<DepthwiseConvWorkerTask<T, TS>> tasks;
// TODO(b/131746020) don't create new heap allocations every time.
// At least we make it a single heap allocation by using reserve().
tasks.reserve(thread_count);
int thread_start = 0;
for (int i = 0; i < thread_count; ++i) {
int thread_end =
thread_start + (thread_dim_size - thread_start) / (thread_count - i);
tasks.emplace_back(params, input_shape, input_data, filter_shape,
filter_data, bias_shape, bias_data, output_shape,
output_data, cpu_flags, thread_start, thread_end,
thread_dim);
thread_start = thread_end;
}
cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
cpu_backend_context);
}
} // namespace optimized_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_MULTITHREAD_H_
tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8.h
+2127
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h
+13442
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_transitional.h
+8129
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/eigen_spatial_convolutions.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
#define EIGEN_USE_CUSTOM_THREAD_POOL
#define EIGEN_USE_THREADS
#define Eigen EigenForTFLite
// NOTE: We need to define our own tensor contraction dispatch method before
// including the unsupported/Eigen/CXX11/Tensor header in order to reduce the
// total number of kernel instantiations.
// If you have trouble simply undef out the reducer macro e.g.
// TFLITE_REDUCE_INSTANTIATIONS, but be aware this will make
// the binary much bigger!
#define TFLITE_REDUCE_INSTANTIATIONS
#if defined(TFLITE_REDUCE_INSTANTIATIONS)
// Override Eigen tensor contraction dispatch method.
#define TENSOR_CONTRACTION_DISPATCH(METHOD, ALIGNMENT, ARGS) \
if (this->m_lhs_inner_dim_contiguous && this->m_rhs_inner_dim_contiguous && \
!this->m_rhs_inner_dim_reordered) { \
METHOD<true, true, false, ALIGNMENT> ARGS; \
} else { \
eigen_assert(false && "Unsupported contraction formats"); \
}
#endif
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h"
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
tensorflow/lite/kernels/internal/optimized/fully_connected_4bit.h
+152
View File
@@ -0,0 +1,152 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_FULLY_CONNECTED_4BIT_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_FULLY_CONNECTED_4BIT_H_
#include <stdint.h>
#ifndef TFLITE_MMAP_DISABLED
#include <sys/mman.h>
#endif
#include <cstdlib>
#include <memory>
#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
#include "tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected.h" // IWYU pragma: export
#elif defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected.h" // IWYU pragma: export
#else
#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference.h" // IWYU pragma: export
#endif
namespace tflite {
namespace optimized_4bit {
// Define 4-bit filter block size: 4x32 (64 bytes)
constexpr int FilterWidth = 4;
constexpr int FilterDepth = 32;
constexpr int kDefaultAlignmentPadding = 63;
struct Deleter {
explicit Deleter(size_t size = 0) : size(size) {}
void operator()(uint8_t* memory) {
if (!memory) {
return;
}
#ifdef TFLITE_MMAP_DISABLED
delete[] memory;
#else
munmap(memory, size);
#endif
}
size_t size;
};
struct OpData4Bit {
int rows_right = 1;
int batch_size = 0;
bool needs_prepack = true;
uint8_t* prepacked_cache = nullptr;
std::unique_ptr<uint8_t[], Deleter> prepacked_cache_buffer;
size_t prepacked_cache_buffer_size = 0;
void AllocatePackedRegion(size_t required_size) {
#ifdef TFLITE_MMAP_DISABLED
uint8_t* region = new uint8_t[required_size];
prepacked_cache_buffer =
std::unique_ptr<uint8_t[], Deleter>(region, Deleter());
#else
uint8_t* region = reinterpret_cast<uint8_t*>(
mmap(nullptr, required_size, PROT_READ | PROT_WRITE,
MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
prepacked_cache_buffer =
std::unique_ptr<uint8_t[], Deleter>(region, Deleter(required_size));
#ifdef MADV_MERGEABLE
madvise(region, required_size, MADV_MERGEABLE);
#endif
#endif
prepacked_cache = reinterpret_cast<uint8_t*>(
(reinterpret_cast<uintptr_t>(prepacked_cache_buffer.get()) +
kDefaultAlignmentPadding) &
~kDefaultAlignmentPadding);
prepacked_cache_buffer_size = required_size;
}
};
namespace api {
/* Prepack lhs matrix into dest.
* Transform tensor from (src_rows, src_cols) to
* (layout_rows / width, layout_cols / depth, width, depth) with possibly
* padding, and interleaving values along depth / 2 dimensions.
* dest should be aligned and allocated before prepack.
*/
inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
int layout_cols, int src_rows, int src_cols, int width,
int depth) {
optimized_4bit::Prepack(dest, tensor, layout_rows, layout_cols, src_rows,
src_cols, width, depth);
}
/* Quantize input floats to 8bit and calculate sum of each column.
* Data in float_data_ptr of shape (n_batch x n_data), is quantized and
* packed into (n_batch / width, n_data / depth, width, data) into
* quantized_data_ptr and input_offsets will contain the product of filter
* zero_point and input.
*/
inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
int n_data, int8_t* quantized_data_ptr,
float* scaling_factors, int width,
int depth, int32_t* input_offsets) {
optimized_4bit::BatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
quantized_data_ptr, scaling_factors,
width, depth, input_offsets);
}
/* Write bias + input offset * filter_scale to output_ptr.
* output_ptr of size (batch_size, output_depth) will have
* output_ptr[output_depth * b + o] =
* bias_ptr[o] + input_offsets[b] * batch_scales[b] * filter_scale[o]
*/
inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
const float* batch_scales,
float* filter_scales,
const float* bias_ptr,
float* output_ptr, int output_depth,
int batch_size) {
optimized_4bit::AssignBiasAndComputeOffsets(
input_offsets, batch_scales, filter_scales, bias_ptr, output_ptr,
output_depth, batch_size);
}
// Compute sum of lhs * rhs columnwise and write output to output_ptr.
inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
int32_t* dst, int output_depth, int batch_size,
int lhs_layout_rows, int lhs_layout_cols,
int rhs_layout_rows, int rhs_layout_cols,
int dst_layout_rows, int dst_layout_cols,
float* output_ptr, const float* scaling_factors,
const float* filter_scales) {
optimized_4bit::RunAndUnpack(
rhs_width, lhs, rhs, dst, output_depth, batch_size, lhs_layout_rows,
lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
dst_layout_cols, output_ptr, scaling_factors, filter_scales);
}
} // namespace api
} // namespace optimized_4bit
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_FULLY_CONNECTED_4BIT_H_
tensorflow/lite/kernels/internal/optimized/im2col_utils.h
+511
View File
@@ -0,0 +1,511 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_
#include <algorithm>
#include <cassert>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_ops {
template <typename T>
inline void ExtractPatchIntoBufferColumn(
const RuntimeShape& input_shape, int w, int h, int b, int kheight,
int kwidth, int stride_width, int stride_height, int pad_width,
int pad_height, int in_width, int in_height, int in_depth,
int single_buffer_length, int buffer_id, const T* in_data,
T* conv_buffer_data, uint8_t zero_byte) {
ruy::profiler::ScopeLabel label("ExtractPatchIntoBufferColumn");
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
// This chunk of code reshapes all the inputs corresponding to
// output (b, h, w) to a column vector in conv_buffer(:, buffer_id).
const int kwidth_times_indepth = kwidth * in_depth;
const int inwidth_times_indepth = in_width * in_depth;
const int ih_ungated_start = h * stride_height - pad_height;
const int ih_ungated_end = (ih_ungated_start + kheight);
const int ih_end = std::min(ih_ungated_end, in_height);
const int iw_ungated_start = w * stride_width - pad_width;
const int iw_ungated_end = (iw_ungated_start + kwidth);
const int iw_end = std::min(iw_ungated_end, in_width);
// If the patch is off the edge of the input image, skip writing those rows
// and columns from the patch into the output array.
const int h_offset = std::max(0, -ih_ungated_start);
const int w_offset = std::max(0, -iw_ungated_start);
const int ih_start = std::max(0, ih_ungated_start);
const int iw_start = std::max(0, iw_ungated_start);
const int single_row_num =
std::max(0, std::min(kwidth - w_offset, in_width - iw_start)) * in_depth;
const int output_row_offset = (buffer_id * single_buffer_length);
int out_offset =
output_row_offset + (h_offset * kwidth + w_offset) * in_depth;
int in_offset = Offset(input_shape, b, ih_start, iw_start, 0);
// Express all of the calculations as padding around the input patch.
const int top_padding = h_offset;
const int bottom_padding = (ih_ungated_end - ih_end);
const int left_padding = w_offset;
const int right_padding = (iw_ungated_end - iw_end);
assert(single_row_num ==
((kwidth - (left_padding + right_padding)) * in_depth));
// Write out zeroes to the elements representing the top rows of the input
// patch that are off the edge of the input image.
if (top_padding > 0) {
const int top_row_elements = (top_padding * kwidth * in_depth);
memset(conv_buffer_data + output_row_offset, zero_byte,
(top_row_elements * sizeof(T)));
}
// If the patch is on the interior of the input image horizontally, just copy
// over the rows sequentially, otherwise add zero padding at the start or end.
if ((left_padding == 0) && (right_padding == 0)) {
for (int ih = ih_start; ih < ih_end; ++ih) {
memcpy(conv_buffer_data + out_offset, in_data + in_offset,
single_row_num * sizeof(T));
out_offset += kwidth_times_indepth;
in_offset += inwidth_times_indepth;
}
} else {
for (int ih = ih_start; ih < ih_end; ++ih) {
if (left_padding > 0) {
const int left_start = (out_offset - (left_padding * in_depth));
memset(conv_buffer_data + left_start, zero_byte,
(left_padding * in_depth * sizeof(T)));
}
memcpy(conv_buffer_data + out_offset, in_data + in_offset,
single_row_num * sizeof(T));
if (right_padding > 0) {
const int right_start = (out_offset + single_row_num);
memset(conv_buffer_data + right_start, zero_byte,
(right_padding * in_depth * sizeof(T)));
}
out_offset += kwidth_times_indepth;
in_offset += inwidth_times_indepth;
}
}
// If the bottom of the patch falls off the input image, pad the values
// representing those input rows with zeroes.
if (bottom_padding > 0) {
const int bottom_row_elements = (bottom_padding * kwidth * in_depth);
const int bottom_start =
output_row_offset +
((top_padding + (ih_end - ih_start)) * kwidth * in_depth);
memset(conv_buffer_data + bottom_start, zero_byte,
(bottom_row_elements * sizeof(T)));
}
}
// Supports per-batch zero_byte for per-batch asymmetric quantized inputs.
template <typename T>
void DilatedIm2col(const ConvParams& params, const RuntimeShape& input_shape,
const T* input_data, const RuntimeShape& filter_shape,
const RuntimeShape& output_shape, T* im2col_data,
const int32_t* zero_bytes, const int zero_bytes_len) {
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const int dilation_width_factor = params.dilation_width_factor;
const int dilation_height_factor = params.dilation_height_factor;
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
// For dilated convolution, the input pixels are not contiguous therefore we
// can't use the same optimizations as Im2Col(). Though note this code would
// work fine for the non-dilated case too (though likely a bit slower).
ruy::profiler::ScopeLabel label("DilatedIm2col");
TFLITE_DCHECK(dilation_width_factor != 1 || dilation_height_factor != 1);
TFLITE_DCHECK(im2col_data);
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
const int input_height = input_shape.Dims(1);
const int input_width = input_shape.Dims(2);
const int input_depth = MatchingDim(input_shape, 3, filter_shape, 3);
const int filter_height = filter_shape.Dims(1);
const int filter_width = filter_shape.Dims(2);
const int output_height = output_shape.Dims(1);
const int output_width = output_shape.Dims(2);
MatchingDim(output_shape, 3, filter_shape, 0);
// Construct the MxN sized im2col matrix.
// The rows M, are sub-ordered B x H x W
const RuntimeShape row_shape({1, batches, output_height, output_width});
// The columns, N, are sub-ordered Kh x Kw x Din
const RuntimeShape col_shape({1, filter_height, filter_width, input_depth});
// Use dimensions M and N to construct dims for indexing directly into im2col
const RuntimeShape im2col_shape(
{1, 1, row_shape.FlatSize(), col_shape.FlatSize()});
// Loop through the output rows (B x H x W)
for (int batch = 0; batch < batches; ++batch) {
const T zero_byte = zero_bytes_len > 1 ? static_cast<T>(zero_bytes[batch])
: static_cast<T>(zero_bytes[0]);
for (int out_y = 0; out_y < output_height; ++out_y) {
for (int out_x = 0; out_x < output_width; ++out_x) {
// Each im2col row is an output pixel. Arrange the input data in this
// row in an order we can conveniently multiply with the filter data.
int row_offset = Offset(row_shape, 0, batch, out_y, out_x);
const int in_x_origin = (out_x * stride_width) - pad_width;
const int in_y_origin = (out_y * stride_height) - pad_height;
// Loop through all the pixels of the filter (Kh x Kw)
for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
const int in_y = in_y_origin + dilation_height_factor * filter_y;
if ((in_y >= 0) && (in_y < input_height)) {
// Filter row is within the input data.
// Loop through all the filter pixels in this row.
for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
const int in_x = in_x_origin + dilation_width_factor * filter_x;
int col_offset = Offset(col_shape, 0, filter_y, filter_x, 0);
T* dst = im2col_data +
Offset(im2col_shape, 0, 0, row_offset, col_offset);
if ((in_x >= 0) && (in_x < input_width)) {
// Filter pixel is within the input, copy the input data.
T const* src =
input_data + Offset(input_shape, batch, in_y, in_x, 0);
memcpy(dst, src, input_depth * sizeof(T));
} else {
// Filter pixel is outside the input, zero it out.
memset(dst, zero_byte, input_depth * sizeof(T));
}
}
} else {
// Filter row is outside the input, zero out the entire filter row.
int col_offset = Offset(col_shape, 0, filter_y, 0, 0);
T* dst = im2col_data +
Offset(im2col_shape, 0, 0, row_offset, col_offset);
memset(dst, zero_byte, filter_width * input_depth * sizeof(T));
}
}
}
}
}
}
template <typename T>
void DilatedIm2col(const ConvParams& params, uint8_t zero_byte,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& filter_shape,
const RuntimeShape& output_shape, T* im2col_data) {
const int32_t zero_point = static_cast<int32_t>(zero_byte);
DilatedIm2col<T>(params, input_shape, input_data, filter_shape, output_shape,
im2col_data, &zero_point, 1);
}
template <typename T>
void Im2col(const ConvParams& params, int kheight, int kwidth,
uint8_t zero_byte, const RuntimeShape& input_shape,
const T* input_data, const RuntimeShape& output_shape,
T* output_data) {
ruy::profiler::ScopeLabel label("Im2col");
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
const int input_depth = input_shape.Dims(3);
const int input_width = input_shape.Dims(2);
const int input_height = input_shape.Dims(1);
const int output_depth = output_shape.Dims(3);
const int output_width = output_shape.Dims(2);
const int output_height = output_shape.Dims(1);
int buffer_id = 0;
// Loop over the output nodes.
for (int b = 0; b < batches; ++b) {
for (int h = 0; h < output_height; ++h) {
for (int w = 0; w < output_width; ++w) {
ExtractPatchIntoBufferColumn(
input_shape, w, h, b, kheight, kwidth, stride_width, stride_height,
pad_width, pad_height, input_width, input_height, input_depth,
output_depth, buffer_id, input_data, output_data, zero_byte);
++buffer_id;
}
}
}
}
template <typename T>
void Im2col(const ConvParams& params, int kheight, int kwidth,
const int32_t* input_offsets, const int input_offsets_size,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& output_shape, T* output_data) {
ruy::profiler::ScopeLabel label("Im2col");
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
TFLITE_DCHECK_EQ(batches, input_offsets_size);
const int input_depth = input_shape.Dims(3);
const int input_width = input_shape.Dims(2);
const int input_height = input_shape.Dims(1);
const int output_depth = output_shape.Dims(3);
const int output_width = output_shape.Dims(2);
const int output_height = output_shape.Dims(1);
int buffer_id = 0;
// Loop over the output nodes.
for (int b = 0; b < batches; ++b) {
uint8_t zero_byte = static_cast<uint8_t>(input_offsets[b]);
for (int h = 0; h < output_height; ++h) {
for (int w = 0; w < output_width; ++w) {
ExtractPatchIntoBufferColumn(
input_shape, w, h, b, kheight, kwidth, stride_width, stride_height,
pad_width, pad_height, input_width, input_height, input_depth,
output_depth, buffer_id, input_data, output_data, zero_byte);
++buffer_id;
}
}
}
}
template <typename T>
inline void ExtractPatchIntoBufferColumn3D(
int b, int d, int h, int w, // Output indexes.
int kdepth, int kheight, int kwidth, // Kernel params.
int stride_depth, int stride_height, int stride_width, // Stride params.
int pad_depth, int pad_height, int pad_width, // Padding params.
int in_depth, int in_height, int in_width, int in_channel, // Input shape.
int output_row_offset, const T* in_data, T* conv_buffer_data,
uint8_t zero_byte) {
ruy::profiler::ScopeLabel label("ExtractPatchIntoBufferColumn3D");
// This chunk of code reshapes all the inputs corresponding to
// output (b, d, h, w) to a column vector in conv_buffer(:, buffer_id).
const int id_ungated_start = d * stride_depth - pad_depth;
const int id_start = std::max(0, id_ungated_start);
const int id_ungated_end = (id_ungated_start + kdepth);
const int id_end = std::min(id_ungated_end, in_depth);
const int ih_ungated_start = h * stride_height - pad_height;
const int ih_start = std::max(0, ih_ungated_start);
const int ih_ungated_end = (ih_ungated_start + kheight);
const int ih_end = std::min(ih_ungated_end, in_height);
const int iw_ungated_start = w * stride_width - pad_width;
const int iw_start = std::max(0, iw_ungated_start);
const int iw_ungated_end = (iw_ungated_start + kwidth);
const int iw_end = std::min(iw_ungated_end, in_width);
// Calculate the padding sizes.
const int d_padding_before = std::max(0, -id_ungated_start);
const int d_padding_after = (id_ungated_end - id_end);
const int h_padding_before = std::max(0, -ih_ungated_start);
const int h_padding_after = (ih_ungated_end - ih_end);
const int w_padding_before = std::max(0, -iw_ungated_start);
const int w_padding_after = (iw_ungated_end - iw_end);
// Memset if there are paddings in the depth dimension.
const int kd_stride_size = kheight * kwidth * in_channel;
const int id_stride_size = in_height * in_width * in_channel;
if (d_padding_before > 0) {
const int d_padding_before_elements = (d_padding_before * kd_stride_size);
memset(conv_buffer_data + output_row_offset, zero_byte,
(d_padding_before_elements * sizeof(T)));
}
if (d_padding_after > 0) {
const int d_padding_after_elements = (d_padding_after * kd_stride_size);
const int bottom_start =
output_row_offset + (kdepth - d_padding_after) * kd_stride_size;
memset(conv_buffer_data + bottom_start, zero_byte,
(d_padding_after_elements * sizeof(T)));
}
// If there are paddings in height or width dimension, menset the entire area
// to take advantage of sequential memory handling performance.
int out_offset = output_row_offset + d_padding_before * kd_stride_size;
if (h_padding_before > 0 || h_padding_after > 0 || w_padding_before > 0 ||
w_padding_after > 0) {
const int middle_elements = (id_end - id_start) * kd_stride_size;
memset(conv_buffer_data + out_offset, zero_byte,
(middle_elements * sizeof(T)));
}
// Copy the valid data from the input tensor.
const int kh_stride_size = kwidth * in_channel;
const int ih_stride_size = in_width * in_channel;
const int h_padding = h_padding_before + h_padding_after;
const int w_padding = w_padding_before + w_padding_after;
const int single_row_num = (kwidth - w_padding) * in_channel;
out_offset +=
h_padding_before * kh_stride_size + w_padding_before * in_channel;
const int in_offset_without_d = b * in_depth * id_stride_size +
ih_start * ih_stride_size +
iw_start * in_channel;
for (int id = id_start; id < id_end; ++id) {
int in_offset = in_offset_without_d + id * id_stride_size;
for (int ih = ih_start; ih < ih_end; ++ih) {
memcpy(conv_buffer_data + out_offset, in_data + in_offset,
single_row_num * sizeof(T));
out_offset += kh_stride_size;
in_offset += ih_stride_size;
}
out_offset += h_padding * kh_stride_size;
}
}
template <typename T>
void Im2col3D(const Conv3DParams& params, int kdepth, int kheight, int kwidth,
uint8_t zero_byte, const RuntimeShape& input_shape,
const T* input_data, const RuntimeShape& im2col_shape,
T* im2col_data) {
ruy::profiler::ScopeLabel label("Im2col3D");
const int stride_depth = params.stride_depth;
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const int pad_depth = params.padding_values.depth;
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 5);
TFLITE_DCHECK_EQ(im2col_shape.DimensionsCount(), 5);
const int batches = MatchingDim(input_shape, 0, im2col_shape, 0);
const int input_depth = input_shape.Dims(1);
const int input_height = input_shape.Dims(2);
const int input_width = input_shape.Dims(3);
const int input_channel = input_shape.Dims(4);
const int output_depth = im2col_shape.Dims(1);
const int output_height = im2col_shape.Dims(2);
const int output_width = im2col_shape.Dims(3);
const int output_channel = im2col_shape.Dims(4);
int buffer_id = 0;
// Loop over the output nodes.
for (int b = 0; b < batches; ++b) {
for (int d = 0; d < output_depth; ++d) {
for (int h = 0; h < output_height; ++h) {
for (int w = 0; w < output_width; ++w) {
ExtractPatchIntoBufferColumn3D(
b, d, h, w, kdepth, kheight, kwidth, stride_depth, stride_height,
stride_width, pad_depth, pad_height, pad_width, input_depth,
input_height, input_width, input_channel, buffer_id, input_data,
im2col_data, zero_byte);
buffer_id += output_channel;
}
}
}
}
}
template <typename T>
inline void DilatedIm2col3D(const Conv3DParams& params, int filter_depth,
int filter_height, int filter_width,
uint8_t zero_byte, const RuntimeShape& input_shape,
const T* input_data,
const RuntimeShape& im2col_shape, T* im2col_data) {
ruy::profiler::ScopeLabel label("DilatedIm2col3D");
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 5);
TFLITE_DCHECK_EQ(im2col_shape.DimensionsCount(), 5);
// Only NDHWC format is currently supported.
const int batches = MatchingDim(input_shape, 0, im2col_shape, 0);
const int input_channels = input_shape.Dims(4);
const int input_width = input_shape.Dims(3);
const int input_height = input_shape.Dims(2);
const int input_depth = input_shape.Dims(1);
const int output_width = im2col_shape.Dims(3);
const int output_height = im2col_shape.Dims(2);
const int output_depth = im2col_shape.Dims(1);
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
const int pad_depth = params.padding_values.depth;
// Construct the MxN sized im2col matrix.
// The rows M, are sub-ordered B x D x H x W.
const RuntimeShape row_shape(
{1, batches, output_depth, output_height, output_width});
// The columns, N, are sub-ordered Kd x Kh x Kw x Din.
const RuntimeShape col_shape(
{1, filter_depth, filter_height, filter_width, input_channels});
// Use dimensions M and N to construct dims for indexing directly into im2col.
const RuntimeShape im2col_reshaped(
{1, 1, row_shape.FlatSize(), col_shape.FlatSize()});
for (int batch = 0; batch < batches; ++batch) {
for (int out_d = 0; out_d < output_depth; ++out_d) {
const int in_d_origin = (out_d * params.stride_depth) - pad_depth;
for (int out_y = 0; out_y < output_height; ++out_y) {
const int in_y_origin = (out_y * params.stride_height) - pad_height;
for (int out_x = 0; out_x < output_width; ++out_x) {
const int in_x_origin = (out_x * params.stride_width) - pad_width;
const int row_offset =
Offset(row_shape, 0, batch, out_d, out_y, out_x);
for (int filter_d = 0; filter_d < filter_depth; ++filter_d) {
const int in_d = in_d_origin + params.dilation_depth * filter_d;
if ((in_d >= 0) && (in_d < input_depth)) {
for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
const int in_y =
in_y_origin + params.dilation_height * filter_y;
if ((in_y >= 0) && (in_y < input_height)) {
for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
const int in_x =
in_x_origin + params.dilation_width * filter_x;
int col_offset =
Offset(col_shape, 0, filter_d, filter_y, filter_x, 0);
T* dst = im2col_data + Offset(im2col_reshaped, 0, 0,
row_offset, col_offset);
if ((in_x >= 0) && (in_x < input_width)) {
// Filter pixel is within the input, copy the input data.
T const* src = input_data + Offset(input_shape, batch,
in_d, in_y, in_x, 0);
memcpy(dst, src, input_depth * sizeof(T));
} else {
// Filter pixel is outside the input, zero it out.
memset(dst, zero_byte, input_depth * sizeof(T));
}
}
} else {
const int col_offset =
Offset(col_shape, 0, filter_d, filter_y, 0, 0);
T* dst = im2col_data + Offset(im2col_reshaped, 0, 0,
row_offset, col_offset);
memset(dst, zero_byte,
filter_width * input_depth * sizeof(T));
}
}
} else {
const int col_offset = Offset(col_shape, 0, filter_d, 0, 0, 0);
T* dst = im2col_data +
Offset(im2col_reshaped, 0, 0, row_offset, col_offset);
memset(dst, zero_byte,
filter_height * filter_width * input_depth * sizeof(T));
}
}
}
}
}
}
}
} // namespace optimized_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/README.md
+8
View File
@@ -0,0 +1,8 @@
This directory contains optimized implementations for int8 fully integer kernels.
Weight filters of convs are expected to be symmetric per-channel quantized in
the range [-127, 127].
Inputs/activations are expected to be asymmetric per-layer quantized in the
range [-128, 127].
THESE ARE EXPERIMENTAL AND PRONE TO CHANGE.
tensorflow/lite/kernels/internal/optimized/integer_ops/add.h
+513
View File
@@ -0,0 +1,513 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_
#include <algorithm>
#include "fixedpoint/fixedpoint.h"
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
// Element-wise add that can often be used for inner loop of broadcast add as
// well as the non-broadcast add.
inline void AddElementwiseInt8(int size, const ArithmeticParams& params,
const int8* input1_data, const int8* input2_data,
int8* output_data) {
ruy::profiler::ScopeLabel label("AddElementwiseInt8/8bit");
int i = 0;
TFLITE_DCHECK_GT(params.input1_offset, -256);
TFLITE_DCHECK_GT(params.input2_offset, -256);
TFLITE_DCHECK_LT(params.input1_offset, 256);
TFLITE_DCHECK_LT(params.input2_offset, 256);
#ifdef USE_NEON
const int8x16_t output_activation_min_vector =
vdupq_n_s8(params.quantized_activation_min);
const int8x16_t output_activation_max_vector =
vdupq_n_s8(params.quantized_activation_max);
const int input1_left_shift = params.left_shift + params.input1_shift;
const int input2_left_shift = params.left_shift + params.input2_shift;
const int32x4_t input1_left_dup = vdupq_n_s32(input1_left_shift);
const int32x4_t input2_left_dup = vdupq_n_s32(input2_left_shift);
const int16x8_t input1_offset_dup = vdupq_n_s16(params.input1_offset);
const int16x8_t input2_offset_dup = vdupq_n_s16(params.input2_offset);
for (; i <= size - 16; i += 16) {
const int8x16_t input1_val_original = vld1q_s8(input1_data + i);
const int8x16_t input2_val_original = vld1q_s8(input2_data + i);
const int16x8_t input1_val_s16_high =
vmovl_s8(vget_high_s8(input1_val_original));
const int16x8_t input1_val_s16_low =
vmovl_s8(vget_low_s8(input1_val_original));
const int16x8_t input2_val_s16_high =
vmovl_s8(vget_high_s8(input2_val_original));
const int16x8_t input2_val_s16_low =
vmovl_s8(vget_low_s8(input2_val_original));
const int16x8_t input1_val_high =
vaddq_s16(input1_val_s16_high, input1_offset_dup);
const int16x8_t input2_val_high =
vaddq_s16(input2_val_s16_high, input2_offset_dup);
const int16x8_t input1_val_low =
vaddq_s16(input1_val_s16_low, input1_offset_dup);
const int16x8_t input2_val_low =
vaddq_s16(input2_val_s16_low, input2_offset_dup);
const int16x4_t input1_val_high_high = vget_high_s16(input1_val_high);
const int16x4_t input1_val_high_low = vget_low_s16(input1_val_high);
const int16x4_t input1_val_low_high = vget_high_s16(input1_val_low);
const int16x4_t input1_val_low_low = vget_low_s16(input1_val_low);
const int16x4_t input2_val_high_high = vget_high_s16(input2_val_high);
const int16x4_t input2_val_high_low = vget_low_s16(input2_val_high);
const int16x4_t input2_val_low_high = vget_high_s16(input2_val_low);
const int16x4_t input2_val_low_low = vget_low_s16(input2_val_low);
int32x4_t x111 = vmovl_s16(input1_val_low_low);
int32x4_t x112 = vmovl_s16(input1_val_low_high);
int32x4_t x121 = vmovl_s16(input1_val_high_low);
int32x4_t x122 = vmovl_s16(input1_val_high_high);
int32x4_t x211 = vmovl_s16(input2_val_low_low);
int32x4_t x212 = vmovl_s16(input2_val_low_high);
int32x4_t x221 = vmovl_s16(input2_val_high_low);
int32x4_t x222 = vmovl_s16(input2_val_high_high);
x111 = vshlq_s32(x111, input1_left_dup);
x112 = vshlq_s32(x112, input1_left_dup);
x121 = vshlq_s32(x121, input1_left_dup);
x122 = vshlq_s32(x122, input1_left_dup);
x211 = vshlq_s32(x211, input2_left_dup);
x212 = vshlq_s32(x212, input2_left_dup);
x221 = vshlq_s32(x221, input2_left_dup);
x222 = vshlq_s32(x222, input2_left_dup);
x111 = vqrdmulhq_n_s32(x111, params.input1_multiplier);
x112 = vqrdmulhq_n_s32(x112, params.input1_multiplier);
x121 = vqrdmulhq_n_s32(x121, params.input1_multiplier);
x122 = vqrdmulhq_n_s32(x122, params.input1_multiplier);
x211 = vqrdmulhq_n_s32(x211, params.input2_multiplier);
x212 = vqrdmulhq_n_s32(x212, params.input2_multiplier);
x221 = vqrdmulhq_n_s32(x221, params.input2_multiplier);
x222 = vqrdmulhq_n_s32(x222, params.input2_multiplier);
int32x4_t s11 = vaddq_s32(x111, x211);
int32x4_t s12 = vaddq_s32(x112, x212);
int32x4_t s21 = vaddq_s32(x121, x221);
int32x4_t s22 = vaddq_s32(x122, x222);
s11 = vqrdmulhq_n_s32(s11, params.output_multiplier);
s12 = vqrdmulhq_n_s32(s12, params.output_multiplier);
s21 = vqrdmulhq_n_s32(s21, params.output_multiplier);
s22 = vqrdmulhq_n_s32(s22, params.output_multiplier);
using gemmlowp::RoundingDivideByPOT;
s11 = RoundingDivideByPOT(s11, -params.output_shift);
s12 = RoundingDivideByPOT(s12, -params.output_shift);
s21 = RoundingDivideByPOT(s21, -params.output_shift);
s22 = RoundingDivideByPOT(s22, -params.output_shift);
const int16x4_t s11_narrowed = vmovn_s32(s11);
const int16x4_t s12_narrowed = vmovn_s32(s12);
const int16x4_t s21_narrowed = vmovn_s32(s21);
const int16x4_t s22_narrowed = vmovn_s32(s22);
const int16x8_t s1 = vaddq_s16(vcombine_s16(s11_narrowed, s12_narrowed),
vdupq_n_s16(params.output_offset));
const int16x8_t s2 = vaddq_s16(vcombine_s16(s21_narrowed, s22_narrowed),
vdupq_n_s16(params.output_offset));
const int8x16_t s = vcombine_s8(vqmovn_s16(s1), vqmovn_s16(s2));
const int8x16_t clamped =
vmaxq_s8(output_activation_min_vector,
vminq_s8(output_activation_max_vector, s));
vst1q_s8(output_data + i, clamped);
}
#endif // NEON
for (; i < size; ++i) {
const int32 input1_val = params.input1_offset + input1_data[i];
const int32 input2_val = params.input2_offset + input2_data[i];
const int32 shifted_input1_val = input1_val * (1 << params.left_shift);
const int32 shifted_input2_val = input2_val * (1 << params.left_shift);
const int32 scaled_input1_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input1_val, params.input1_multiplier, params.input1_shift);
const int32 scaled_input2_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input2_val, params.input2_multiplier, params.input2_shift);
const int32 raw_sum = scaled_input1_val + scaled_input2_val;
const int32 raw_output =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
raw_sum, params.output_multiplier, params.output_shift) +
params.output_offset;
const int32 clamped_output =
std::min(params.quantized_activation_max,
std::max(params.quantized_activation_min, raw_output));
output_data[i] = static_cast<int8>(clamped_output);
}
}
// Element-wise add is used for the non-broadcast add.
inline void AddElementwiseInt16(int size, const ArithmeticParams& params,
const int16* input1_data,
const int16* input2_data, int16* output_data) {
ruy::profiler::ScopeLabel label("AddElementwiseInt16/16bit");
int i = 0;
TFLITE_DCHECK_GT(params.input1_offset, -32768);
TFLITE_DCHECK_GT(params.input2_offset, -32768);
TFLITE_DCHECK_LT(params.input1_offset, 32768);
TFLITE_DCHECK_LT(params.input2_offset, 32768);
#ifdef __AVX2__
const int32_t input1_left_shift = params.left_shift + params.input1_shift;
const int32_t input2_left_shift = params.left_shift + params.input2_shift;
const __m256i input1_offset = _mm256_set1_epi32(params.input1_offset);
const __m256i input2_offset = _mm256_set1_epi32(params.input2_offset);
const __m256i output_offset = _mm256_set1_epi32(params.output_offset);
const __m256i clamp_max_v =
_mm256_set1_epi32(params.quantized_activation_max);
const __m256i clamp_min_v =
_mm256_set1_epi32(params.quantized_activation_min);
for (; i <= size - 16; i += 16) {
const __m256i input1_val_original =
_mm256_loadu_si256(reinterpret_cast<__m256i const*>(input1_data + i));
const __m256i input2_val_original =
_mm256_loadu_si256(reinterpret_cast<__m256i const*>(input2_data + i));
__m256i s11 =
_mm256_cvtepi16_epi32(_mm256_castsi256_si128(input1_val_original));
__m256i s12 =
_mm256_cvtepi16_epi32(_mm256_extracti128_si256(input1_val_original, 1));
__m256i s21 =
_mm256_cvtepi16_epi32(_mm256_castsi256_si128(input2_val_original));
__m256i s22 =
_mm256_cvtepi16_epi32(_mm256_extracti128_si256(input2_val_original, 1));
s11 = _mm256_add_epi32(s11, input1_offset);
s12 = _mm256_add_epi32(s12, input1_offset);
s21 = _mm256_add_epi32(s21, input2_offset);
s22 = _mm256_add_epi32(s22, input2_offset);
s11 = avx2_utils::MultiplyByQuantizedMultiplier(
s11, params.input1_multiplier, input1_left_shift);
s12 = avx2_utils::MultiplyByQuantizedMultiplier(
s12, params.input1_multiplier, input1_left_shift);
s21 = avx2_utils::MultiplyByQuantizedMultiplier(
s21, params.input2_multiplier, input2_left_shift);
s22 = avx2_utils::MultiplyByQuantizedMultiplier(
s22, params.input2_multiplier, input2_left_shift);
__m256i s1 = _mm256_add_epi32(s11, s21);
__m256i s2 = _mm256_add_epi32(s12, s22);
s1 = avx2_utils::MultiplyByQuantizedMultiplier(s1, params.output_multiplier,
params.output_shift);
s2 = avx2_utils::MultiplyByQuantizedMultiplier(s2, params.output_multiplier,
params.output_shift);
s1 = _mm256_add_epi32(s1, output_offset);
s2 = _mm256_add_epi32(s2, output_offset);
s1 = _mm256_min_epi32(s1, clamp_max_v);
s1 = _mm256_max_epi32(s1, clamp_min_v);
s2 = _mm256_min_epi32(s2, clamp_max_v);
s2 = _mm256_max_epi32(s2, clamp_min_v);
avx2_utils::CastInt32ToInt16AndStore(output_data + i, s1);
avx2_utils::CastInt32ToInt16AndStore(output_data + i + 8, s2);
}
#elif defined(USE_NEON)
const int32x4_t output_activation_min_vector =
vdupq_n_s32(params.quantized_activation_min);
const int32x4_t output_activation_max_vector =
vdupq_n_s32(params.quantized_activation_max);
const int input1_left_shift = params.left_shift + params.input1_shift;
const int input2_left_shift = params.left_shift + params.input2_shift;
const int32x4_t input1_left_dup = vdupq_n_s32(input1_left_shift);
const int32x4_t input2_left_dup = vdupq_n_s32(input2_left_shift);
const int32x4_t input1_offset_dup = vdupq_n_s32(params.input1_offset);
const int32x4_t input2_offset_dup = vdupq_n_s32(params.input2_offset);
const int32x4_t output_offset_dup = vdupq_n_s32(params.output_offset);
// Use the size 16 batch as it is effective on pixel 3/4.
for (; i <= size - 16; i += 16) {
const int16x8_t input11_val_original = vld1q_s16(input1_data + i);
const int16x8_t input12_val_original = vld1q_s16(input2_data + i);
const int16x8_t input21_val_original = vld1q_s16(input1_data + 8 + i);
const int16x8_t input22_val_original = vld1q_s16(input2_data + 8 + i);
int32x4_t x111 = vmovl_s16(vget_low_s16(input11_val_original));
int32x4_t x112 = vmovl_s16(vget_high_s16(input11_val_original));
int32x4_t x121 = vmovl_s16(vget_low_s16(input12_val_original));
int32x4_t x122 = vmovl_s16(vget_high_s16(input12_val_original));
int32x4_t x211 = vmovl_s16(vget_low_s16(input21_val_original));
int32x4_t x212 = vmovl_s16(vget_high_s16(input21_val_original));
int32x4_t x221 = vmovl_s16(vget_low_s16(input22_val_original));
int32x4_t x222 = vmovl_s16(vget_high_s16(input22_val_original));
x111 = vaddq_s32(x111, input1_offset_dup);
x112 = vaddq_s32(x112, input1_offset_dup);
x121 = vaddq_s32(x121, input2_offset_dup);
x122 = vaddq_s32(x122, input2_offset_dup);
x211 = vaddq_s32(x211, input1_offset_dup);
x212 = vaddq_s32(x212, input1_offset_dup);
x221 = vaddq_s32(x221, input2_offset_dup);
x222 = vaddq_s32(x222, input2_offset_dup);
x111 = vshlq_s32(x111, input1_left_dup);
x112 = vshlq_s32(x112, input1_left_dup);
x121 = vshlq_s32(x121, input2_left_dup);
x122 = vshlq_s32(x122, input2_left_dup);
x211 = vshlq_s32(x211, input1_left_dup);
x212 = vshlq_s32(x212, input1_left_dup);
x221 = vshlq_s32(x221, input2_left_dup);
x222 = vshlq_s32(x222, input2_left_dup);
x111 = vqrdmulhq_n_s32(x111, params.input1_multiplier);
x112 = vqrdmulhq_n_s32(x112, params.input1_multiplier);
x121 = vqrdmulhq_n_s32(x121, params.input2_multiplier);
x122 = vqrdmulhq_n_s32(x122, params.input2_multiplier);
x211 = vqrdmulhq_n_s32(x211, params.input1_multiplier);
x212 = vqrdmulhq_n_s32(x212, params.input1_multiplier);
x221 = vqrdmulhq_n_s32(x221, params.input2_multiplier);
x222 = vqrdmulhq_n_s32(x222, params.input2_multiplier);
int32x4_t s11 = vaddq_s32(x111, x121);
int32x4_t s12 = vaddq_s32(x112, x122);
int32x4_t s21 = vaddq_s32(x211, x221);
int32x4_t s22 = vaddq_s32(x212, x222);
s11 = vqrdmulhq_n_s32(s11, params.output_multiplier);
s12 = vqrdmulhq_n_s32(s12, params.output_multiplier);
s21 = vqrdmulhq_n_s32(s21, params.output_multiplier);
s22 = vqrdmulhq_n_s32(s22, params.output_multiplier);
using gemmlowp::RoundingDivideByPOT;
s11 = RoundingDivideByPOT(s11, -params.output_shift);
s12 = RoundingDivideByPOT(s12, -params.output_shift);
s21 = RoundingDivideByPOT(s21, -params.output_shift);
s22 = RoundingDivideByPOT(s22, -params.output_shift);
s11 = vaddq_s32(s11, output_offset_dup);
s12 = vaddq_s32(s12, output_offset_dup);
s21 = vaddq_s32(s21, output_offset_dup);
s22 = vaddq_s32(s22, output_offset_dup);
s11 = vmaxq_s32(output_activation_min_vector,
vminq_s32(output_activation_max_vector, s11));
s12 = vmaxq_s32(output_activation_min_vector,
vminq_s32(output_activation_max_vector, s12));
s21 = vmaxq_s32(output_activation_min_vector,
vminq_s32(output_activation_max_vector, s21));
s22 = vmaxq_s32(output_activation_min_vector,
vminq_s32(output_activation_max_vector, s22));
const int16x8_t s1 = vcombine_s16(vqmovn_s32(s11), vqmovn_s32(s12));
const int16x8_t s2 = vcombine_s16(vqmovn_s32(s21), vqmovn_s32(s22));
vst1q_s16(output_data + i, s1);
vst1q_s16(output_data + 8 + i, s2);
}
#endif // NEON
for (; i < size; ++i) {
const int32 input1_val = params.input1_offset + input1_data[i];
const int32 input2_val = params.input2_offset + input2_data[i];
const int32 shifted_input1_val = input1_val * (1 << params.left_shift);
const int32 shifted_input2_val = input2_val * (1 << params.left_shift);
const int32 scaled_input1_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input1_val, params.input1_multiplier, params.input1_shift);
const int32 scaled_input2_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input2_val, params.input2_multiplier, params.input2_shift);
const int32 raw_sum = scaled_input1_val + scaled_input2_val;
const int32 raw_output =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
raw_sum, params.output_multiplier, params.output_shift) +
params.output_offset;
const int32 clamped_output =
std::min(params.quantized_activation_max,
std::max(params.quantized_activation_min, raw_output));
output_data[i] = static_cast<int16>(clamped_output);
}
}
// Scalar-broadcast add that can be used for inner loop of more general
// broadcast add, so that, for example, scalar-broadcast with batch will still
// be fast.
inline void AddScalarBroadcast(int size, const ArithmeticParams& params,
int8 input1_data, const int8* input2_data,
int8* output_data) {
using gemmlowp::RoundingDivideByPOT;
ruy::profiler::ScopeLabel label("AddScalarBroadcastInt8/8bit");
TFLITE_DCHECK_GT(params.input1_offset, -256);
TFLITE_DCHECK_GT(params.input2_offset, -256);
TFLITE_DCHECK_LT(params.input1_offset, 256);
TFLITE_DCHECK_LT(params.input2_offset, 256);
int i = 0;
#ifdef USE_NEON
const int32x4_t left_shift_dup = vdupq_n_s32(params.left_shift);
const int8x8_t output_activation_min_vector =
vdup_n_s8(params.quantized_activation_min);
const int8x8_t output_activation_max_vector =
vdup_n_s8(params.quantized_activation_max);
// Process broadcast scalar.
const int8x8_t input1_val_original = vdup_n_s8(input1_data);
const int16x8_t input1_val_s16 = vmovl_s8(input1_val_original);
const int16x8_t input1_val =
vaddq_s16(input1_val_s16, vdupq_n_s16(params.input1_offset));
const int16x4_t input1_val_high = vget_high_s16(input1_val);
const int16x4_t input1_val_low = vget_low_s16(input1_val);
int32x4_t x11 = vmovl_s16(input1_val_low);
int32x4_t x12 = vmovl_s16(input1_val_high);
x11 = vshlq_s32(x11, left_shift_dup);
x12 = vshlq_s32(x12, left_shift_dup);
x11 = vqrdmulhq_n_s32(x11, params.input1_multiplier);
x12 = vqrdmulhq_n_s32(x12, params.input1_multiplier);
const int32x4_t input1_shift_dup = vdupq_n_s32(params.input1_shift);
x11 = vshlq_s32(x11, input1_shift_dup);
x12 = vshlq_s32(x12, input1_shift_dup);
for (; i <= size - 8; i += 8) {
const int8x8_t input2_val_original = vld1_s8(input2_data + i);
const int16x8_t input2_val_s16 = vmovl_s8(input2_val_original);
const int16x8_t input2_val =
vaddq_s16(input2_val_s16, vdupq_n_s16(params.input2_offset));
const int16x4_t input2_val_high = vget_high_s16(input2_val);
const int16x4_t input2_val_low = vget_low_s16(input2_val);
int32x4_t x21 = vmovl_s16(input2_val_low);
int32x4_t x22 = vmovl_s16(input2_val_high);
x21 = vshlq_s32(x21, left_shift_dup);
x22 = vshlq_s32(x22, left_shift_dup);
x21 = vqrdmulhq_n_s32(x21, params.input2_multiplier);
x22 = vqrdmulhq_n_s32(x22, params.input2_multiplier);
const int32x4_t input2_shift_dup = vdupq_n_s32(params.input2_shift);
x21 = vshlq_s32(x21, input2_shift_dup);
x22 = vshlq_s32(x22, input2_shift_dup);
int32x4_t s1 = vaddq_s32(x11, x21);
int32x4_t s2 = vaddq_s32(x12, x22);
s1 = vqrdmulhq_n_s32(s1, params.output_multiplier);
s2 = vqrdmulhq_n_s32(s2, params.output_multiplier);
s1 = RoundingDivideByPOT(s1, -params.output_shift);
s2 = RoundingDivideByPOT(s2, -params.output_shift);
const int16x4_t s1_narrowed = vmovn_s32(s1);
const int16x4_t s2_narrowed = vmovn_s32(s2);
const int16x8_t s = vaddq_s16(vcombine_s16(s1_narrowed, s2_narrowed),
vdupq_n_s16(params.output_offset));
const int8x8_t clamped =
vmax_s8(output_activation_min_vector,
vmin_s8(output_activation_max_vector, vqmovn_s16(s)));
vst1_s8(output_data + i, clamped);
}
#endif // NEON
if (i < size) {
// Process broadcast scalar.
const int32 input1_val = params.input1_offset + input1_data;
const int32 shifted_input1_val = input1_val * (1 << params.left_shift);
const int32 scaled_input1_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input1_val, params.input1_multiplier, params.input1_shift);
for (; i < size; ++i) {
const int32 input2_val = params.input2_offset + input2_data[i];
const int32 shifted_input2_val = input2_val * (1 << params.left_shift);
const int32 scaled_input2_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input2_val, params.input2_multiplier,
params.input2_shift);
const int32 raw_sum = scaled_input1_val + scaled_input2_val;
const int32 raw_output =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
raw_sum, params.output_multiplier, params.output_shift) +
params.output_offset;
const int32 clamped_output =
std::min(params.quantized_activation_max,
std::max(params.quantized_activation_min, raw_output));
output_data[i] = static_cast<int8>(clamped_output);
}
}
}
inline void Add(const ArithmeticParams& params,
const RuntimeShape& input1_shape, const int8* input1_data,
const RuntimeShape& input2_shape, const int8* input2_data,
const RuntimeShape& output_shape, int8* output_data) {
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
ruy::profiler::ScopeLabel label("AddInt8/8bit");
const int flat_size =
MatchingElementsSize(input1_shape, input2_shape, output_shape);
TFLITE_DCHECK_GT(params.input1_offset, -256);
TFLITE_DCHECK_GT(params.input2_offset, -256);
TFLITE_DCHECK_LT(params.input1_offset, 256);
TFLITE_DCHECK_LT(params.input2_offset, 256);
AddElementwiseInt8(flat_size, params, input1_data, input2_data, output_data);
}
inline void Add(const ArithmeticParams& params,
const RuntimeShape& input1_shape, const int16* input1_data,
const RuntimeShape& input2_shape, const int16* input2_data,
const RuntimeShape& output_shape, int16* output_data) {
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
ruy::profiler::ScopeLabel label("AddInt16/16bit");
const int flat_size =
MatchingElementsSize(input1_shape, input2_shape, output_shape);
TFLITE_DCHECK_GT(params.input1_offset, -32768);
TFLITE_DCHECK_GT(params.input2_offset, -32768);
TFLITE_DCHECK_LT(params.input1_offset, 32768);
TFLITE_DCHECK_LT(params.input2_offset, 32768);
AddElementwiseInt16(flat_size, params, input1_data, input2_data, output_data);
}
inline void BroadcastAddDispatch(const ArithmeticParams& params,
const RuntimeShape& input1_shape,
const int8* input1_data,
const RuntimeShape& input2_shape,
const int8* input2_data,
const RuntimeShape& output_shape,
int8* output_data) {
if (params.broadcast_category == BroadcastableOpCategory::kGenericBroadcast) {
return reference_integer_ops::BroadcastAdd6DSlow(
params, input1_shape, input1_data, input2_shape, input2_data,
output_shape, output_data);
}
optimized_ops::BinaryBroadcastFiveFold(
params, input1_shape, input1_data, input2_shape, input2_data,
output_shape, output_data, AddElementwiseInt8, AddScalarBroadcast);
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/conv.h
+130
View File
@@ -0,0 +1,130 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_CONV_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_CONV_H_
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/cpu_backend_gemm.h"
#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/optimized/im2col_utils.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
// Fixed-point per-channel-quantization convolution reference kernel.
template <typename InputScalar, typename DstScalar>
inline void ConvPerChannel(
const ConvParams& params, const int32* output_multiplier,
const int32* output_shift, const RuntimeShape& input_shape,
const InputScalar* input_data, const RuntimeShape& filter_shape,
const int8* filter_data, const RuntimeShape& bias_shape,
const int32* bias_data, const RuntimeShape& output_shape,
DstScalar* output_data, const RuntimeShape& im2col_shape,
InputScalar* im2col_data, CpuBackendContext* cpu_backend_context) {
ruy::profiler::ScopeLabel label("Conv/8bit");
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const int dilation_width_factor = params.dilation_width_factor;
const int dilation_height_factor = params.dilation_height_factor;
const int32 input_offset = params.input_offset;
const int32 output_offset = params.output_offset;
// Set min and max value of the output.
const int32 output_activation_min = params.quantized_activation_min;
const int32 output_activation_max = params.quantized_activation_max;
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const InputScalar* gemm_input_data = nullptr;
const RuntimeShape* gemm_input_shape = nullptr;
const int filter_width = filter_shape.Dims(2);
const int filter_height = filter_shape.Dims(1);
const bool need_dilated_im2col =
dilation_width_factor != 1 || dilation_height_factor != 1;
const bool need_im2col = stride_width != 1 || stride_height != 1 ||
filter_width != 1 || filter_height != 1;
const int8 input_zero_point = -input_offset;
const uint8 zero_point_byte =
*reinterpret_cast<const uint8*>(&input_zero_point);
if (need_dilated_im2col) {
TFLITE_DCHECK(im2col_data);
optimized_ops::DilatedIm2col(params, zero_point_byte, input_shape,
input_data, filter_shape, output_shape,
im2col_data);
gemm_input_data = im2col_data;
gemm_input_shape = &im2col_shape;
} else if (need_im2col) {
TFLITE_DCHECK(im2col_data);
optimized_ops::Im2col(params, filter_height, filter_width, zero_point_byte,
input_shape, input_data, im2col_shape, im2col_data);
gemm_input_data = im2col_data;
gemm_input_shape = &im2col_shape;
} else {
TFLITE_DCHECK(!im2col_data);
gemm_input_data = input_data;
gemm_input_shape = &input_shape;
}
const int gemm_input_rows = gemm_input_shape->Dims(3);
const int gemm_input_cols = FlatSizeSkipDim(*gemm_input_shape, 3);
const int filter_rows = filter_shape.Dims(0);
const int filter_cols = FlatSizeSkipDim(filter_shape, 0);
const int output_rows = output_shape.Dims(3);
// See b/79927784.
// const int output_cols = FlatSizeSkipDim(output_shape, 3);
const int output_cols =
output_shape.Dims(0) * output_shape.Dims(1) * output_shape.Dims(2);
TFLITE_DCHECK_EQ(output_rows, filter_rows);
TFLITE_DCHECK_EQ(output_cols, gemm_input_cols);
TFLITE_DCHECK_EQ(filter_cols, gemm_input_rows);
TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_rows);
cpu_backend_gemm::MatrixParams<int8> lhs_params;
lhs_params.rows = filter_rows;
lhs_params.cols = filter_cols;
lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
lhs_params.zero_point = 0; // filter is symmetric-quantized
cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
rhs_params.rows = gemm_input_rows;
rhs_params.cols = gemm_input_cols;
rhs_params.order = cpu_backend_gemm::Order::kColMajor;
rhs_params.zero_point = -input_offset;
cpu_backend_gemm::MatrixParams<DstScalar> dst_params;
dst_params.rows = output_rows;
dst_params.cols = output_cols;
dst_params.order = cpu_backend_gemm::Order::kColMajor;
dst_params.zero_point = output_offset;
cpu_backend_gemm::GemmParams<
int32, DstScalar,
cpu_backend_gemm::QuantizationFlavor::kIntegerWithPerRowMultiplier>
gemm_params;
gemm_params.bias = bias_data;
gemm_params.clamp_min = output_activation_min;
gemm_params.clamp_max = output_activation_max;
gemm_params.multiplier_fixedpoint_perchannel = output_multiplier;
gemm_params.multiplier_exponent_perchannel = output_shift;
cpu_backend_gemm::Gemm(lhs_params, filter_data, rhs_params, gemm_input_data,
dst_params, output_data, gemm_params,
cpu_backend_context);
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_CONV_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h
+2026
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_3x3_filter.h
+3083
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid.h
+511
View File
@@ -0,0 +1,511 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_H_
#include <algorithm>
#include <memory>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_3x3_filter_common.h"
#include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h"
#include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h"
#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
namespace depthwise_conv {
// Initializes the accumulator buffer with zeros.
inline void DepthwiseConvInitAccBuffer(int num_output_pixels, int output_depth,
int32* acc_buffer) {
memset(acc_buffer, 0,
sizeof(acc_buffer[0]) * output_depth * num_output_pixels);
}
// Base DWConv Implementation used with both static and dynamic
// accumulator buffers.
// Initializes the accumulator buffer with bias values.
static void DoDepthwiseConvHybridGeneral(
const DepthwiseParams& params, const float* input_scales,
const RuntimeShape& input_shape, const int8* input_data,
const RuntimeShape& filter_shape, const int8* filter_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data,
const float* per_channel_scales, const int32_t* input_offsets,
int thread_start, int thread_end, int thread_dim, int32* acc_buffer,
int32 acc_buffer_size) {
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
const int depth_multiplier = params.depth_multiplier;
const float output_activation_min = params.float_activation_min;
const float output_activation_max = params.float_activation_max;
const int dilation_width_factor = params.dilation_width_factor;
const int dilation_height_factor = params.dilation_height_factor;
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
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 input_depth = input_shape.Dims(3);
const int filter_height = filter_shape.Dims(1);
const int filter_width = filter_shape.Dims(2);
const int output_rows = output_shape.Dims(1);
const int output_width = output_shape.Dims(2);
TFLITE_DCHECK_GE(acc_buffer_size, output_depth);
const int kOutputPixelsInAccBuffer = acc_buffer_size / output_depth;
const int kAccBufferActualSize = kOutputPixelsInAccBuffer * output_depth;
TFLITE_DCHECK_LE(kOutputPixelsInAccBuffer * output_depth,
kAccBufferActualSize);
TFLITE_DCHECK_LE(kAccBufferActualSize, acc_buffer_size);
TFLITE_DCHECK_GE(kOutputPixelsInAccBuffer, 1);
TFLITE_DCHECK(thread_dim == 0 || thread_dim == 1);
// row_accum_func will point to the core accumulation function to be used
// for this DepthwiseConvHybrid op.
using row_accum_func_t = decltype(&QuantizedDepthwiseConvAccumRowGeneric);
row_accum_func_t row_accum_func = nullptr;
#define TFMINI_USE_DEPTHWISECONV_KERNEL(ALLOW_STRIDED, FIXED_INPUT_DEPTH, \
FIXED_DEPTH_MULTIPLIER) \
if (!row_accum_func && (stride_width == 1 || ALLOW_STRIDED) && \
(input_depth == FIXED_INPUT_DEPTH || FIXED_INPUT_DEPTH == 0) && \
depth_multiplier == FIXED_DEPTH_MULTIPLIER) { \
row_accum_func = \
QuantizedDepthwiseConvAccumRow<ALLOW_STRIDED, FIXED_INPUT_DEPTH, \
FIXED_DEPTH_MULTIPLIER>; \
}
#ifdef USE_NEON
// We go over our list of kernels by decreasing order of preference
// for the cases where multiple kernels could apply.
// Start with the fastest kernels: AllowStrided=false, fixed input depth.
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 1, 2)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 2)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 2)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 1, 4)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 1)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 4)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 8, 1)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 8)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 1)
TFMINI_USE_DEPTHWISECONV_KERNEL(false, 12, 1)
// Next come the strided kernels: AllowStrided=true, fixed input depth.
// They are a bit less efficient, but allow stride!=1.
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 2)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 16, 1)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 16)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 20)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 32)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 8)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 1)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 2, 1)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 4, 1)
// Finally, the kernels allowing a variable input depth,
// these are the least efficient but most general kernels.
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 1)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 2)
TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 3)
#endif // USE_NEON
// No matching fast kernel found, use slow fallback.
if (!row_accum_func) {
row_accum_func = QuantizedDepthwiseConvAccumRowGeneric;
}
#undef TFMINI_USE_DEPTHWISECONV_KERNEL
const int input_height_stride = input_shape.Dims(3) * input_shape.Dims(2);
const int input_batch_stride = input_height_stride * input_shape.Dims(1);
const int filter_height_stride = filter_shape.Dims(3) * filter_shape.Dims(2);
// Now that we have determined row_accum_func, we can start work.
int batch_start = 0;
int batch_end = batches;
int row_start = 0;
int row_end = output_rows;
int output_ptr_offset = 0;
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;
output_ptr_offset = batch_start * FlatSizeSkipDim(output_shape, 0);
break;
case 1:
TFLITE_DCHECK_GE(thread_start, 0);
TFLITE_DCHECK_LE(thread_end, output_rows);
row_start = thread_start;
row_end = thread_end;
output_ptr_offset = row_start * output_width * output_depth;
break;
}
float* output_ptr = output_data + output_ptr_offset;
int batch_step =
(output_rows + row_start - row_end) * output_width * output_depth;
for (int b = batch_start; b < batch_end; ++b) {
float input_scale = input_scales[b];
int32_t input_offset = input_offsets[b];
for (int out_y = row_start; out_y < row_end; ++out_y) {
const int in_y_origin = (out_y * stride_height) - pad_height;
const int filter_y_start =
std::max(0, (-in_y_origin + dilation_height_factor - 1) /
dilation_height_factor);
const int filter_y_end =
std::min(filter_height,
(input_height - in_y_origin + dilation_height_factor - 1) /
dilation_height_factor);
for (int out_x_buffer_start = 0; out_x_buffer_start < output_width;
out_x_buffer_start += kOutputPixelsInAccBuffer) {
const int out_x_buffer_end = std::min(
output_width, out_x_buffer_start + kOutputPixelsInAccBuffer);
// We call a 'pixel' a group of activation that share all but the
// 'depth'/'channel' coordinate. num_output_pixels is the number of
// output pixels that we will accumulate in this loop iteration.
const int num_output_pixels = out_x_buffer_end - out_x_buffer_start;
DepthwiseConvInitAccBuffer(num_output_pixels, output_depth, acc_buffer);
// Accumulation loop. Most of the time should be spent in here.
for (int filter_y = filter_y_start; filter_y < filter_y_end;
++filter_y) {
const int in_y = in_y_origin + dilation_height_factor * filter_y;
row_accum_func(
stride_width, dilation_width_factor, input_depth, input_width,
input_data + in_y * input_height_stride + b * input_batch_stride,
-input_offset, pad_width, depth_multiplier, filter_width,
filter_data + filter_y * filter_height_stride, out_x_buffer_start,
out_x_buffer_end, output_depth, acc_buffer);
}
// Finished accumulating int32 values. Just store them as float values
gemmlowp::ScopedProfilingLabel label("store");
const int num_output_values = output_depth * num_output_pixels;
int c = 0;
while (c < output_depth) {
int target_output_depth = output_depth;
#ifdef USE_NEON
const float32x4_t output_activation_min_vec =
vdupq_n_f32(output_activation_min);
const float32x4_t output_activation_max_vec =
vdupq_n_f32(output_activation_max);
const float32x4_t input_scale_32x4 = vdupq_n_f32(input_scale);
for (; c <= output_depth - 4; c += 4) {
if ((c + 4) > output_depth) {
break;
}
const float32x4_t channel_scale_32x4 =
vld1q_f32(per_channel_scales + c);
const float32x4_t bias_32x4 = vld1q_f32(bias_data + c);
for (int n = 0; n < num_output_pixels; ++n) {
int loc = n * output_depth + c;
int32x4_t acc = vld1q_s32(acc_buffer + loc);
float32x4_t float_acc = vcvtq_f32_s32(acc);
float_acc = vmulq_f32(float_acc, channel_scale_32x4);
float_acc = vmulq_f32(float_acc, input_scale_32x4);
float_acc = vaddq_f32(float_acc, bias_32x4);
float_acc = vmaxq_f32(float_acc, output_activation_min_vec);
float_acc = vminq_f32(float_acc, output_activation_max_vec);
vst1q_f32(output_ptr + loc, float_acc);
}
}
#endif // USE_NEON
for (; c < target_output_depth; c++) {
for (int n = 0; n < num_output_pixels; ++n) {
int loc = n * output_depth + c;
int32 acc = acc_buffer[loc];
float float_acc = acc * input_scale * per_channel_scales[c];
float_acc += bias_data[c];
float_acc = std::max(float_acc, output_activation_min);
float_acc = std::min(float_acc, output_activation_max);
output_ptr[loc] = float_acc;
}
}
}
output_ptr += num_output_values;
}
}
output_ptr += batch_step;
}
}
// Utilize the base implementation of DWConv with a stack allocated accumulator
// buffer. The static allocation limits the number of depthwise channels that
// can be processed to kStaticAccBufferMaxSize.
static void DoDepthwiseConvHybridGeneralStatic(
const DepthwiseParams& params, const float* input_scales,
const RuntimeShape& input_shape, const int8* input_data,
const RuntimeShape& filter_shape, const int8* filter_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data,
const float* per_channel_scales, const int32_t* input_offsets,
int thread_start, int thread_end, int thread_dim) {
static const int kStaticAccBufferMaxSize = 2048;
int32 stack_acc_buffer[kStaticAccBufferMaxSize];
DoDepthwiseConvHybridGeneral(
params, input_scales, input_shape, input_data, filter_shape, filter_data,
bias_shape, bias_data, output_shape, output_data, per_channel_scales,
input_offsets, thread_start, thread_end, thread_dim, stack_acc_buffer,
kStaticAccBufferMaxSize);
}
// This DWConv function uses static memory for accumulation by default for upto
// kStaticAccBufferMaxSize channels. Beyound that, a dynamic buffer is used on
// a per call basis. The function errors out if number of channels is larger
// than kStaticAccBufferMaxSize and TF_LITE_STATIC_MEMORY is defined.
inline void DepthwiseConvHybridGeneral(
const DepthwiseParams& params, const float* input_scales,
const RuntimeShape& input_shape, const int8* input_data,
const RuntimeShape& filter_shape, const int8* filter_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data,
const float* per_channel_scales, const int32_t* input_offsets,
int thread_start, int thread_end, int thread_dim) {
#ifndef TF_LITE_STATIC_MEMORY
static const int kStaticAccBufferMaxSize = 2048;
const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
if (kStaticAccBufferMaxSize < output_depth) {
std::unique_ptr<int32[]> heap_acc_buffer(new int32[output_depth]);
DoDepthwiseConvHybridGeneral(
params, input_scales, input_shape, input_data, filter_shape,
filter_data, bias_shape, bias_data, output_shape, output_data,
per_channel_scales, input_offsets, thread_start, thread_end, thread_dim,
heap_acc_buffer.get(), output_depth);
return;
}
#endif
DoDepthwiseConvHybridGeneralStatic(
params, input_scales, input_shape, input_data, filter_shape, filter_data,
bias_shape, bias_data, output_shape, output_data, per_channel_scales,
input_offsets, thread_start, thread_end, thread_dim);
}
} // namespace depthwise_conv
template <DepthwiseConvOutputRounding kOutputRounding>
inline void DepthwiseConvHybridWithRounding(
const DepthwiseParams& params, const float* input_scales,
const RuntimeShape& input_shape, const int8* input_data,
const RuntimeShape& filter_shape, const int8* filter_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data,
const float* per_channel_scales, const int32_t* input_offsets,
int thread_start, int thread_end, int thread_dim) {
gemmlowp::ScopedProfilingLabel label("DepthwiseConvHybridInt8/8bit");
const int depth_multiplier = params.depth_multiplier;
const int dilation_width_factor = params.dilation_width_factor;
const int dilation_height_factor = params.dilation_height_factor;
TFLITE_DCHECK_GE(dilation_width_factor, 1);
TFLITE_DCHECK_GE(dilation_height_factor, 1);
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
const int input_depth = input_shape.Dims(3);
TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier);
TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth);
// Enable for arm64 except for the Nvidia Linux 4 Tegra (L4T) running on
// Jetson TX-2. This compiler does not support the offsetof() macro.
#if defined(__aarch64__) && !defined(GOOGLE_L4T)
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
// Call kernel optimized for depthwise convolutions using 3x3 filters if
// parameters are supported.
if (optimized_ops::depthwise_conv::Fast3x3FilterKernelSupported<
optimized_ops::depthwise_conv::QuantizationType::kNonPerChannelUint8>(
input_shape, filter_shape, stride_width, stride_height,
dilation_width_factor, dilation_height_factor, pad_width, pad_height,
depth_multiplier, output_shape, 0, nullptr)) {
gemmlowp::ScopedProfilingLabel specialized_label(
"DepthwiseConvHybridInt8/8bit/3x3");
optimized_ops::depthwise_conv::DepthwiseConvHybrid3x3FilterPerChannel<
DepthwiseConvOutputRounding::kUpward>(
params, input_scales, input_shape, input_data,
filter_shape, filter_data, bias_shape, bias_data, output_shape,
output_data, per_channel_scales, input_offsets,
thread_start, thread_end, thread_dim);
return;
}
#endif
gemmlowp::ScopedProfilingLabel specialized_label(
"DepthwiseConvHybridInt8/8bit/General");
depthwise_conv::DepthwiseConvHybridGeneral(
params, input_scales, input_shape, input_data,
filter_shape, filter_data, bias_shape, bias_data, output_shape,
output_data, per_channel_scales, input_offsets,
thread_start, thread_end, thread_dim);
}
inline void DepthwiseConvHybridImpl(
const DepthwiseParams& params, const float* input_scales,
const RuntimeShape& input_shape, const int8* input_data,
const RuntimeShape& filter_shape, const int8* filter_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data,
const float* per_channel_scales, const int32_t* input_offsets,
int thread_start, int thread_end, int thread_dim) {
return DepthwiseConvHybridWithRounding<
DepthwiseConvOutputRounding::kAwayFromZero>(
params, input_scales, input_shape, input_data,
filter_shape, filter_data, bias_shape, bias_data, output_shape,
output_data, per_channel_scales, input_offsets,
thread_start, thread_end, thread_dim);
}
template <typename T, typename TS>
struct DepthwiseConvHybridWorkerTask : cpu_backend_threadpool::Task {
DepthwiseConvHybridWorkerTask(const DepthwiseParams& params,
const float* input_scales,
const RuntimeShape& input_shape,
const T* input_data,
const RuntimeShape& filter_shape,
const T* filter_data,
const RuntimeShape& bias_shape,
const TS* bias_data,
const RuntimeShape& output_shape,
float* output_data,
const float* per_channel_scales,
const int32_t* input_offsets,
int thread_start, int thread_end,
int thread_dim)
: params(params),
input_scales(input_scales),
input_shape(input_shape),
input_data(input_data),
filter_shape(filter_shape),
filter_data(filter_data),
bias_shape(bias_shape),
bias_data(bias_data),
output_shape(output_shape),
output_data(output_data),
per_channel_scales(per_channel_scales),
input_offsets(input_offsets),
thread_start(thread_start),
thread_end(thread_end),
thread_dim(thread_dim) {}
void Run() override {
DepthwiseConvHybridImpl(params, input_scales, input_shape,
input_data, filter_shape, filter_data,
bias_shape, bias_data, output_shape,
output_data, per_channel_scales, input_offsets,
thread_start, thread_end, thread_dim);
}
private:
const DepthwiseParams& params;
const float* input_scales;
const RuntimeShape& input_shape;
const T* input_data;
const RuntimeShape& filter_shape;
const T* filter_data;
const RuntimeShape& bias_shape;
const TS* bias_data;
const RuntimeShape& output_shape;
float* output_data;
const float* per_channel_scales;
const int32_t* input_offsets;
int thread_start;
int thread_end;
int thread_dim;
};
inline void DepthwiseConvHybridPerChannel(
const DepthwiseParams& params, const float* input_scales,
const RuntimeShape& input_shape, const int8* input_data,
const RuntimeShape& filter_shape, const int8* filter_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data,
const float* per_channel_scales, int32_t* input_offsets,
CpuBackendContext* cpu_backend_context) {
gemmlowp::ScopedProfilingLabel label("DepthwiseConvHybridInt8");
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const int output_batches = output_shape.Dims(0);
const int output_rows = output_shape.Dims(1);
int thread_count_batch = HowManyConvThreads(output_shape, filter_shape, 0);
int thread_count_row = HowManyConvThreads(output_shape, filter_shape, 1);
int thread_dim, thread_count, thread_dim_size;
if (thread_count_batch > thread_count_row) {
thread_dim = 0;
thread_dim_size = output_batches;
thread_count = thread_count_batch;
} else {
thread_dim = 1;
thread_dim_size = output_rows;
thread_count = thread_count_row;
}
const int max_threads = cpu_backend_context->max_num_threads();
thread_count = std::max(1, std::min(thread_count, max_threads));
if (thread_count == 1) {
DepthwiseConvHybridImpl(params, input_scales, input_shape,
input_data, filter_shape, filter_data, bias_shape,
bias_data, output_shape, output_data,
per_channel_scales, input_offsets,
/*thread_start=*/0, /*thread_end=*/output_rows,
/*thread_dim=*/1);
} else {
std::vector<DepthwiseConvHybridWorkerTask<int8, float>> tasks;
// TODO(b/131746020) don't create new heap allocations every time.
// At least we make it a single heap allocation by using reserve().
tasks.reserve(thread_count);
int thread_start = 0;
for (int i = 0; i < thread_count; ++i) {
int thread_end =
thread_start + (thread_dim_size - thread_start) / (thread_count - i);
tasks.emplace_back(params, input_scales, input_shape,
input_data, filter_shape, filter_data, bias_shape,
bias_data, output_shape, output_data,
per_channel_scales, input_offsets, thread_start,
thread_end, thread_dim);
thread_start = thread_end;
}
cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
cpu_backend_context);
}
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h
+3243
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/integer_ops/fully_connected.h
+172
View File
@@ -0,0 +1,172 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_FULLY_CONNECTED_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_FULLY_CONNECTED_H_
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/cpu_backend_gemm.h"
#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/reference/integer_ops/fully_connected.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
template <typename InputScalar, typename DstScalar>
inline void FullyConnectedPerChannel(
const FullyConnectedParams& params, const int32_t* output_multiplier,
const int* output_shift, const RuntimeShape& input_shape,
const InputScalar* input_data, const RuntimeShape& filter_shape,
const int8_t* filter_data, const RuntimeShape& bias_shape,
const int32_t* bias_data, const RuntimeShape& output_shape,
DstScalar* output_data, CpuBackendContext* cpu_backend_context) {
ruy::profiler::ScopeLabel label("FullyConnectedInt8/8bit");
const int32_t input_offset = params.input_offset;
const int32_t output_offset = params.output_offset;
const int32_t output_activation_min = params.quantized_activation_min;
const int32_t output_activation_max = params.quantized_activation_max;
TFLITE_DCHECK_GE(filter_shape.DimensionsCount(), 2);
TFLITE_DCHECK_GE(output_shape.DimensionsCount(), 1);
// TODO(b/62193649): This really should be:
// const int batches = ArraySize(output_dims, 1);
// but the current --variable_batch hack consists in overwriting the 3rd
// dimension with the runtime batch size, as we don't keep track for each
// array of which dimension is the batch dimension in it.
const int output_dim_count = output_shape.DimensionsCount();
const int filter_dim_count = filter_shape.DimensionsCount();
const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
const int filter_rows = filter_shape.Dims(filter_dim_count - 2);
const int filter_cols = filter_shape.Dims(filter_dim_count - 1);
TFLITE_DCHECK_EQ(filter_shape.FlatSize(), filter_rows * filter_cols);
const int output_rows = output_shape.Dims(output_dim_count - 1);
TFLITE_DCHECK_EQ(output_rows, filter_rows);
if (bias_data) {
TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_rows);
}
const bool use_caching =
(cpu_backend_context != nullptr) && cpu_backend_context->use_caching();
cpu_backend_gemm::MatrixParams<int8_t> lhs_params;
lhs_params.rows = filter_rows;
lhs_params.cols = filter_cols;
lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
lhs_params.zero_point = 0;
lhs_params.cache_policy =
use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.lhs_cacheable)
: cpu_backend_gemm::CachePolicy::kNeverCache;
cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
rhs_params.rows = filter_cols;
rhs_params.cols = batches;
rhs_params.order = cpu_backend_gemm::Order::kColMajor;
rhs_params.zero_point = -input_offset;
rhs_params.cache_policy =
use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.rhs_cacheable)
: cpu_backend_gemm::CachePolicy::kNeverCache;
cpu_backend_gemm::MatrixParams<DstScalar> dst_params;
dst_params.rows = filter_rows;
dst_params.cols = batches;
dst_params.order = cpu_backend_gemm::Order::kColMajor;
dst_params.zero_point = output_offset;
cpu_backend_gemm::GemmParams<
int32_t, DstScalar,
cpu_backend_gemm::QuantizationFlavor::kIntegerWithPerRowMultiplier>
gemm_params;
gemm_params.bias = bias_data;
gemm_params.clamp_min = output_activation_min;
gemm_params.clamp_max = output_activation_max;
gemm_params.multiplier_fixedpoint_perchannel = output_multiplier;
gemm_params.multiplier_exponent_perchannel = output_shift;
cpu_backend_gemm::Gemm(lhs_params, filter_data, rhs_params, input_data,
dst_params, output_data, gemm_params,
cpu_backend_context);
}
template <typename InputScalar, typename DstScalar>
inline void FullyConnected(
const FullyConnectedParams& params, const RuntimeShape& input_shape,
const InputScalar* input_data, const RuntimeShape& filter_shape,
const int8_t* filter_data, const RuntimeShape& bias_shape,
const int32_t* bias_data, const RuntimeShape& output_shape,
DstScalar* output_data, CpuBackendContext* cpu_backend_context) {
ruy::profiler::ScopeLabel label("FullyConnectedInt8/8bit");
const int32_t input_offset = params.input_offset;
const int32_t filter_offset = params.weights_offset;
const int32_t output_offset = params.output_offset;
const int32_t output_multiplier = params.output_multiplier;
const int output_shift = params.output_shift;
const int32_t output_activation_min = params.quantized_activation_min;
const int32_t output_activation_max = params.quantized_activation_max;
TFLITE_DCHECK_GE(filter_shape.DimensionsCount(), 2);
TFLITE_DCHECK_GE(output_shape.DimensionsCount(), 1);
// TODO(b/62193649): This really should be:
// const int batches = ArraySize(output_dims, 1);
// but the current --variable_batch hack consists in overwriting the 3rd
// dimension with the runtime batch size, as we don't keep track for each
// array of which dimension is the batch dimension in it.
const int output_dim_count = output_shape.DimensionsCount();
const int filter_dim_count = filter_shape.DimensionsCount();
const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
const int filter_rows = filter_shape.Dims(filter_dim_count - 2);
const int filter_cols = filter_shape.Dims(filter_dim_count - 1);
TFLITE_DCHECK_EQ(filter_shape.FlatSize(), filter_rows * filter_cols);
const int output_rows = output_shape.Dims(output_dim_count - 1);
TFLITE_DCHECK_EQ(output_rows, filter_rows);
if (bias_data) {
TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_rows);
}
const bool use_caching =
(cpu_backend_context != nullptr) && cpu_backend_context->use_caching();
cpu_backend_gemm::MatrixParams<int8_t> lhs_params;
lhs_params.rows = filter_rows;
lhs_params.cols = filter_cols;
lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
lhs_params.zero_point = -filter_offset;
lhs_params.cache_policy =
use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.lhs_cacheable)
: cpu_backend_gemm::CachePolicy::kNeverCache;
cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
rhs_params.rows = filter_cols;
rhs_params.cols = batches;
rhs_params.order = cpu_backend_gemm::Order::kColMajor;
rhs_params.zero_point = -input_offset;
rhs_params.cache_policy =
use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.rhs_cacheable)
: cpu_backend_gemm::CachePolicy::kNeverCache;
cpu_backend_gemm::MatrixParams<DstScalar> dst_params;
dst_params.rows = filter_rows;
dst_params.cols = batches;
dst_params.order = cpu_backend_gemm::Order::kColMajor;
dst_params.zero_point = output_offset;
cpu_backend_gemm::GemmParams<int32_t, DstScalar> gemm_params;
gemm_params.bias = bias_data;
gemm_params.clamp_min = output_activation_min;
gemm_params.clamp_max = output_activation_max;
gemm_params.multiplier_fixedpoint = output_multiplier;
gemm_params.multiplier_exponent = output_shift;
cpu_backend_gemm::Gemm(lhs_params, filter_data, rhs_params, input_data,
dst_params, output_data, gemm_params,
cpu_backend_context);
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_FULLY_CONNECTED_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/leaky_relu.h
+113
View File
@@ -0,0 +1,113 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LEAKY_RELU_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LEAKY_RELU_H_
#include <algorithm>
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
inline void QuantizeLeakyRelu(const LeakyReluParams& params,
const RuntimeShape& input_shape,
const int16* input_data,
const RuntimeShape& output_shape,
int16* output_data) {
const int flat_size = MatchingFlatSize(input_shape, output_shape);
const int32_t quantized_min = std::numeric_limits<int16>::min();
const int32_t quantized_max = std::numeric_limits<int16>::max();
int i = 0;
#ifdef __AVX2__
const __m256i input_offset = _mm256_set1_epi32(params.input_offset);
const __m256i output_offset = _mm256_set1_epi32(params.output_offset);
const __m256i output_muliplier_identity =
_mm256_set1_epi32(params.output_multiplier_identity);
const __m256i output_shift_identity =
_mm256_set1_epi32(params.output_shift_identity);
const __m256i output_multiplier_alpha =
_mm256_set1_epi32(params.output_multiplier_alpha);
const __m256i output_shift_alpha =
_mm256_set1_epi32(params.output_shift_alpha);
const __m256i clamp_max_v = _mm256_set1_epi32(quantized_max);
const __m256i clamp_min_v = _mm256_set1_epi32(quantized_min);
for (; i <= flat_size - 16; i += 16) {
const __m256i input =
_mm256_loadu_si256(reinterpret_cast<__m256i const*>(input_data + i));
__m256i input_low = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(input));
__m256i input_high =
_mm256_cvtepi16_epi32(_mm256_extracti128_si256(input, 1));
input_low = _mm256_sub_epi32(input_low, input_offset);
input_high = _mm256_sub_epi32(input_high, input_offset);
const __m256i zeros = _mm256_setzero_si256();
const __m256i input_low_mask = _mm256_cmpgt_epi32(input_low, zeros);
const __m256i input_high_mask = _mm256_cmpgt_epi32(input_high, zeros);
const __m256i input_low_output_multiplier = avx2_utils::mm256_blendv_epi32(
output_multiplier_alpha, output_muliplier_identity, input_low_mask);
const __m256i input_low_output_shift = avx2_utils::mm256_blendv_epi32(
output_shift_alpha, output_shift_identity, input_low_mask);
const __m256i input_high_output_multiplier = avx2_utils::mm256_blendv_epi32(
output_multiplier_alpha, output_muliplier_identity, input_high_mask);
const __m256i input_high_output_shift = avx2_utils::mm256_blendv_epi32(
output_shift_alpha, output_shift_identity, input_high_mask);
input_low = avx2_utils::MultiplyByQuantizedMultiplier(
input_low, input_low_output_multiplier, input_low_output_shift);
input_high = avx2_utils::MultiplyByQuantizedMultiplier(
input_high, input_high_output_multiplier, input_high_output_shift);
input_low = _mm256_add_epi32(input_low, output_offset);
input_high = _mm256_add_epi32(input_high, output_offset);
input_low = _mm256_min_epi32(input_low, clamp_max_v);
input_low = _mm256_max_epi32(input_low, clamp_min_v);
input_high = _mm256_min_epi32(input_high, clamp_max_v);
input_high = _mm256_max_epi32(input_high, clamp_min_v);
avx2_utils::CastInt32ToInt16AndStore(output_data + i, input_low);
avx2_utils::CastInt32ToInt16AndStore(output_data + i + 8, input_high);
}
#endif // __AVX2__
for (; i < flat_size; ++i) {
const int32_t input_value = input_data[i] - params.input_offset;
int32_t unclamped_output;
if (input_value >= 0) {
unclamped_output = params.output_offset +
MultiplyByQuantizedMultiplier(
input_value, params.output_multiplier_identity,
params.output_shift_identity);
} else {
unclamped_output = params.output_offset +
MultiplyByQuantizedMultiplier(
input_value, params.output_multiplier_alpha,
params.output_shift_alpha);
}
const int16 clamped_output =
std::min(quantized_max, std::max(quantized_min, unclamped_output));
output_data[i] = static_cast<int16>(clamped_output);
}
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LEAKY_RELU_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/lut.h
+72
View File
@@ -0,0 +1,72 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LUT_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LUT_H_
#include <cstdint>
#if __aarch64__ && __clang__
#include <arm_neon.h>
#endif
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
namespace tflite {
namespace optimized_integer_ops {
inline void LookupTable(const uint8_t* input_data, int num_elements,
const uint8_t* lut, uint8_t* output_data) {
int i = 0;
#if __aarch64__ && __clang__
// This code uses ARM64-only instructions.
// TODO(b/143709993): Port to ARMv7
// Load the tables into registers. (4*4 128-bit registers)
uint8x16x4_t table[4];
table[0] = vld1q_u8_x4(lut + 16 * 4 * 0);
table[1] = vld1q_u8_x4(lut + 16 * 4 * 1);
table[2] = vld1q_u8_x4(lut + 16 * 4 * 2);
table[3] = vld1q_u8_x4(lut + 16 * 4 * 3);
// Vectorized loop; process uint8x16_t (16 elements) at a time.
constexpr int vectorized_16_loop_step = 16;
const int vectorized_16_loop_end =
num_elements / vectorized_16_loop_step * vectorized_16_loop_step;
for (; i < vectorized_16_loop_end; i += vectorized_16_loop_step) {
uint8x16_t input = vld1q_u8(input_data + i);
uint8x16_t output = optimized_ops::aarch64_lookup_vector(table, input);
vst1q_u8(output_data + i, output);
}
// Postamble and non-ARM64 code: simple for loop.
#endif
for (; i < num_elements; ++i) {
output_data[i] = lut[input_data[i]];
}
}
// LUTPopulate<int8_t> has ordered the LUT so that indexing it with an
// int8_t is just done by casting it to an uint8_t. We can thus reuse the uint8
// LookupTable function.
inline void LookupTable(const int8_t* input_data, int num_elements,
const int8_t* lut, int8_t* output_data) {
LookupTable(reinterpret_cast<const uint8_t*>(input_data), num_elements,
reinterpret_cast<const uint8_t*>(lut),
reinterpret_cast<uint8_t*>(output_data));
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LUT_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/mean.h
+251
View File
@@ -0,0 +1,251 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MEAN_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MEAN_H_
#include <algorithm>
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
namespace tflite {
namespace optimized_integer_ops {
inline void MeanImpl(const tflite::MeanParams& op_params,
const RuntimeShape& input_shape, const int8_t* input_data,
int32 multiplier, int32 shift, int32 bias,
const RuntimeShape& output_shape, int8_t* output_data,
int start_depth, int end_depth) {
ruy::profiler::ScopeLabel label("Mean4D/Int8/MeanImpl");
// Current implementation only supports dimension equals 4 and simultaneous
// reduction over width and height.
const int output_batch = output_shape.Dims(0);
const int output_height = output_shape.Dims(2);
const int output_width = output_shape.Dims(2);
const int input_height = input_shape.Dims(1);
const int input_width = input_shape.Dims(2);
TFLITE_CHECK_EQ(op_params.axis_count, 2);
TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
(op_params.axis[0] == 2 && op_params.axis[1] == 1));
TFLITE_CHECK_EQ(output_height, 1);
TFLITE_CHECK_EQ(output_width, 1);
constexpr static int32_t kMinValue = std::numeric_limits<int8_t>::min();
constexpr static int32_t kMaxValue = std::numeric_limits<int8_t>::max();
#ifdef USE_NEON
const int32x4_t bias_dup = vdupq_n_s32(bias);
const int32x4_t min_dup = vdupq_n_s32(kMinValue);
const int32x4_t max_dup = vdupq_n_s32(kMaxValue);
#endif // USE_NEON
for (int out_b = 0; out_b < output_batch; ++out_b) {
int out_d = start_depth;
#ifdef USE_NEON
for (; out_d <= end_depth - 16; out_d += 16) {
int32x4x4_t temp_sum;
temp_sum.val[0] = vdupq_n_s32(0);
temp_sum.val[1] = vdupq_n_s32(0);
temp_sum.val[2] = vdupq_n_s32(0);
temp_sum.val[3] = vdupq_n_s32(0);
for (int in_h = 0; in_h < input_height; ++in_h) {
for (int in_w = 0; in_w < input_width; ++in_w) {
const int8_t* input_data_ptr =
input_data + Offset(input_shape, out_b, in_h, in_w, out_d);
int8x16_t input_data_val = vld1q_s8(input_data_ptr);
int16x8_t input_data_low_shift =
vmovl_s8(vget_low_s8(input_data_val));
int16x8_t input_data_high_shift =
vmovl_s8(vget_high_s8(input_data_val));
int32x4_t input_low_low =
vmovl_s16(vget_low_s16(input_data_low_shift));
int32x4_t input_high_low =
vmovl_s16(vget_high_s16(input_data_low_shift));
int32x4_t input_low_high =
vmovl_s16(vget_low_s16(input_data_high_shift));
int32x4_t input_high_high =
vmovl_s16(vget_high_s16(input_data_high_shift));
temp_sum.val[0] = vaddq_s32(temp_sum.val[0], input_low_low);
temp_sum.val[1] = vaddq_s32(temp_sum.val[1], input_high_low);
temp_sum.val[2] = vaddq_s32(temp_sum.val[2], input_low_high);
temp_sum.val[3] = vaddq_s32(temp_sum.val[3], input_high_high);
}
}
temp_sum =
MultiplyByQuantizedMultiplier4Rows(temp_sum, multiplier, shift);
temp_sum.val[0] = vaddq_s32(temp_sum.val[0], bias_dup);
temp_sum.val[1] = vaddq_s32(temp_sum.val[1], bias_dup);
temp_sum.val[2] = vaddq_s32(temp_sum.val[2], bias_dup);
temp_sum.val[3] = vaddq_s32(temp_sum.val[3], bias_dup);
temp_sum.val[0] = vminq_s32(vmaxq_s32(temp_sum.val[0], min_dup), max_dup);
temp_sum.val[1] = vminq_s32(vmaxq_s32(temp_sum.val[1], min_dup), max_dup);
temp_sum.val[2] = vminq_s32(vmaxq_s32(temp_sum.val[2], min_dup), max_dup);
temp_sum.val[3] = vminq_s32(vmaxq_s32(temp_sum.val[3], min_dup), max_dup);
int16x4_t narrowed_low_low = vmovn_s32(temp_sum.val[0]);
int16x4_t narrowed_high_low = vmovn_s32(temp_sum.val[1]);
int16x4_t narrowed_low_high = vmovn_s32(temp_sum.val[2]);
int16x4_t narrowed_high_high = vmovn_s32(temp_sum.val[3]);
int16x8_t combined_low =
vcombine_s16(narrowed_low_low, narrowed_high_low);
int16x8_t combined_high =
vcombine_s16(narrowed_low_high, narrowed_high_high);
int8x8_t narrowed_low = vmovn_s16(combined_low);
int8x8_t narrowed_high = vmovn_s16(combined_high);
int8x16_t combined_output = vcombine_s8(narrowed_low, narrowed_high);
int8_t* output_data_ptr =
output_data + Offset(output_shape, out_b, 0, 0, out_d);
vst1q_s8(output_data_ptr, combined_output);
}
#endif // USE_NEON
for (; out_d < end_depth; ++out_d) {
int acc = 0;
for (int in_h = 0; in_h < input_height; ++in_h) {
for (int in_w = 0; in_w < input_width; ++in_w) {
acc += input_data[Offset(input_shape, out_b, in_h, in_w, out_d)];
}
}
acc = MultiplyByQuantizedMultiplier(acc, multiplier, shift);
acc += bias;
acc = std::min(std::max(acc, kMinValue), kMaxValue);
output_data[Offset(output_shape, out_b, 0, 0, out_d)] =
static_cast<int8_t>(acc);
}
}
}
struct MeanWorkerTask : cpu_backend_threadpool::Task {
MeanWorkerTask(const tflite::MeanParams& op_params,
const RuntimeShape& input_shape, const int8_t* input_data,
int32 multiplier, int32 shift, int32 bias,
const RuntimeShape& output_shape, int8_t* output_data,
int start_height, int end_height)
: op_params(op_params),
input_shape(input_shape),
input_data(input_data),
multiplier(multiplier),
shift(shift),
bias(bias),
output_shape(output_shape),
output_data(output_data),
start_height(start_height),
end_height(end_height) {}
void Run() override {
MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
output_shape, output_data, start_height, end_height);
}
private:
const tflite::MeanParams& op_params;
const RuntimeShape& input_shape;
const int8_t* input_data;
int32 multiplier;
int32 shift;
int32 bias;
const RuntimeShape& output_shape;
int8_t* output_data;
int start_height;
int end_height;
};
inline void Mean(const tflite::MeanParams& op_params,
const RuntimeShape& unextended_input_shape,
const int8_t* input_data, int32 input_zero_point,
float input_scale, const RuntimeShape& unextended_output_shape,
int8_t* output_data, int32 output_zero_point,
float output_scale, CpuBackendContext* cpu_backend_context) {
ruy::profiler::ScopeLabel label("Mean4D/Int8");
// Current implementation only supports dimension equals 4 and simultaneous
// reduction over width and height.
TFLITE_CHECK_EQ(unextended_input_shape.DimensionsCount(), 4);
TFLITE_CHECK_LE(unextended_output_shape.DimensionsCount(), 4);
const RuntimeShape input_shape =
RuntimeShape::ExtendedShape(4, unextended_input_shape);
const RuntimeShape output_shape =
RuntimeShape::ExtendedShape(4, unextended_output_shape);
const int output_height = output_shape.Dims(1);
const int output_width = output_shape.Dims(2);
const int output_depth = output_shape.Dims(3);
TFLITE_CHECK_EQ(op_params.axis_count, 2);
TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
(op_params.axis[0] == 2 && op_params.axis[1] == 1));
TFLITE_CHECK_EQ(output_height, 1);
TFLITE_CHECK_EQ(output_width, 1);
const int input_height = input_shape.Dims(1);
const int input_width = input_shape.Dims(2);
const float num_elements_in_axis = input_width * input_height;
float temp = input_zero_point * input_scale / output_scale;
temp = temp > 0 ? temp + 0.5f : temp - 0.5f;
int32_t bias = output_zero_point - static_cast<int32_t>(temp);
float real_scale = input_scale / (num_elements_in_axis * output_scale);
int32 multiplier, shift;
QuantizeMultiplier(real_scale, &multiplier, &shift);
constexpr int kMinDepthPerThread = 8;
int thread_count = output_depth / kMinDepthPerThread;
thread_count = thread_count > 0 ? thread_count : 1;
const int capped_thread_count =
std::min(thread_count, cpu_backend_context->max_num_threads());
if (capped_thread_count == 1) {
MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
output_shape, output_data, 0, output_depth);
} else {
// Instead parallel for batch, we loop for the output_depth since batch
// is typical 1.
std::vector<MeanWorkerTask> tasks;
// TODO(b/131746020) don't create new heap allocations every time.
// At least we make it a single heap allocation by using reserve().
tasks.reserve(capped_thread_count);
int depth_start = 0;
for (int i = 0; i < capped_thread_count; ++i) {
// Try to distribute the tasks as even as possible.
int depth_end = depth_start +
(output_depth - depth_start) / (capped_thread_count - i);
tasks.emplace_back(op_params, input_shape, input_data, multiplier, shift,
bias, output_shape, output_data, depth_start,
depth_end);
depth_start = depth_end;
}
cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
cpu_backend_context);
}
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MEAN_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/mul.h
+266
View File
@@ -0,0 +1,266 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_
#include <algorithm>
#include "fixedpoint/fixedpoint.h"
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/lite/kernels/internal/reference/integer_ops/mul.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
// Element-wise mul that can often be used for inner loop of broadcast Mul as
// well as the non-broadcast Mul.
inline void MulElementwise(int size, const ArithmeticParams& params,
const int8* input1_data, const int8* input2_data,
int8* output_data) {
ruy::profiler::ScopeLabel label("MulElementwiseInt8/8bit");
int i = 0;
TFLITE_DCHECK_GT(params.input1_offset, -256);
TFLITE_DCHECK_LT(params.input1_offset, 256);
TFLITE_DCHECK_GT(params.input2_offset, -256);
TFLITE_DCHECK_LT(params.input2_offset, 256);
TFLITE_DCHECK_GT(params.output_offset, -256);
TFLITE_DCHECK_LT(params.output_offset, 256);
#ifdef USE_NEON
const int16x8_t input1_offset_vector = vdupq_n_s16(params.input1_offset);
const int16x8_t input2_offset_vector = vdupq_n_s16(params.input2_offset);
const int16x8_t output_offset_vector = vdupq_n_s16(params.output_offset);
const auto output_activation_min_vector =
vdupq_n_s8(params.quantized_activation_min);
const auto output_activation_max_vector =
vdupq_n_s8(params.quantized_activation_max);
const int left_shift = std::max(0, params.output_shift);
const int right_shift = std::max(0, -params.output_shift);
const int32x4_t left_shift_vec = vdupq_n_s32(left_shift);
for (; i <= size - 16; i += 16) {
// We load / store 16 at a time, multiplying as four sets of 4 int32s.
const int8x16_t input1_val_original = vld1q_s8(input1_data + i);
const int8x16_t input2_val_original = vld1q_s8(input2_data + i);
const int16x8_t input1_val_s16_high =
vmovl_s8(vget_high_s8(input1_val_original));
const int16x8_t input1_val_s16_low =
vmovl_s8(vget_low_s8(input1_val_original));
const int16x8_t input2_val_s16_high =
vmovl_s8(vget_high_s8(input2_val_original));
const int16x8_t input2_val_s16_low =
vmovl_s8(vget_low_s8(input2_val_original));
const int16x8_t input1_val_high =
vaddq_s16(input1_val_s16_high, input1_offset_vector);
const int16x8_t input2_val_high =
vaddq_s16(input2_val_s16_high, input2_offset_vector);
const int16x8_t input1_val_low =
vaddq_s16(input1_val_s16_low, input1_offset_vector);
const int16x8_t input2_val_low =
vaddq_s16(input2_val_s16_low, input2_offset_vector);
const int16x4_t input1_val_high_high = vget_high_s16(input1_val_high);
const int16x4_t input1_val_high_low = vget_low_s16(input1_val_high);
const int16x4_t input1_val_low_high = vget_high_s16(input1_val_low);
const int16x4_t input1_val_low_low = vget_low_s16(input1_val_low);
const int16x4_t input2_val_high_high = vget_high_s16(input2_val_high);
const int16x4_t input2_val_high_low = vget_low_s16(input2_val_high);
const int16x4_t input2_val_low_high = vget_high_s16(input2_val_low);
const int16x4_t input2_val_low_low = vget_low_s16(input2_val_low);
auto p1 = vmull_s16(input2_val_high_high, input1_val_high_high);
auto p2 = vmull_s16(input2_val_high_low, input1_val_high_low);
auto p3 = vmull_s16(input2_val_low_high, input1_val_low_high);
auto p4 = vmull_s16(input2_val_low_low, input1_val_low_low);
p1 = vshlq_s32(p1, left_shift_vec);
p2 = vshlq_s32(p2, left_shift_vec);
p3 = vshlq_s32(p3, left_shift_vec);
p4 = vshlq_s32(p4, left_shift_vec);
p1 = vqrdmulhq_n_s32(p1, params.output_multiplier);
p2 = vqrdmulhq_n_s32(p2, params.output_multiplier);
p3 = vqrdmulhq_n_s32(p3, params.output_multiplier);
p4 = vqrdmulhq_n_s32(p4, params.output_multiplier);
using gemmlowp::RoundingDivideByPOT;
p1 = RoundingDivideByPOT(p1, right_shift);
p2 = RoundingDivideByPOT(p2, right_shift);
p3 = RoundingDivideByPOT(p3, right_shift);
p4 = RoundingDivideByPOT(p4, right_shift);
const auto p1_narrowed = vqmovn_s32(p1);
const auto p2_narrowed = vqmovn_s32(p2);
const auto p3_narrowed = vqmovn_s32(p3);
const auto p4_narrowed = vqmovn_s32(p4);
const int16x8_t p_part1 =
vaddq_s16(vcombine_s16(p2_narrowed, p1_narrowed), output_offset_vector);
const int16x8_t p_part2 =
vaddq_s16(vcombine_s16(p4_narrowed, p3_narrowed), output_offset_vector);
const int8x16_t p = vcombine_s8(vqmovn_s16(p_part2), vqmovn_s16(p_part1));
const auto clamped = vmaxq_s8(output_activation_min_vector,
vminq_s8(output_activation_max_vector, p));
vst1q_s8(output_data + i, clamped);
}
#endif // NEON
for (; i < size; ++i) {
const int32 input1_val = params.input1_offset + input1_data[i];
const int32 input2_val = params.input2_offset + input2_data[i];
const int32 unclamped_result =
params.output_offset +
MultiplyByQuantizedMultiplier(input1_val * input2_val,
params.output_multiplier,
params.output_shift);
const int32 clamped_output =
std::min(params.quantized_activation_max,
std::max(params.quantized_activation_min, unclamped_result));
output_data[i] = static_cast<int8>(clamped_output);
}
}
// Broadcast mul that can often be used for inner loop of broadcast Mul.
inline void MulSimpleBroadcast(int size, const ArithmeticParams& params,
const int8 broadcast_value,
const int8* input2_data, int8* output_data) {
ruy::profiler::ScopeLabel label("BroadMulSimpleBroadcastInt8/8bit");
const int16 input1_val = params.input1_offset + broadcast_value;
int i = 0;
TFLITE_DCHECK_GT(params.input1_offset, -256);
TFLITE_DCHECK_LT(params.input1_offset, 256);
TFLITE_DCHECK_GT(params.input2_offset, -256);
TFLITE_DCHECK_LT(params.input2_offset, 256);
TFLITE_DCHECK_GT(params.output_offset, -256);
TFLITE_DCHECK_LT(params.output_offset, 256);
#ifdef USE_NEON
const auto input2_offset_vector = vdupq_n_s16(params.input2_offset);
const auto output_offset_vector = vdupq_n_s16(params.output_offset);
const auto output_activation_min_vector =
vdupq_n_s8(params.quantized_activation_min);
const auto output_activation_max_vector =
vdupq_n_s8(params.quantized_activation_max);
const int left_shift = std::max(0, params.output_shift);
const int right_shift = std::max(0, -params.output_shift);
const int32x4_t left_shift_vec = vdupq_n_s32(left_shift);
for (; i <= size - 16; i += 16) {
// We load / store 16 at a time, multiplying as four sets of 4 int32s.
const auto input2_val_original = vld1q_s8(input2_data + i);
const auto input2_val_s16_high =
vmovl_s8(vget_high_s8(input2_val_original));
const auto input2_val_s16_low = vmovl_s8(vget_low_s8(input2_val_original));
const auto input2_val_high =
vaddq_s16(input2_val_s16_high, input2_offset_vector);
const auto input2_val_low =
vaddq_s16(input2_val_s16_low, input2_offset_vector);
const auto input2_val_low_low = vget_low_s16(input2_val_low);
const auto input2_val_low_high = vget_high_s16(input2_val_low);
const auto input2_val_high_low = vget_low_s16(input2_val_high);
const auto input2_val_high_high = vget_high_s16(input2_val_high);
auto p1 = vmull_n_s16(input2_val_high_high, input1_val);
auto p2 = vmull_n_s16(input2_val_high_low, input1_val);
auto p3 = vmull_n_s16(input2_val_low_high, input1_val);
auto p4 = vmull_n_s16(input2_val_low_low, input1_val);
p1 = vshlq_s32(p1, left_shift_vec);
p2 = vshlq_s32(p2, left_shift_vec);
p3 = vshlq_s32(p3, left_shift_vec);
p4 = vshlq_s32(p4, left_shift_vec);
p1 = vqrdmulhq_n_s32(p1, params.output_multiplier);
p2 = vqrdmulhq_n_s32(p2, params.output_multiplier);
p3 = vqrdmulhq_n_s32(p3, params.output_multiplier);
p4 = vqrdmulhq_n_s32(p4, params.output_multiplier);
using gemmlowp::RoundingDivideByPOT;
p1 = RoundingDivideByPOT(p1, right_shift);
p2 = RoundingDivideByPOT(p2, right_shift);
p3 = RoundingDivideByPOT(p3, right_shift);
p4 = RoundingDivideByPOT(p4, right_shift);
const auto p1_narrowed = vqmovn_s32(p1);
const auto p2_narrowed = vqmovn_s32(p2);
const auto p3_narrowed = vqmovn_s32(p3);
const auto p4_narrowed = vqmovn_s32(p4);
const int16x8_t p_part1 =
vaddq_s16(vcombine_s16(p2_narrowed, p1_narrowed), output_offset_vector);
const int16x8_t p_part2 =
vaddq_s16(vcombine_s16(p4_narrowed, p3_narrowed), output_offset_vector);
const int8x16_t p = vcombine_s8(vqmovn_s16(p_part2), vqmovn_s16(p_part1));
const auto clamped = vmaxq_s8(output_activation_min_vector,
vminq_s8(output_activation_max_vector, p));
vst1q_s8(output_data + i, clamped);
}
#endif // NEON
for (; i < size; ++i) {
const int32 input2_val = params.input2_offset + input2_data[i];
const int32 unclamped_result =
params.output_offset +
MultiplyByQuantizedMultiplier(input1_val * input2_val,
params.output_multiplier,
params.output_shift);
const int32 clamped_output =
std::min(params.quantized_activation_max,
std::max(params.quantized_activation_min, unclamped_result));
output_data[i] = static_cast<int8>(clamped_output);
}
}
inline void Mul(const ArithmeticParams& params,
const RuntimeShape& input1_shape, const int8* input1_data,
const RuntimeShape& input2_shape, const int8* input2_data,
const RuntimeShape& output_shape, int8* output_data) {
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
ruy::profiler::ScopeLabel label("MulInt8/8bit");
const int flat_size =
MatchingElementsSize(input1_shape, input2_shape, output_shape);
MulElementwise(flat_size, params, input1_data, input2_data, output_data);
}
inline void BroadcastMulDispatch(const ArithmeticParams& params,
const RuntimeShape& input1_shape,
const int8* input1_data,
const RuntimeShape& input2_shape,
const int8* input2_data,
const RuntimeShape& output_shape,
int8* output_data) {
if (params.broadcast_category == BroadcastableOpCategory::kGenericBroadcast) {
return reference_integer_ops::BroadcastMul6DSlow(
params, input1_shape, input1_data, input2_shape, input2_data,
output_shape, output_data);
}
optimized_ops::BinaryBroadcastFiveFold(
params, input1_shape, input1_data, input2_shape, input2_data,
output_shape, output_data, MulElementwise, MulSimpleBroadcast);
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h
+278
View File
@@ -0,0 +1,278 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_
#include <string.h>
#include <algorithm>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/cppmath.h"
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#include "tensorflow/lite/kernels/internal/optimized/im2col_utils.h"
#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
#include "tensorflow/lite/kernels/internal/quantization_util.h"
#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
#include "tensorflow/lite/kernels/internal/strided_slice_logic.h"
#include "tensorflow/lite/kernels/internal/tensor_utils.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
inline void MaxPool(const PoolParams& params, const RuntimeShape& input_shape,
const int8_t* input_data, const RuntimeShape& output_shape,
int8_t* output_data) {
ruy::profiler::ScopeLabel label("MaxPool/8bit");
// Here, and in other pooling ops, in order to maintain locality of reference,
// to minimize some recalculations, and to load into NEON vector registers, we
// use an inner loop down the depth. Since depths can be large and hence we
// would need arbitrarily large temporary storage, we divide the work up into
// depth tranches just within the batch loop.
static constexpr int kPoolingAccTrancheSize = 256;
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
const int depth = MatchingDim(input_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 stride_height = params.stride_height;
const int stride_width = params.stride_width;
int8_t acc[kPoolingAccTrancheSize];
for (int batch = 0; batch < batches; ++batch) {
// We proceed through the depth in tranches (see comment above). The
// depth_base is the depth at the beginning of the tranche. The
// tranche_depth is the depth dimension of the tranche.
for (int depth_base = 0; depth_base < depth;
depth_base += kPoolingAccTrancheSize) {
const int tranche_depth =
std::min(depth - depth_base, kPoolingAccTrancheSize);
for (int out_y = 0; out_y < output_height; ++out_y) {
for (int out_x = 0; out_x < output_width; ++out_x) {
const int in_x_origin =
(out_x * stride_width) - params.padding_values.width;
const int in_y_origin =
(out_y * stride_height) - params.padding_values.height;
const int filter_x_start = std::max(0, -in_x_origin);
const int filter_x_end =
std::min(params.filter_width, input_width - in_x_origin);
const int filter_y_start = std::max(0, -in_y_origin);
const int filter_y_end =
std::min(params.filter_height, input_height - in_y_origin);
memset(acc, params.quantized_activation_min,
tranche_depth * sizeof(acc[0]));
const int8_t* input_ptr =
input_data + depth_base +
depth * (in_x_origin +
input_width * (in_y_origin + input_height * batch));
for (int fy = filter_y_start; fy < filter_y_end; fy++) {
const int8_t* input_row_ptr =
input_ptr + depth * (fy * input_width + filter_x_start);
for (int fx = filter_x_start; fx < filter_x_end; fx++) {
const int8_t* input_channel_ptr = input_row_ptr;
int channel = 0;
#ifdef USE_NEON
for (; channel <= tranche_depth - 16; channel += 16) {
int8x16_t acc_reg = vld1q_s8(acc + channel);
int8x16_t input_reg = vld1q_s8(input_channel_ptr);
input_channel_ptr += 16;
acc_reg = vmaxq_s8(acc_reg, input_reg);
vst1q_s8(acc + channel, acc_reg);
}
for (; channel <= tranche_depth - 8; channel += 8) {
int8x8_t acc_reg = vld1_s8(acc + channel);
int8x8_t input_reg = vld1_s8(input_channel_ptr);
input_channel_ptr += 8;
acc_reg = vmax_s8(acc_reg, input_reg);
vst1_s8(acc + channel, acc_reg);
}
#endif
for (; channel < tranche_depth; ++channel) {
acc[channel] = std::max(acc[channel], *input_channel_ptr++);
}
input_row_ptr += depth;
}
}
int8_t* output_ptr = output_data + Offset(output_shape, batch, out_y,
out_x, depth_base);
int channel = 0;
#ifdef USE_NEON
for (; channel <= tranche_depth - 16; channel += 16) {
int8x16_t a = vld1q_s8(acc + channel);
a = vminq_s8(a, vdupq_n_s8(params.quantized_activation_max));
a = vmaxq_s8(a, vdupq_n_s8(params.quantized_activation_min));
vst1q_s8(output_ptr + channel, a);
}
for (; channel <= tranche_depth - 8; channel += 8) {
int8x8_t a = vld1_s8(acc + channel);
a = vmin_s8(a, vdup_n_s8(params.quantized_activation_max));
a = vmax_s8(a, vdup_n_s8(params.quantized_activation_min));
vst1_s8(output_ptr + channel, a);
}
#endif
for (; channel < tranche_depth; ++channel) {
int8_t a = acc[channel];
a = std::max<int8_t>(a, params.quantized_activation_min);
a = std::min<int8_t>(a, params.quantized_activation_max);
output_ptr[channel] = static_cast<int8_t>(a);
}
}
}
}
}
}
inline bool AveragePool(const PoolParams& params,
const RuntimeShape& input_shape,
const int8_t* input_data,
const RuntimeShape& output_shape, int8_t* output_data) {
ruy::profiler::ScopeLabel label("AveragePool/8bitWith32bitAccumulator");
// Here, and in other pooling ops, in order to maintain locality of reference,
// to minimize some recalculations, and to load into NEON vector registers, we
// use an inner loop down the depth. Since depths can be large and hence we
// would need arbitrarily large temporary storage, we divide the work up into
// depth tranches just within the batch loop.
static constexpr int kPoolingAccTrancheSize = 256;
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
const int depth = MatchingDim(input_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 stride_height = params.stride_height;
const int stride_width = params.stride_width;
int32_t acc[kPoolingAccTrancheSize];
for (int batch = 0; batch < batches; ++batch) {
// We proceed through the depth in tranches (see comment above). The
// depth_base is the depth at the beginning of the tranche. The
// tranche_depth is the depth dimension of the tranche.
for (int depth_base = 0; depth_base < depth;
depth_base += kPoolingAccTrancheSize) {
const int tranche_depth =
std::min(depth - depth_base, kPoolingAccTrancheSize);
for (int out_y = 0; out_y < output_height; ++out_y) {
for (int out_x = 0; out_x < output_width; ++out_x) {
const int in_x_origin =
(out_x * stride_width) - params.padding_values.width;
const int in_y_origin =
(out_y * stride_height) - params.padding_values.height;
const int filter_x_start = std::max(0, -in_x_origin);
const int filter_x_end =
std::min(params.filter_width, input_width - in_x_origin);
const int filter_y_start = std::max(0, -in_y_origin);
const int filter_y_end =
std::min(params.filter_height, input_height - in_y_origin);
const int filter_count =
(filter_x_end - filter_x_start) * (filter_y_end - filter_y_start);
if (filter_count == 0) return false;
memset(acc, 0, tranche_depth * sizeof(acc[0]));
const int8_t* input_ptr =
input_data + depth_base +
depth * (in_x_origin +
input_width * (in_y_origin + input_height * batch));
for (int fy = filter_y_start; fy < filter_y_end; fy++) {
const int8_t* input_row_ptr =
input_ptr + depth * (fy * input_width + filter_x_start);
for (int fx = filter_x_start; fx < filter_x_end; fx++) {
const int8_t* input_channel_ptr = input_row_ptr;
int channel = 0;
#ifdef USE_NEON
for (; channel <= tranche_depth - 16; channel += 16) {
int16x4_t acc_reg[4];
int8x16_t input_reg = vld1q_s8(input_channel_ptr);
input_channel_ptr += 16;
acc_reg[0] = vget_low_s16(vmovl_s8(vget_low_s8(input_reg)));
acc_reg[1] = vget_high_s16(vmovl_s8(vget_low_s8(input_reg)));
acc_reg[2] = vget_low_s16(vmovl_s8(vget_high_s8(input_reg)));
acc_reg[3] = vget_high_s16(vmovl_s8(vget_high_s8(input_reg)));
for (int i = 0; i < 4; i++) {
vst1q_s32(
acc + channel + 4 * i,
vaddw_s16(vld1q_s32(acc + channel + 4 * i), acc_reg[i]));
}
}
for (; channel <= tranche_depth - 8; channel += 8) {
int16x4_t acc_reg[2];
int16x8_t input_reg = vmovl_s8(vld1_s8(input_channel_ptr));
input_channel_ptr += 8;
acc_reg[0] = vget_low_s16(input_reg);
acc_reg[1] = vget_high_s16(input_reg);
for (int i = 0; i < 2; i++) {
vst1q_s32(
acc + channel + 4 * i,
vaddw_s16(vld1q_s32(acc + channel + 4 * i), acc_reg[i]));
}
}
#endif
for (; channel < tranche_depth; ++channel) {
acc[channel] += *input_channel_ptr++;
}
input_row_ptr += depth;
}
}
int8_t* output_ptr = output_data + Offset(output_shape, batch, out_y,
out_x, depth_base);
int channel = 0;
#ifdef USE_NEON
for (; channel <= tranche_depth - 8; channel += 8) {
int16_t buf[8];
for (int i = 0; i < 8; i++) {
buf[i] =
acc[channel + i] > 0
? (acc[channel + i] + filter_count / 2) / filter_count
: (acc[channel + i] - filter_count / 2) / filter_count;
}
int8x8_t buf8 = vqmovn_s16(vld1q_s16(buf));
buf8 = vmin_s8(buf8, vdup_n_s8(params.quantized_activation_max));
buf8 = vmax_s8(buf8, vdup_n_s8(params.quantized_activation_min));
vst1_s8(output_ptr + channel, buf8);
}
#endif
for (; channel < tranche_depth; ++channel) {
int16_t a = acc[channel] > 0
? (acc[channel] + filter_count / 2) / filter_count
: (acc[channel] - filter_count / 2) / filter_count;
a = std::max<int16_t>(a, params.quantized_activation_min);
a = std::min<int16_t>(a, params.quantized_activation_max);
output_ptr[channel] = static_cast<int8_t>(a);
}
}
}
}
}
return true;
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/sub.h
+237
View File
@@ -0,0 +1,237 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_SUB_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_SUB_H_
#include <algorithm>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
#include "tensorflow/lite/kernels/internal/reference/sub.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_integer_ops {
inline void SubElementwiseInt16(int size, const ArithmeticParams& params,
const int16* input1_data,
const int16* input2_data, int16* output_data) {
ruy::profiler::ScopeLabel label("SubElementwiseInt16/16bit");
int i = 0;
TFLITE_DCHECK_GT(params.input1_offset, -32768);
TFLITE_DCHECK_GT(params.input2_offset, -32768);
TFLITE_DCHECK_LT(params.input1_offset, 32768);
TFLITE_DCHECK_LT(params.input2_offset, 32768);
#ifdef __AVX2__
const int32_t input1_left_shift = params.left_shift + params.input1_shift;
const int32_t input2_left_shift = params.left_shift + params.input2_shift;
const __m256i input1_offset = _mm256_set1_epi32(params.input1_offset);
const __m256i input2_offset = _mm256_set1_epi32(params.input2_offset);
const __m256i output_offset = _mm256_set1_epi32(params.output_offset);
const __m256i clamp_max_v =
_mm256_set1_epi32(params.quantized_activation_max);
const __m256i clamp_min_v =
_mm256_set1_epi32(params.quantized_activation_min);
for (; i <= size - 16; i += 16) {
const __m256i input1_val_original =
_mm256_loadu_si256(reinterpret_cast<__m256i const*>(input1_data + i));
const __m256i input2_val_original =
_mm256_loadu_si256(reinterpret_cast<__m256i const*>(input2_data + i));
__m256i s11 =
_mm256_cvtepi16_epi32(_mm256_castsi256_si128(input1_val_original));
__m256i s12 =
_mm256_cvtepi16_epi32(_mm256_extracti128_si256(input1_val_original, 1));
__m256i s21 =
_mm256_cvtepi16_epi32(_mm256_castsi256_si128(input2_val_original));
__m256i s22 =
_mm256_cvtepi16_epi32(_mm256_extracti128_si256(input2_val_original, 1));
s11 = _mm256_add_epi32(s11, input1_offset);
s12 = _mm256_add_epi32(s12, input1_offset);
s21 = _mm256_add_epi32(s21, input2_offset);
s22 = _mm256_add_epi32(s22, input2_offset);
s11 = avx2_utils::MultiplyByQuantizedMultiplier(
s11, params.input1_multiplier, input1_left_shift);
s12 = avx2_utils::MultiplyByQuantizedMultiplier(
s12, params.input1_multiplier, input1_left_shift);
s21 = avx2_utils::MultiplyByQuantizedMultiplier(
s21, params.input2_multiplier, input2_left_shift);
s22 = avx2_utils::MultiplyByQuantizedMultiplier(
s22, params.input2_multiplier, input2_left_shift);
__m256i s1 = _mm256_sub_epi32(s11, s21);
__m256i s2 = _mm256_sub_epi32(s12, s22);
s1 = avx2_utils::MultiplyByQuantizedMultiplier(s1, params.output_multiplier,
params.output_shift);
s2 = avx2_utils::MultiplyByQuantizedMultiplier(s2, params.output_multiplier,
params.output_shift);
s1 = _mm256_add_epi32(s1, output_offset);
s2 = _mm256_add_epi32(s2, output_offset);
s1 = _mm256_min_epi32(s1, clamp_max_v);
s1 = _mm256_max_epi32(s1, clamp_min_v);
s2 = _mm256_min_epi32(s2, clamp_max_v);
s2 = _mm256_max_epi32(s2, clamp_min_v);
avx2_utils::CastInt32ToInt16AndStore(output_data + i, s1);
avx2_utils::CastInt32ToInt16AndStore(output_data + i + 8, s2);
}
#endif // __AVX2__
for (; i < size; ++i) {
const int32_t input1_val = params.input1_offset + input1_data[i];
const int32_t input2_val = params.input2_offset + input2_data[i];
const int32_t shifted_input1_val = input1_val * (1 << params.left_shift);
const int32_t shifted_input2_val = input2_val * (1 << params.left_shift);
const int32_t scaled_input1_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input1_val, params.input1_multiplier, params.input1_shift);
const int32_t scaled_input2_val =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
shifted_input2_val, params.input2_multiplier, params.input2_shift);
const int32_t raw_sum = scaled_input1_val - scaled_input2_val;
const int32_t raw_output =
MultiplyByQuantizedMultiplierSmallerThanOneExp(
raw_sum, params.output_multiplier, params.output_shift) +
params.output_offset;
const int32_t clamped_output =
std::min(params.quantized_activation_max,
std::max(params.quantized_activation_min, raw_output));
output_data[i] = static_cast<int16>(clamped_output);
}
}
inline void BroadcastSubFiveFold(const ArithmeticParams& unswitched_params,
const RuntimeShape& input1_shape,
const int16* unswitched_input1_data,
const RuntimeShape& input2_shape,
const int16* unswitched_input2_data,
const RuntimeShape& output_shape,
int16* output_data) {
ruy::profiler::ScopeLabel label("BroadcastSubFiveFold/16bit");
ArithmeticParams switched_params = unswitched_params;
switched_params.input1_offset = unswitched_params.input2_offset;
switched_params.input1_multiplier = unswitched_params.input2_multiplier;
switched_params.input1_shift = unswitched_params.input2_shift;
switched_params.input2_offset = unswitched_params.input1_offset;
switched_params.input2_multiplier = unswitched_params.input1_multiplier;
switched_params.input2_shift = unswitched_params.input1_shift;
const bool use_unswitched =
unswitched_params.broadcast_category ==
tflite::BroadcastableOpCategory::kFirstInputBroadcastsFast;
const ArithmeticParams& params =
use_unswitched ? unswitched_params : switched_params;
const int16_t* input1_data =
use_unswitched ? unswitched_input1_data : unswitched_input2_data;
const int16_t* input2_data =
use_unswitched ? unswitched_input2_data : unswitched_input1_data;
int16_t* output_data_ptr = output_data;
const int16_t* input1_data_ptr = input1_data;
const int16_t* input2_data_reset = input2_data;
// In the fivefold pattern, y0, y2 and y4 are not broadcast, and so shared
// between input shapes. y3 for input 1 is always broadcast, and so the
// dimension there is 1, whereas optionally y1 might be broadcast for input 2.
// The flatsize for each inputs are as below.
// input1.shape.FlatSize = y0 * y1 * y2 * y4,
// input2.shape.FlatSize = y0 * y2 * y3 * y4.
const int y0 = params.broadcast_shape[0];
const int y1 = params.broadcast_shape[1];
const int y2 = params.broadcast_shape[2];
const int y3 = params.broadcast_shape[3];
const int y4 = params.broadcast_shape[4];
for (int i0 = 0; i0 < y0; ++i0) {
const int16_t* input2_data_ptr = nullptr;
for (int i1 = 0; i1 < y1; ++i1) {
input2_data_ptr = input2_data_reset;
for (int i2 = 0; i2 < y2; ++i2) {
for (int i3 = 0; i3 < y3; ++i3) {
if (use_unswitched) {
SubElementwiseInt16(y4, params, input1_data_ptr, input2_data_ptr,
output_data_ptr);
} else {
// When input1 and input2 are switched, calculate (input2 - input1)
// and use unswitched_params as we switch the switched input here.
SubElementwiseInt16(y4, unswitched_params, input2_data_ptr,
input1_data_ptr, output_data_ptr);
}
input2_data_ptr += y4;
output_data_ptr += y4;
}
// We have broadcast y4 of input1 data y3 times, and now move on.
input1_data_ptr += y4;
}
}
// We have broadcast y2*y3*y4 of input2 data y1 times, and now move on.
input2_data_reset = input2_data_ptr;
}
}
inline void Sub(const ArithmeticParams& params,
const RuntimeShape& input1_shape, const int16* input1_data,
const RuntimeShape& input2_shape, const int16* input2_data,
const RuntimeShape& output_shape, int16* output_data) {
ruy::profiler::ScopeLabel label("SubInt16/16bit");
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
TFLITE_DCHECK_GT(params.input1_offset, -32768);
TFLITE_DCHECK_GT(params.input2_offset, -32768);
TFLITE_DCHECK_LT(params.input1_offset, 32768);
TFLITE_DCHECK_LT(params.input2_offset, 32768);
const int flat_size =
MatchingElementsSize(input1_shape, input2_shape, output_shape);
SubElementwiseInt16(flat_size, params, input1_data, input2_data, output_data);
}
inline void BroadcastSubDispatch(const ArithmeticParams& params,
const RuntimeShape& input1_shape,
const int16* input1_data,
const RuntimeShape& input2_shape,
const int16* input2_data,
const RuntimeShape& output_shape,
int16* output_data) {
ruy::profiler::ScopeLabel label("BroadcastSubDispatchInt16/16bit");
TFLITE_DCHECK_LE(params.quantized_activation_min,
params.quantized_activation_max);
TFLITE_DCHECK_GT(params.input1_offset, -32768);
TFLITE_DCHECK_GT(params.input2_offset, -32768);
TFLITE_DCHECK_LT(params.input1_offset, 32768);
TFLITE_DCHECK_LT(params.input2_offset, 32768);
if (params.broadcast_category == BroadcastableOpCategory::kGenericBroadcast) {
return reference_ops::BroadcastQuantSubSlow(
params, input1_shape, input1_data, input2_shape, input2_data,
output_shape, output_data);
}
BroadcastSubFiveFold(params, input1_shape, input1_data, input2_shape,
input2_data, output_shape, output_data);
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_SUB_H_
tensorflow/lite/kernels/internal/optimized/integer_ops/transpose_conv.h
+118
View File
@@ -0,0 +1,118 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_TRANSPOSE_CONV_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_TRANSPOSE_CONV_H_
#include <algorithm>
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
namespace tflite {
namespace optimized_integer_ops {
// TransposeConvV2 expect the weights in HWOI order.
template <typename InputScalar, typename DestinationScalar>
inline void TransposeConvV2(
const ConvParams& params, const int32* output_multiplier,
const int32* output_shift, const RuntimeShape& input_shape,
const InputScalar* input_data,
const RuntimeShape& hwoi_ordered_filter_shape,
const int8_t* hwoi_ordered_filter_data, const RuntimeShape& bias_shape,
const int32* bias_data, const RuntimeShape& output_shape,
DestinationScalar* output_data, const RuntimeShape& col2im_shape,
int32_t* col2im_data, int32_t* scratch_data,
CpuBackendContext* cpu_backend_context) {
ruy::profiler::ScopeLabel label("TransposeConvV2/int8");
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(hwoi_ordered_filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK(col2im_data);
TFLITE_DCHECK(hwoi_ordered_filter_data);
const int batch_size = MatchingDim(input_shape, 0, output_shape, 0);
const int input_image_size = input_shape.Dims(1) * input_shape.Dims(2);
const int output_height = output_shape.Dims(1);
const int output_width = output_shape.Dims(2);
const int output_image_size = output_height * output_width;
const int input_depth =
MatchingDim(input_shape, 3, hwoi_ordered_filter_shape, 3);
const int output_depth =
MatchingDim(output_shape, 3, hwoi_ordered_filter_shape, 2);
const int input_offset = input_image_size * input_depth;
const int output_offset = output_image_size * output_depth;
const int filter_height = hwoi_ordered_filter_shape.Dims(0);
const int filter_width = hwoi_ordered_filter_shape.Dims(1);
const int padding_top = params.padding_values.height;
const int padding_bottom =
params.padding_values.height + params.padding_values.height_offset;
const int padding_left = params.padding_values.width;
const int padding_right =
params.padding_values.width + params.padding_values.width_offset;
const int stride_height = params.stride_height;
const int stride_width = params.stride_width;
const int32 output_activation_min = params.quantized_activation_min;
const int32 output_activation_max = params.quantized_activation_max;
const int hwoi_ordered_filter_total_size =
filter_height * filter_width * output_depth;
cpu_backend_gemm::MatrixParams<int8_t> lhs_params;
lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
lhs_params.rows = hwoi_ordered_filter_total_size;
lhs_params.cols = input_depth;
// Since our weight is symmetric quantized, the zp will always be 0.
lhs_params.zero_point = 0;
int32_t* scratch_data_p = scratch_data;
std::fill_n(scratch_data, output_offset * batch_size, static_cast<int32>(0));
for (int i = 0; i < batch_size; ++i) {
cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
rhs_params.order = cpu_backend_gemm::Order::kColMajor;
rhs_params.rows = input_depth;
rhs_params.cols = input_image_size;
rhs_params.zero_point = -params.input_offset;
cpu_backend_gemm::MatrixParams<int32_t> dst_params;
dst_params.order = cpu_backend_gemm::Order::kColMajor;
dst_params.rows = hwoi_ordered_filter_total_size;
dst_params.cols = input_image_size;
cpu_backend_gemm::GemmParams<int32_t, int32_t> gemm_params;
cpu_backend_gemm::Gemm(lhs_params, hwoi_ordered_filter_data, rhs_params,
input_data + input_offset * i, dst_params,
col2im_data, gemm_params, cpu_backend_context);
optimized_ops::Col2im(
col2im_data, output_depth, output_height, output_width, filter_height,
filter_width, padding_top, padding_left, padding_bottom, padding_right,
stride_height, stride_width, scratch_data_p);
scratch_data_p += output_offset;
}
scratch_data_p = scratch_data;
optimized_ops::BiasAdd(scratch_data_p, bias_data, batch_size, output_height,
output_width, output_depth);
optimized_ops::Quantize(output_multiplier, output_shift, output_depth,
output_shape.FlatSize(), params.output_offset,
output_activation_min, output_activation_max,
scratch_data, output_data);
}
} // namespace optimized_integer_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_TRANSPOSE_CONV_H_
tensorflow/lite/kernels/internal/optimized/legacy_optimized_ops.h
+5022
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/multithreaded_conv.h
+184
View File
@@ -0,0 +1,184 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_MULTITHREADED_CONV_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_MULTITHREADED_CONV_H_
#include <assert.h>
#include <stdint.h>
#include <sys/types.h>
#include <algorithm>
#include <cmath>
#include <limits>
#include <memory>
#include <tuple>
#include <type_traits>
#include "tensorflow/lite/core/c/builtin_op_data.h"
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/optimized/eigen_spatial_convolutions.h"
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace multithreaded_ops {
// Shorthands for the types we need when interfacing with the EigenTensor
// library.
typedef Eigen::TensorMap<
Eigen::Tensor<float, 2, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>
EigenMatrix;
typedef Eigen::TensorMap<
const Eigen::Tensor<float, 2, Eigen::RowMajor, Eigen::DenseIndex>,
Eigen::Aligned>
ConstEigenMatrix;
typedef Eigen::TensorMap<
Eigen::Tensor<float, 4, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>
EigenTensor;
typedef Eigen::TensorMap<
const Eigen::Tensor<float, 4, Eigen::RowMajor, Eigen::DenseIndex>,
Eigen::Aligned>
ConstEigenTensor;
// Utility functions we need for the EigenTensor API.
template <typename Device, typename T>
struct MatMulConvFunctor {
// Computes on device "d": out = in0 * in1, where * is matrix
// multiplication.
void operator()(
const Device& d, EigenMatrix out, ConstEigenMatrix in0,
ConstEigenMatrix in1,
const Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1>& dim_pair) {
out.device(d) = in0.contract(in1, dim_pair);
}
};
template <class T>
class EigenTensorConvFunctor {
private:
Eigen::PaddingType RuntimePadding2EigenPadding(PaddingType padding) {
switch (padding) {
case PaddingType::kValid:
return Eigen::PADDING_VALID;
case PaddingType::kSame:
return Eigen::PADDING_SAME;
case PaddingType::kNone:
assert(false); // should never get here.
return Eigen::PADDING_VALID;
}
return Eigen::PADDING_SAME; // Prevent compiler warning about missing
// return
}
public:
void operator()(const Eigen::ThreadPoolDevice& device, const T* input_data,
int input_batches, int input_height, int input_width,
int input_depth, const T* filter_data, int filter_height,
int filter_width, int filter_count, int stride_rows,
int stride_cols, int pad_width, int pad_height,
PaddingType padding, T* output_data, int output_height,
int output_width) {
const bool is_1x1_kernel = (filter_height == 1 && filter_width == 1 &&
stride_rows == 1 && stride_cols == 1);
if (is_1x1_kernel) {
// For 1x1 kernel, the 2D convolution is reduced to matrix
// multiplication.
const int conv_width = output_height * output_width;
Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> dim_pair;
dim_pair[0] = Eigen::IndexPair<Eigen::DenseIndex>(1, 0);
EigenMatrix output(output_data, input_batches * conv_width, filter_count);
ConstEigenMatrix input(input_data, input_batches * conv_width,
input_depth);
ConstEigenMatrix filter(filter_data, input_depth, filter_count);
MatMulConvFunctor<Eigen::ThreadPoolDevice, T>()(device, output, input,
filter, dim_pair);
} else if (filter_height == input_height && filter_width == input_width &&
pad_width == 0 && pad_height == 0) {
// If the input data and filter have the same height/width,
// the 2D convolution is reduced to matrix multiplication.
const int k = // Length of reduction dimension.
filter_width * filter_height * input_depth;
Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> dim_pair;
dim_pair[0] = Eigen::IndexPair<Eigen::DenseIndex>(1, 0);
EigenMatrix output(output_data, input_batches, filter_count);
ConstEigenMatrix input(input_data, input_batches, k);
ConstEigenMatrix filter(filter_data, k, filter_count);
MatMulConvFunctor<Eigen::ThreadPoolDevice, T>()(device, output, input,
filter, dim_pair);
} else {
EigenTensor output(output_data, input_batches, output_height,
output_width, filter_count);
ConstEigenTensor input(input_data, input_batches, input_height,
input_width, input_depth);
ConstEigenTensor filter(filter_data, filter_height, filter_width,
input_depth, filter_count);
output.device(device) =
Eigen::SpatialConvolution(input, filter, stride_cols, stride_rows,
RuntimePadding2EigenPadding(padding));
}
}
};
inline void Conv(const Eigen::ThreadPoolDevice& device,
const ConvParams& params, const RuntimeShape& input_shape,
const float* input_data, const RuntimeShape& filter_shape,
const float* filter_data, const RuntimeShape& bias_shape,
const float* bias_data, const RuntimeShape& output_shape,
float* output_data, const RuntimeShape& im2col_shape,
float* im2col_data) {
// Nest profiling under "Conv", to aggregate with other kernels.
ruy::profiler::ScopeLabel label("Conv");
ruy::profiler::ScopeLabel inner_label("Multithreaded EigenTensor");
// im2col data should not be generated for the multi-thread supporting case.
TFLITE_DCHECK(!im2col_data);
(void)im2col_shape;
const int stride_width = params.stride_width;
const int stride_height = params.stride_height;
const PaddingType padding = params.padding_type;
const int pad_width = params.padding_values.width;
const int pad_height = params.padding_values.height;
const float output_activation_min = params.float_activation_min;
const float output_activation_max = params.float_activation_max;
TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
const int batches = MatchingDim(input_shape, 0, output_shape, 0);
const int input_depth = MatchingDim(input_shape, 3, filter_shape, 3);
const int output_depth = MatchingDim(filter_shape, 0, output_shape, 3);
const int input_height = input_shape.Dims(1);
const int input_width = input_shape.Dims(2);
const int filter_height = filter_shape.Dims(1);
const int filter_width = filter_shape.Dims(2);
const int output_height = output_shape.Dims(1);
const int output_width = output_shape.Dims(2);
EigenTensorConvFunctor<float> conv_functor;
conv_functor(device, input_data, batches, input_height, input_width,
input_depth, filter_data, filter_height, filter_width,
output_depth, stride_height, stride_width, pad_height, pad_width,
padding, output_data, output_height, output_width);
optimized_ops::AddBiasAndEvalActivationFunction(
output_activation_min, output_activation_max, bias_shape, bias_data,
output_shape, output_data);
}
} // namespace multithreaded_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_MULTITHREADED_CONV_H_
tensorflow/lite/kernels/internal/optimized/neon_check.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
#define USE_NEON
#include <arm_neon.h> // IWYU pragma: export
#endif
#if defined __GNUC__ && defined __SSE4_1__ && !defined TF_LITE_DISABLE_X86_NEON
#define USE_NEON
#include "NEON_2_SSE.h" // IWYU pragma: export
#endif
// NEON_OR_PORTABLE(SomeFunc, args) calls NeonSomeFunc(args) if USE_NEON is
// defined, PortableSomeFunc(args) otherwise.
#ifdef USE_NEON
// Always use Neon code
#define NEON_OR_PORTABLE(funcname, ...) Neon##funcname(__VA_ARGS__)
#else
// No NEON available: Use Portable code
#define NEON_OR_PORTABLE(funcname, ...) Portable##funcname(__VA_ARGS__)
#endif // defined(USE_NEON)
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_
tensorflow/lite/kernels/internal/optimized/neon_tensor_utils.cc
+2806
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/neon_tensor_utils.h
+334
View File
@@ -0,0 +1,334 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
#include "tensorflow/lite/kernels/internal/optimized/neon_tensor_utils_impl.h"
#include "tensorflow/lite/kernels/internal/reference/portable_tensor_utils_impl.h"
namespace tflite {
namespace tensor_utils {
void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
int m_cols, const float* vector,
int n_batch, float* result) {
NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vector, n_batch, result);
}
void MatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* scaling_factors,
int n_batch,
float* __restrict__ result) {
NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vectors, scaling_factors, n_batch, result);
}
void MatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* scaling_factors,
int n_batch, int32_t* scratch,
float* __restrict__ result,
CpuBackendContext* context) {
NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vectors, scaling_factors, n_batch, scratch, result, context);
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors, const float* scaling_factors,
int n_batch, float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
bool* compute_row_sums, CpuBackendContext* context) {
NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vectors, scaling_factors, n_batch, result, per_channel_scale,
input_offset, scratch, row_sums, compute_row_sums, context);
}
void SparseMatrixBatchVectorMultiplyAccumulate1x4(
const float* __restrict__ matrix, const int32_t* __restrict__ segments,
const int32_t* __restrict__ indices, int m_rows, int m_cols,
const float* __restrict__ vector, int n_batch, float* __restrict__ result) {
NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x4, matrix,
segments, indices, m_rows, m_cols, vector, n_batch, result);
}
void SparseMatrixBatchVectorMultiplyAccumulate(
const float* __restrict__ matrix, const uint8_t* __restrict__ ledger,
int m_rows, int m_cols, const float* __restrict__ vector, int n_batch,
float* __restrict__ result) {
NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
m_rows, m_cols, vector, n_batch, result);
}
void SparseMatrixBatchVectorMultiplyAccumulate1x16(
const int8_t* __restrict__ matrix, const int32_t* __restrict__ segments,
const int32_t* __restrict__ indices, int m_rows, int m_cols,
const int8_t* __restrict__ vector, const int32_t* __restrict__ bias_vector,
int n_batch, const int32_t input_offset, const int32_t output_multiplier,
const int32_t output_shift, const int32_t* per_channel_scale,
const int32_t* per_channel_shift, const int32_t output_offset,
const int32_t output_activation_min, const int32_t output_activation_max,
int8_t* __restrict__ result) {
NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x16, matrix,
segments, indices, m_rows, m_cols, vector, bias_vector,
n_batch, input_offset, output_multiplier, output_shift,
per_channel_scale, per_channel_shift, output_offset,
output_activation_min, output_activation_max, result);
}
void SparseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const uint8_t* ledger, const int m_rows,
const int m_cols, const int8_t* __restrict__ vectors,
const float* scaling_factors, int n_batch, float* __restrict__ result,
const float* per_channel_scale) {
NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
m_rows, m_cols, vectors, scaling_factors, n_batch, result,
per_channel_scale);
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* input, const int32_t* bias,
const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
int32_t* scratch, int16_t* output, CpuBackendContext* context) {
NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, input, bias,
input_to_gate_weights, multiplier, shift, n_batch, n_input,
n_output, output_zp, scratch, output, context);
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* input, const int32_t* bias,
const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
int32_t* scratch, int8_t* output, CpuBackendContext* context) {
NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, input, bias,
input_to_gate_weights, multiplier, shift, n_batch, n_input,
n_output, output_zp, scratch, output, context);
}
void MatrixBatchVectorMultiply(const int8_t* input, int32_t input_zeropoint,
const int8_t* input_to_gate_weights,
int32_t input_to_gate_effective_scale_a,
int32_t input_to_gate_effective_scale_b,
int32_t n_batch, int32_t n_input, int32_t n_cell,
int8_t* gate_output, int8_t gate_output_zp) {
PortableMatrixBatchVectorMultiply(
input, input_zeropoint, input_to_gate_weights,
input_to_gate_effective_scale_a, input_to_gate_effective_scale_b, n_batch,
n_input, n_cell, gate_output, gate_output_zp);
}
void MatrixBatchVectorMultiply(const int16_t* hidden,
const int8_t* hidden_to_output_weights,
int32_t proj_effective_scale_a,
int32_t proj_effective_scale_b,
const int32_t* gate_bias, int32_t n_batch,
int32_t n_hidden, int32_t n_output,
int32_t output_zp, int8_t* proj_output) {
PortableMatrixBatchVectorMultiply(hidden, hidden_to_output_weights,
proj_effective_scale_a,
proj_effective_scale_b, gate_bias, n_batch,
n_hidden, n_output, output_zp, proj_output);
}
void MatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar,
int32_t n_row, int32_t n_col,
int32_t* output) {
NEON_OR_PORTABLE(MatrixScalarMultiplyAccumulate, matrix, scalar, n_row, n_col,
output);
}
void ApplyLayerNorm(const int16_t* input, const int16_t* layer_norm_weights,
const int32_t* bias, int32_t layer_norm_scale_a,
int32_t layer_norm_scale_b, int32_t variance_limit,
int n_batch, int n_input, int16_t* output) {
NEON_OR_PORTABLE(ApplyLayerNorm, input, layer_norm_weights, bias,
layer_norm_scale_a, layer_norm_scale_b, variance_limit,
n_batch, n_input, output);
}
void ApplyLayerNormFloat(const int16_t* input,
const int16_t* layer_norm_weights,
int32_t layer_norm_scale_a, int32_t layer_norm_scale_b,
const int32_t* bias, int n_batch, int n_input,
int16_t* output) {
PortableApplyLayerNormFloat(input, layer_norm_weights, layer_norm_scale_a,
layer_norm_scale_b, bias, n_batch, n_input,
output);
}
void ApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input,
int16_t* output) {
NEON_OR_PORTABLE(ApplySigmoid, input, n_batch, n_input, output);
}
void ApplySigmoidFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
int16_t* output) {
PortableApplySigmoidFloat(input, n_batch, n_input, output);
}
void ApplyTanh(int32_t integer_bits, const int16_t* input, int32_t n_batch,
int32_t n_input, int16_t* output) {
NEON_OR_PORTABLE(ApplyTanh, integer_bits, input, n_batch, n_input, output);
}
void ApplyTanhFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
int32_t integer_bits, int16_t* output) {
PortableApplyTanhFloat(input, n_batch, n_input, integer_bits, output);
}
void CwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
int n_input, int shift, int16_t* output) {
NEON_OR_PORTABLE(CwiseMul, input_1, input_2, n_batch, n_input, shift, output);
}
void CwiseMul(const int16_t* input_1, const int16_t* input_2,
int32_t multiplier, int shift, int n_batch, int n_input,
int32_t output_zp, int8_t* output) {
NEON_OR_PORTABLE(CwiseMul, input_1, input_2, multiplier, shift, n_batch,
n_input, output_zp, output);
}
void CwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch,
int n_input, int16_t* output) {
NEON_OR_PORTABLE(CwiseAdd, input_1, input_2, n_batch, n_input, output);
}
void CwiseClipping(float* vector, const int v_size,
const float clipping_value) {
NEON_OR_PORTABLE(CwiseClipping, vector, v_size, clipping_value);
}
void CwiseClipping(int16_t* vector, const int v_size,
const int16_t clipping_value) {
NEON_OR_PORTABLE(CwiseClipping, vector, v_size, clipping_value);
}
void CwiseClipping(int8_t* vector, const int v_size,
const int8_t clipping_value) {
NEON_OR_PORTABLE(CwiseClipping, vector, v_size, clipping_value);
}
void BatchVectorBatchVectorDotProduct(const int16_t* vector1,
const int16_t* vector2, int v_size,
int n_batch, int32_t* result) {
PortableBatchVectorBatchVectorDotProduct(vector1, vector2, v_size, n_batch,
result);
}
void VectorBatchVectorCwiseProductAccumulate(const int16_t* vector, int v_size,
const int16_t* batch_vector,
int n_batch, int32_t multiplier,
int shift, int16_t* result) {
NEON_OR_PORTABLE(VectorBatchVectorCwiseProductAccumulate, vector, v_size,
batch_vector, n_batch, multiplier, shift, result);
}
float VectorVectorDotProduct(const float* vector1, const float* vector2,
int v_size) {
return NEON_OR_PORTABLE(VectorVectorDotProduct, vector1, vector2, v_size);
}
void Sub1Vector(const float* vector, int v_size, float* result) {
NEON_OR_PORTABLE(Sub1Vector, vector, v_size, result);
}
void Sub1Vector(const int16_t* vector, int v_size, int16_t* result) {
NEON_OR_PORTABLE(Sub1Vector, vector, v_size, result);
}
// Check if all entries of a vector are zero for float.
bool IsZeroVector(const float* vector, int v_size) {
return NEON_OR_PORTABLE(IsZeroVector, vector, v_size);
}
// Check if all entries of a vector are zero for int8.
bool IsZeroVector(const int8_t* vector, int v_size) {
return NEON_OR_PORTABLE(IsZeroVector, vector, v_size);
}
void VectorScalarMultiply(const int8_t* vector, int v_size, float scale,
float* result) {
NEON_OR_PORTABLE(VectorScalarMultiply, vector, v_size, scale, result);
}
void SymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float* min_value,
float* max_value, float* scaling_factor) {
NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
min_value, max_value, scaling_factor);
}
void SymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float min_value,
float max_value, float* scaling_factor) {
NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
min_value, max_value, scaling_factor);
}
void AsymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float* scaling_factor,
int32_t* offset) {
NEON_OR_PORTABLE(AsymmetricQuantizeFloats, values, size, quantized_values,
scaling_factor, offset);
}
void ReductionSumVector(const float* input_vector, float* output_vector,
int output_size, int reduction_size) {
NEON_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
reduction_size);
}
void ReductionSumVector(const int32_t* input_vector, int32_t* output_vector,
int output_size, int reduction_size) {
PortableReductionSumVector(input_vector, output_vector, output_size,
reduction_size);
}
void ReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
int output_size, int reduction_size) {
NEON_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
reduction_size);
}
void MeanStddevNormalization(const float* __restrict__ input_vector,
float* __restrict__ output_vector, int v_size,
int n_batch) {
NEON_OR_PORTABLE(MeanStddevNormalization, input_vector, output_vector, v_size,
n_batch);
}
void TwoGateSaturatingAdd(const int8_t* input, int8_t input_zp,
const int8_t* recurrent, int8_t recurrent_zp,
int32_t input_effective_scale_a,
int32_t input_effective_scale_b,
int32_t recurrent_effective_scale_a,
int32_t recurrent_effective_scale_b, int32_t n_batch,
int32_t n_cell, int16_t* output) {
PortableTwoGateSaturatingAdd(
input, input_zp, recurrent, recurrent_zp, input_effective_scale_a,
input_effective_scale_b, recurrent_effective_scale_a,
recurrent_effective_scale_b, n_batch, n_cell, output);
}
} // namespace tensor_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_
tensorflow/lite/kernels/internal/optimized/neon_tensor_utils_impl.h
+198
View File
@@ -0,0 +1,198 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_IMPL_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_IMPL_H_
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
#if defined(_MSC_VER)
#define __restrict__ __restrict
#endif
namespace tflite {
namespace tensor_utils {
#ifdef USE_NEON
// Multiply a matrix by a batch vector, and store results in a batch-size
// vector.
void NeonMatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
int m_cols, const float* vector,
int n_batch, float* result);
// Matrix multiplication for quantized values using symmetric quantization.
void NeonMatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* scaling_factors,
int n_batch,
float* __restrict__ result);
// Same as above but with a scratch buffer and CpuBackendContext for the
// int8 x int8 -> int32 accumulation computation
void NeonMatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* scaling_factors,
int n_batch, int32_t* scratch,
float* __restrict__ result,
CpuBackendContext* context);
// Matrix multiplication for quantized values using asymmetric quantization.
void NeonMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors, const float* scaling_factors,
int n_batch, float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
bool* compute_row_sums, CpuBackendContext* context);
void NeonApplyLayerNorm(const int16_t* input, const int16_t* layer_norm_weights,
const int32_t* bias, int32_t layer_norm_scale_a,
int32_t layer_norm_scale_b, int32_t variance_limit,
int n_batch, int n_input, int16_t* output);
void NeonApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input,
int16_t* output);
void NeonApplyTanh(int32_t integer_bits, const int16_t* input, int32_t n_batch,
int32_t n_input, int16_t* output);
void NeonCwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
int n_input, int shift, int16_t* output);
void NeonCwiseMul(const int16_t* input_1, const int16_t* input_2,
int32_t multiplier, int shift, int n_batch, int n_input,
int32_t output_zp, int8_t* output);
void NeonCwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch,
int n_input, int16_t* output);
void NeonCwiseClipping(float* vector, const int v_size,
const float clipping_value);
void NeonCwiseClipping(int16_t* vector, const int v_size,
const int16_t clipping_value);
void NeonCwiseClipping(int8_t* vector, const int v_size,
const int8_t clipping_value);
void NeonMatrixBatchVectorMultiplyAccumulate(
const int8_t* input, const int32_t* bias,
const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
int32_t* scratch, int8_t* output, CpuBackendContext* context);
void NeonMatrixBatchVectorMultiplyAccumulate(
const int8_t* input, const int32_t* bias,
const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
int32_t* scratch, int16_t* output, CpuBackendContext* context);
void NeonMatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar,
int32_t n_row, int32_t n_col,
int32_t* output);
void NeonSparseMatrixBatchVectorMultiplyAccumulate1x4(
const float* __restrict__ matrix, const int32_t* __restrict__ segments,
const int32_t* __restrict__ indices, int m_rows, int m_cols,
const float* __restrict__ vector, int n_batch, float* __restrict__ result);
// Multiply a matrix by a batch vector, and store results in a batch-size
// vector. Sparse version.
void NeonSparseMatrixBatchVectorMultiplyAccumulate(
const float* __restrict__ matrix, const uint8_t* __restrict__ ledger,
int m_rows, int m_cols, const float* __restrict__ vector, int n_batch,
float* __restrict__ result);
// Multiplies a symmetric quantized matrix by a quantized batch vector. The
// matrix is stored in sparse format.
void NeonSparseMatrixBatchVectorMultiplyAccumulate1x16(
const int8_t* __restrict__ matrix, const int32_t* __restrict__ segments,
const int32_t* __restrict__ indices, int m_rows, int m_cols,
const int8_t* __restrict__ vector, const int32_t* __restrict__ bias_vector,
int n_batch, const int32_t input_offset, const int32_t output_multiplier,
int32_t output_shift, const int32_t* per_channel_scale,
const int32_t* per_channel_shift, int32_t output_offset,
const int32_t output_activation_min, const int32_t output_activation_max,
int8_t* __restrict__ result);
// Matrix multiplication for quantized values using symmetric quantization.
// Sparse version.
void NeonSparseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const uint8_t* ledger, const int m_rows,
const int m_cols, const int8_t* __restrict__ vectors,
const float* scaling_factors, int n_batch, float* __restrict__ result,
const float* per_channel_scale);
// Dot product of two vectors.
float NeonVectorVectorDotProduct(const float* vector1, const float* vector2,
int v_size);
// Compute "1.0f - elements of vector" (used in CIFG).
void NeonSub1Vector(const float* vector, int v_size, float* result);
void NeonSub1Vector(const int16_t* vector, int v_size, int16_t* result);
// Multiply all elements of vector with a scalar.
void NeonVectorScalarMultiply(const int8_t* vector, int v_size, float scale,
float* result);
// Check if all entries of a vector are zero.
bool NeonIsZeroVector(const float* vector, int v_size);
// Check if all entries of a vector are zero.
bool NeonIsZeroVector(const int8_t* vector, int v_size);
// Symmetric quantizer.
void NeonSymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float* min,
float* max, float* scaling_factor);
// Symmetric quantizer.
void NeonSymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float min, float max,
float* scaling_factor);
// Asymmetric quantizer.
void NeonAsymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values,
float* scaling_factor, int32_t* offset);
// Reduce-sum on a float input vector:
// input_vector: float pointer to input vector.
// output_vector: float pointer to vector.
// output_size: output vector size.
// reduction_size: number of consecutive elements from input vector which are
// added to get one element of output.
void NeonReductionSumVector(const float* input_vector, float* output_vector,
int output_size, int reduction_size);
void NeonReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
int output_size, int reduction_size);
void NeonVectorBatchVectorCwiseProductAccumulate(
const int16_t* vector, int v_size, const int16_t* batch_vector, int n_batch,
int32_t multiplier, int shift, int16_t* result);
// Layer norm for each batch.
void NeonMeanStddevNormalization(const float* __restrict__ input_vector,
float* __restrict__ output_vector, int v_size,
int n_batch);
#endif // USE_NEON
} // namespace tensor_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_IMPL_H_
tensorflow/lite/kernels/internal/optimized/optimized_4bit_test.cc
+597
View File
@@ -0,0 +1,597 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include <cstdint>
#include <cstdlib>
#include <functional>
#include <memory>
#include <random>
#include <tuple>
#include <type_traits>
#include <vector>
#include <gtest/gtest.h>
#include "tensorflow/lite/kernels/internal/optimized/fully_connected_4bit.h"
namespace tflite {
namespace {
std::mt19937 random_engine(2023);
std::uniform_real_distribution<float> real_dist(0.f, 1.f);
std::uniform_int_distribution<int32_t> int_dist(-7, 7);
struct TestPack {
TestPack(std::vector<int8_t> src_data, int src_rows, int src_cols, int width,
int depth)
: src_data(src_data),
src_rows(src_rows),
src_cols(src_cols),
width(width),
depth(depth),
rows((src_rows + (width - 1)) & ~(width - 1)),
cols((src_cols + (depth - 1)) & ~(depth - 1)),
// Must be vector-aligned.
packed_data_buffer(
[=]() -> uint8_t* {
void* ptr;
if (posix_memalign(&ptr, 64, rows * cols + padding)) {
abort();
}
return static_cast<uint8_t*>(ptr);
}(),
[](uint8_t* ptr) { free(ptr); }) {}
void Prepack() {
packed_data = packed_data_buffer.get();
optimized_4bit::Prepack(packed_data, src_data.data(), rows, cols, src_rows,
src_cols, width, depth);
}
std::vector<uint8_t> AsVector() {
int size = rows * cols / 2;
std::vector<uint8_t> values(size);
for (int i = 0; i < size; i++) {
values[i] = packed_data[i];
}
packed_data = nullptr;
return values;
}
std::vector<int8_t> src_data;
uint8_t* packed_data;
int src_rows;
int src_cols;
int width;
int depth;
int rows;
int cols;
int padding = optimized_4bit::kDefaultAlignmentPadding;
std::unique_ptr<uint8_t, std::function<void(uint8_t*)>> packed_data_buffer;
};
class RunPackTests
: public ::testing::TestWithParam<::testing::tuple<int, int>> {};
TEST_P(RunPackTests, RunPackTests) {
auto params = GetParam();
int src_rows = std::get<0>(params);
int src_cols = std::get<1>(params);
std::vector<int8_t> test_data;
test_data.reserve(src_rows * src_cols / 2);
for (int i = 0; i < src_rows; ++i) {
int stride = optimized_4bit::FilterDepth / 4;
int strides = src_cols / stride / 2;
int v = -7;
int l = 0;
for (int j = 0; j < strides; j++) {
for (int k = 0; k < stride; k++) { // 8
int lower = static_cast<uint8_t>(v) & UINT8_C(15);
int upper = static_cast<uint8_t>(v) << 4;
test_data.push_back(upper | lower);
l++;
}
v++;
}
while (l < (src_cols / 2)) {
int lower = static_cast<uint8_t>(v) & UINT8_C(15);
test_data.push_back(lower << 4 | lower);
l++;
}
}
TestPack test(test_data, src_rows, src_cols, optimized_4bit::FilterWidth,
optimized_4bit::FilterDepth);
test.Prepack();
std::vector<uint8_t> result = test.AsVector();
int outer_rows = test.rows / optimized_4bit::FilterWidth;
int outer_cols = test.cols / optimized_4bit::FilterDepth;
int k = 0;
for (int i = 0; i < outer_rows; ++i) {
int v = -7;
for (int j = 0; j < outer_cols; ++j) {
for (int w = 0; w < optimized_4bit::FilterWidth; w++) {
int c = 0;
for (; c < optimized_4bit::FilterDepth / 2; c++) {
uint8_t res = result[k++];
uint8_t res0 = res >> 4;
uint8_t res1 = res & UINT8_C(15);
int res00 = res0 - 7;
int res11 = res1 - 7;
if ((i * optimized_4bit::FilterWidth + w) < src_rows) {
if ((j * optimized_4bit::FilterDepth / 2 + c) < src_cols / 2) {
EXPECT_EQ(res00, v % 8);
}
if ((j * optimized_4bit::FilterDepth / 2 + c + 16) < src_cols / 2) {
EXPECT_EQ(res11, v + 1 % 8);
}
}
}
}
v += 2;
}
}
}
INSTANTIATE_TEST_SUITE_P(RunPackTests, RunPackTests,
::testing::ValuesIn({
std::make_tuple(4, 32),
std::make_tuple(4, 46),
std::make_tuple(4, 56),
std::make_tuple(5, 64),
std::make_tuple(5, 72),
std::make_tuple(5, 80),
std::make_tuple(5, 84),
}));
struct TestQuantize {
TestQuantize(std::vector<float> src_data, int src_rows, int src_cols,
int width, int depth)
: src_data(src_data),
src_rows(src_rows),
src_cols(src_cols),
width(width),
depth(depth) {
rows = (src_rows + (width - 1)) & ~(width - 1);
cols = (src_cols + (depth - 1)) & ~(depth - 1);
scaling_factors.assign(rows, 1.0);
input_offsets.assign(rows, 0);
output_data.assign(rows * cols, 0);
}
void BatchQuantizeFloats4Bit() {
optimized_4bit::BatchQuantizeFloats4Bit(
src_data.data(), src_rows, src_cols, output_data.data(),
scaling_factors.data(), width, depth, input_offsets.data());
}
std::vector<float> src_data;
int src_rows;
int src_cols;
int rows;
int cols;
int width;
int depth;
std::vector<int8_t> output_data;
std::vector<float> scaling_factors;
std::vector<int32_t> input_offsets;
};
class RunQuantizeInputTests
: public ::testing::TestWithParam<::testing::tuple<int, int, int>> {};
TEST_P(RunQuantizeInputTests, RunQuantizeInputsTests) {
auto params = GetParam();
int width = std::get<0>(params);
int src_rows = std::get<1>(params);
int src_cols = std::get<2>(params);
std::vector<float> test_data;
test_data.reserve(src_rows * src_cols);
float v = -127.0;
for (int i = 0; i < src_rows; ++i) {
for (int j = 0; j < src_cols; ++j) {
test_data.push_back(v / (i + 1));
v = -v;
}
}
TestQuantize test(test_data, src_rows, src_cols, width,
optimized_4bit::FilterDepth);
test.BatchQuantizeFloats4Bit();
int8_t* result = test.output_data.data();
int k = 0;
int outer_rows = test.rows / width;
int outer_cols = test.cols / optimized_4bit::FilterDepth;
for (int i = 0; i < outer_rows; ++i) {
for (int j = 0; j < outer_cols; ++j) {
for (int w = 0; w < width; w++) {
int c = 0;
v = -127;
for (; c < optimized_4bit::FilterDepth; c++) {
int8_t res = result[k++];
int res0 = static_cast<int>(res);
if ((i * width + w) < src_rows) {
if ((j * optimized_4bit::FilterDepth + c) < src_cols) {
EXPECT_EQ(res0, v);
}
}
v = -v;
}
}
v += 2;
}
}
for (int i = 0; i < test.rows; i++) {
if (i >= src_rows) {
continue;
}
EXPECT_EQ(test.input_offsets[i], 0);
EXPECT_NEAR(test.scaling_factors[i], 1.0 / (1 + i), 1e-3);
}
}
INSTANTIATE_TEST_SUITE_P(RunQuantizeInputTests, RunQuantizeInputTests,
::testing::ValuesIn({
std::make_tuple(1, 1, 32),
std::make_tuple(1, 3, 46),
std::make_tuple(1, 9, 64),
std::make_tuple(1, 25, 72),
std::make_tuple(2, 2, 32),
std::make_tuple(2, 3, 46),
std::make_tuple(2, 9, 64),
std::make_tuple(2, 25, 72),
std::make_tuple(4, 4, 32),
std::make_tuple(4, 5, 46),
std::make_tuple(4, 9, 64),
std::make_tuple(4, 25, 72),
}));
struct TestAssignBiasAndComputeOffset {
TestAssignBiasAndComputeOffset(std::vector<float> output_data,
std::vector<int32_t> input_offsets,
std::vector<float> input_scales,
std::vector<float> filter_scales,
std::vector<float> bias, int output_rows,
int output_cols, bool use_bias)
: output_data(output_data),
input_offsets(input_offsets),
input_scales(input_scales),
filter_scales(filter_scales),
bias(bias),
output_rows(output_rows),
output_cols(output_cols),
use_bias(use_bias) {}
void AssignBiasAndComputeOffsets() {
optimized_4bit::AssignBiasAndComputeOffsets(
input_offsets.data(), input_scales.data(), filter_scales.data(),
use_bias ? bias.data() : nullptr, output_data.data(), output_cols,
output_rows);
}
std::vector<float> output_data;
std::vector<int32_t> input_offsets;
std::vector<float> input_scales;
std::vector<float> filter_scales;
std::vector<float> bias;
int output_rows;
int output_cols;
bool use_bias;
};
class RunAssignBiasAndOffsetsTests
: public ::testing::TestWithParam<::testing::tuple<int, int, bool>> {};
TEST_P(RunAssignBiasAndOffsetsTests, RunAssignBiasAndOffsetssTests) {
auto params = GetParam();
int output_rows = std::get<0>(params);
int output_cols = std::get<1>(params);
bool use_bias = std::get<2>(params);
std::vector<float> test_data(output_rows * output_cols, 0);
std::vector<float> test_input_scales(output_rows);
std::vector<int32_t> test_input_offsets(output_rows);
std::vector<float> test_filter_scales(output_cols);
std::vector<float> test_bias(output_cols);
for (int i = 0; i < output_rows; ++i) {
test_input_scales[i] = real_dist(random_engine);
test_input_offsets[i] = int_dist(random_engine);
}
for (int i = 0; i < output_cols; ++i) {
test_filter_scales[i] = real_dist(random_engine);
test_bias[i] = real_dist(random_engine);
}
TestAssignBiasAndComputeOffset test(
test_data, test_input_offsets, test_input_scales, test_filter_scales,
test_bias, output_rows, output_cols, use_bias);
test.AssignBiasAndComputeOffsets();
float* result = test.output_data.data();
for (int i = 0; i < output_rows; ++i) {
for (int j = 0; j < output_cols; ++j) {
float val = result[i * output_cols + j];
float expected = use_bias ? test_bias[j] : 0;
expected +=
test_input_offsets[i] * test_input_scales[i] * test_filter_scales[j];
EXPECT_NEAR(val, expected, 1e-3);
}
}
}
INSTANTIATE_TEST_SUITE_P(RunAssignBiasAndOffsetsTests,
RunAssignBiasAndOffsetsTests,
::testing::ValuesIn({
std::make_tuple(1, 8, true),
std::make_tuple(4, 17, false),
std::make_tuple(4, 17, true),
std::make_tuple(11, 33, false),
std::make_tuple(11, 33, true),
}));
struct TestUnpack {
TestUnpack(std::vector<int32_t> src_data, std::vector<float> input_scales,
std::vector<float> filter_scales, int src_rows, int src_cols,
int output_rows, int output_cols)
: src_data(src_data),
input_scales(input_scales),
filter_scales(filter_scales),
src_rows(src_rows),
src_cols(src_cols),
output_rows(output_rows),
output_cols(output_cols) {
output_data.assign(output_rows * output_cols, 0.0);
}
template <int Depth, int Width>
void Unpack() {
optimized_4bit::Unpack<Depth, Width>(
output_data.data(), src_data.data(), output_rows, output_cols,
input_scales.data(), filter_scales.data(), src_rows, src_cols);
}
std::vector<int32_t> src_data;
std::vector<float> input_scales;
std::vector<float> filter_scales;
std::vector<float> output_data;
int src_rows;
int src_cols;
int output_rows;
int output_cols;
};
class RunUnpackTests
: public ::testing::TestWithParam<::testing::tuple<int, int, int>> {};
TEST_P(RunUnpackTests, RunUnpackTests) {
auto params = GetParam();
int src_rows = std::get<0>(params);
int src_cols = std::get<1>(params);
// In this case, we only unpack 1 rhs row,
// so the batch_size and accumulator rows must match.
int output_rows = src_rows;
int output_cols = std::get<2>(params);
std::vector<float> test_input_scales(src_rows);
std::vector<float> test_filter_scales(src_cols);
std::vector<int32_t> test_data;
test_data.reserve(src_rows * src_cols);
int outer_cols = src_cols / optimized_4bit::FilterWidth;
int outer_rows = src_rows;
for (int j = 0; j < outer_cols; ++j) {
for (int i = 0; i < outer_rows; ++i) {
for (int k = 0; k < optimized_4bit::FilterWidth; ++k) {
test_data.push_back(i);
}
}
}
for (int i = 0; i < src_rows; ++i) {
test_input_scales[i] = real_dist(random_engine);
}
for (int i = 0; i < src_cols; ++i) {
test_filter_scales[i] = real_dist(random_engine);
}
TestUnpack test(test_data, test_input_scales, test_filter_scales, src_rows,
src_cols, output_rows, output_cols);
test.Unpack<4, 1>();
std::vector<float> result = test.output_data;
for (int i = 0; i < output_rows; ++i) {
for (int j = 0; j < output_cols; ++j) {
float res = result[i * output_cols + j];
EXPECT_EQ(res, i * test_input_scales[i] * test_filter_scales[j]);
}
}
}
INSTANTIATE_TEST_SUITE_P(RunUnpackTests, RunUnpackTests,
::testing::ValuesIn({
std::make_tuple(1, 8, 5),
std::make_tuple(3, 4, 4),
}));
class RunKernelTests
: public ::testing::TestWithParam<::testing::tuple<int, int, int, int>> {};
TEST_P(RunKernelTests, RunKernelTests) {
auto params = GetParam();
int rhs_width = std::get<0>(params);
int lhs_layout_rows = std::get<1>(params);
int rhs_layout_rows = std::get<2>(params);
int lhs_layout_cols = std::get<3>(params);
int rhs_layout_cols = lhs_layout_cols;
std::vector<uint8_t> test_lhs(lhs_layout_rows * lhs_layout_cols / 2, 0.0);
std::vector<int8_t> test_rhs(rhs_layout_rows * rhs_layout_cols, 0.0);
std::vector<int32_t> test_accum(lhs_layout_rows * rhs_layout_rows, 0.0);
int lhs_outer_rows = lhs_layout_rows / optimized_4bit::FilterWidth;
int lhs_outer_cols = lhs_layout_cols / optimized_4bit::FilterDepth;
// pack lhs
for (int i = 0; i < lhs_outer_rows; ++i) {
for (int j = 0; j < lhs_outer_cols; ++j) {
for (int k = 0; k < optimized_4bit::FilterWidth; ++k) {
for (int l = 0; l < optimized_4bit::FilterDepth / 2; ++l) {
uint8_t u = static_cast<uint8_t>(int_dist(random_engine) + 7);
uint8_t v = static_cast<uint8_t>(int_dist(random_engine) + 7);
int lower = static_cast<uint8_t>(v) & UINT8_C(15);
int upper = static_cast<uint8_t>(u) << 4;
int cluster_index = (i * lhs_outer_cols + j) *
optimized_4bit::FilterDepth *
optimized_4bit::FilterWidth / 2;
int index = cluster_index + k * (optimized_4bit::FilterDepth / 2);
test_lhs[index + l] = (upper | lower);
}
}
}
}
int rhs_outer_rows = rhs_layout_rows / rhs_width;
int rhs_outer_cols = rhs_layout_cols / optimized_4bit::FilterDepth;
for (int i = 0; i < rhs_outer_rows; ++i) {
for (int j = 0; j < rhs_outer_cols; ++j) {
for (int k = 0; k < rhs_width; ++k) {
for (int l = 0; l < optimized_4bit::FilterDepth; ++l) {
int8_t u = static_cast<int8_t>(int_dist(random_engine));
int cluster_index = (i * rhs_outer_cols + j) *
optimized_4bit::FilterDepth * rhs_width;
int index = cluster_index + k * optimized_4bit::FilterDepth;
test_rhs[index + l] = u;
}
}
}
}
int index = 0;
std::vector<int32_t> expected_accum(lhs_layout_rows * rhs_layout_rows, 0.0);
int outer_cols = rhs_outer_cols;
for (int i = 0; i < lhs_outer_rows; ++i) {
for (int j = 0; j < rhs_outer_rows; ++j) {
for (int k = 0; k < rhs_width; ++k) {
for (int l = 0; l < optimized_4bit::FilterWidth; ++l) {
for (int m = 0; m < outer_cols; ++m) {
for (int n = 0; n < optimized_4bit::FilterDepth; ++n) {
int right_index = ((j * outer_cols + m) * rhs_width + k) *
optimized_4bit::FilterDepth;
int8_t rhs = test_rhs[right_index + n];
int left_index =
((i * outer_cols + m) * optimized_4bit::FilterWidth + l) *
(optimized_4bit::FilterDepth / 2);
uint8_t lhs = 0;
if (n < optimized_4bit::FilterDepth / 2) {
int a = n % 16;
lhs = static_cast<uint8_t>(test_lhs[left_index + a] >> 4);
} else {
int a = n % 16;
lhs = static_cast<uint8_t>(test_lhs[left_index + a] &
UINT8_C(15));
}
int accum_index = ((i * rhs_outer_rows + j) * rhs_width + k) *
optimized_4bit::FilterWidth;
expected_accum[accum_index + l] += rhs * lhs;
}
}
}
}
}
}
index = 0;
switch (rhs_width) {
#if defined(FC_4BIT_NEON) && defined(__aarch64__)
case 4:
optimized_4bit::RunKernel<optimized_4bit::FilterWidth, 4,
optimized_4bit::FilterDepth>(
test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
lhs_layout_rows);
break;
case 2:
optimized_4bit::RunKernel<optimized_4bit::FilterWidth, 2,
optimized_4bit::FilterDepth>(
test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
lhs_layout_rows);
break;
#endif
case 1:
[[fallthrough]];
default:
optimized_4bit::RunKernel<optimized_4bit::FilterWidth, 1,
optimized_4bit::FilterDepth>(
test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
lhs_layout_rows);
break;
}
for (int i = 0; i < (rhs_layout_rows * lhs_layout_rows); ++i) {
int32_t expected_val = expected_accum[i];
int32_t val = test_accum[i];
EXPECT_EQ(val, expected_val);
}
}
INSTANTIATE_TEST_SUITE_P(
RunKernelTests, RunKernelTests, ::testing::ValuesIn({
std::make_tuple(1, 4, 1, 32), std::make_tuple(1, 8, 1, 32),
std::make_tuple(1, 16, 1, 32), std::make_tuple(1, 4, 1, 64),
std::make_tuple(1, 8, 1, 64), std::make_tuple(1, 16, 1, 64),
std::make_tuple(1, 4, 5, 64), std::make_tuple(1, 8, 9, 64),
std::make_tuple(1, 16, 17, 64),
#if defined(FC_4BIT_NEON) && defined(__aarch64__)
std::make_tuple(2, 8, 2, 32), std::make_tuple(2, 16, 2, 32),
std::make_tuple(2, 4, 4, 64), std::make_tuple(2, 8, 4, 64),
std::make_tuple(2, 16, 4, 64), std::make_tuple(2, 4, 4, 64),
std::make_tuple(2, 8, 8, 64), std::make_tuple(2, 16, 16, 64),
std::make_tuple(4, 4, 4, 32), std::make_tuple(4, 8, 4, 32),
std::make_tuple(4, 16, 4, 32), std::make_tuple(4, 4, 8, 64),
std::make_tuple(4, 8, 8, 64), std::make_tuple(4, 16, 8, 64),
std::make_tuple(4, 4, 8, 64), std::make_tuple(4, 8, 12, 64),
std::make_tuple(4, 16, 32, 64),
#endif
}));
template <typename T>
class OverflowKernelTests : public ::testing::Test {};
using RhsWidths = ::testing::Types<std::integral_constant<int, 1>
#if defined(FC_4BIT_NEON) && defined(__aarch64__)
,
std::integral_constant<int, 2>,
std::integral_constant<int, 4>
#endif
>;
TYPED_TEST_SUITE(OverflowKernelTests, RhsWidths);
TYPED_TEST(OverflowKernelTests, OverflowTest) {
constexpr int rhs_width = TypeParam::value;
int lhs_layout_rows = 4;
int rhs_layout_rows = rhs_width;
int lhs_layout_cols = 1024;
int rhs_layout_cols = 1024;
// 0xff = two packed int4s with value of 15.
std::vector<uint8_t> test_lhs(lhs_layout_rows * lhs_layout_cols / 2, 0xff);
std::vector<int8_t> test_rhs(rhs_layout_rows * rhs_layout_cols, 127);
std::vector<int32_t> test_accum(lhs_layout_rows * rhs_layout_rows);
optimized_4bit::RunKernel<optimized_4bit::FilterWidth, rhs_width,
optimized_4bit::FilterDepth>(
test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
lhs_layout_rows);
for (int i = 0; i < (rhs_layout_rows * lhs_layout_rows); ++i) {
EXPECT_EQ(test_accum[i], 1950720);
}
}
} // namespace
} // namespace tflite
tensorflow/lite/kernels/internal/optimized/optimized_ops.h
+8126
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/optimized_ops_utils.h
+91
View File
@@ -0,0 +1,91 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
#include "Eigen/Core" // from @eigen_archive
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/lite/kernels/internal/runtime_shape.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_ops {
// Make a local VectorMap typedef allowing to map a float array
// as a Eigen vector expression. The std::conditional here is to
// construct the suitable Eigen type for the constness of the
// data. Indeed, for const data, we need to produce
// Eigen::Map<const Eigen::Matrix<float, ...>>
// and not the more straightforward
// Eigen::Map<Eigen::Matrix<const float, ...>>
template <typename Scalar>
using VectorMap = typename std::conditional<
std::is_const<Scalar>::value,
Eigen::Map<const Eigen::Matrix<typename std::remove_const<Scalar>::type,
Eigen::Dynamic, 1>>,
Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, 1>>>::type;
template <typename Scalar>
VectorMap<Scalar> MapAsVector(Scalar* data, const RuntimeShape& shape) {
const int size = shape.FlatSize();
return VectorMap<Scalar>(data, size, 1);
}
// Make a local VectorMap typedef allowing to map a float array
// as a Eigen matrix expression. The same explanation as for VectorMap
// above also applies here.
template <typename Scalar>
using MatrixMap = typename std::conditional<
std::is_const<Scalar>::value,
Eigen::Map<const Eigen::Matrix<typename std::remove_const<Scalar>::type,
Eigen::Dynamic, Eigen::Dynamic>>,
Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>>::type;
template <typename Scalar>
MatrixMap<Scalar> MapAsMatrixWithLastDimAsRows(Scalar* data,
const RuntimeShape& shape) {
const int dims_count = shape.DimensionsCount();
const int rows = shape.Dims(dims_count - 1);
const int cols = FlatSizeSkipDim(shape, dims_count - 1);
return MatrixMap<Scalar>(data, rows, cols);
}
template <typename Scalar>
MatrixMap<Scalar> MapAsMatrixWithFirstDimAsCols(Scalar* data,
const RuntimeShape& shape) {
const int cols = shape.Dims(0);
const int rows = FlatSizeSkipDim(shape, 0);
return MatrixMap<Scalar>(data, rows, cols);
}
template <typename Scalar>
using ArrayMap = typename std::conditional<
std::is_const<Scalar>::value,
Eigen::Map<const Eigen::Array<typename std::remove_const<Scalar>::type,
Eigen::Dynamic, Eigen::Dynamic>>,
Eigen::Map<Eigen::Array<Scalar, Eigen::Dynamic, Eigen::Dynamic>>>::type;
template <typename Scalar>
ArrayMap<Scalar> MapAsArrayWithLastDimAsRows(Scalar* data,
const RuntimeShape& shape) {
const int dims_count = shape.DimensionsCount();
const int rows = shape.Dims(dims_count - 1);
const int cols = FlatSizeSkipDim(shape, dims_count - 1);
return ArrayMap<Scalar>(data, rows, cols);
}
} // namespace optimized_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
tensorflow/lite/kernels/internal/optimized/reduce.h
+808
View File
@@ -0,0 +1,808 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
#include <stdint.h>
#include <algorithm>
#include <limits>
#include <vector>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops_utils.h"
#include "tensorflow/lite/kernels/internal/optimized/reduce_utils.h"
#include "tensorflow/lite/kernels/internal/reduce_common.h"
#include "tensorflow/lite/kernels/internal/reference/reduce.h"
#include "tensorflow/lite/kernels/internal/runtime_shape.h"
#include "tensorflow/lite/kernels/internal/types.h"
#include "tensorflow/lite/kernels/kernel_util.h"
namespace tflite {
namespace optimized_ops {
inline void MeanImpl(const tflite::MeanParams& op_params,
const RuntimeShape& input_shape, const uint8_t* input_data,
int32 multiplier, int32 shift, int32 bias,
const RuntimeShape& output_shape, uint8_t* output_data,
int start_depth, int end_depth) {
ruy::profiler::ScopeLabel label("Mean4D/Uint8/MeanImpl");
// Current implementation only supports dimension equals 4 and simultaneous
// reduction over width and height.
const int output_batch = output_shape.Dims(0);
const int output_height = output_shape.Dims(2);
const int output_width = output_shape.Dims(2);
const int input_height = input_shape.Dims(1);
const int input_width = input_shape.Dims(2);
TFLITE_CHECK_EQ(op_params.axis_count, 2);
TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
(op_params.axis[0] == 2 && op_params.axis[1] == 1));
TFLITE_CHECK_EQ(output_height, 1);
TFLITE_CHECK_EQ(output_width, 1);
constexpr int32_t kMinValue = std::numeric_limits<uint8_t>::min();
constexpr int32_t kMaxValue = std::numeric_limits<uint8_t>::max();
#ifdef USE_NEON
const int32x4_t bias_dup = vdupq_n_s32(bias);
const int32x4_t min_dup = vdupq_n_s32(kMinValue);
const int32x4_t max_dup = vdupq_n_s32(kMaxValue);
#endif // USE_NEON
for (int out_b = 0; out_b < output_batch; ++out_b) {
int out_d = start_depth;
#ifdef USE_NEON
for (; out_d <= end_depth - 16; out_d += 16) {
int32x4x4_t temp_sum;
temp_sum.val[0] = vdupq_n_s32(0);
temp_sum.val[1] = vdupq_n_s32(0);
temp_sum.val[2] = vdupq_n_s32(0);
temp_sum.val[3] = vdupq_n_s32(0);
for (int in_h = 0; in_h < input_height; ++in_h) {
for (int in_w = 0; in_w < input_width; ++in_w) {
const uint8_t* input_data_ptr =
input_data + Offset(input_shape, out_b, in_h, in_w, out_d);
uint8x16_t input_data_val = vld1q_u8(input_data_ptr);
int16x8_t input_data_low_shift =
vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(input_data_val)));
int16x8_t input_data_high_shift =
vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(input_data_val)));
int32x4_t input_low_low =
vmovl_s16(vget_low_s16(input_data_low_shift));
int32x4_t input_high_low =
vmovl_s16(vget_high_s16(input_data_low_shift));
int32x4_t input_low_high =
vmovl_s16(vget_low_s16(input_data_high_shift));
int32x4_t input_high_high =
vmovl_s16(vget_high_s16(input_data_high_shift));
temp_sum.val[0] = vaddq_s32(temp_sum.val[0], input_low_low);
temp_sum.val[1] = vaddq_s32(temp_sum.val[1], input_high_low);
temp_sum.val[2] = vaddq_s32(temp_sum.val[2], input_low_high);
temp_sum.val[3] = vaddq_s32(temp_sum.val[3], input_high_high);
}
}
temp_sum =
MultiplyByQuantizedMultiplier4Rows(temp_sum, multiplier, shift);
temp_sum.val[0] = vaddq_s32(temp_sum.val[0], bias_dup);
temp_sum.val[1] = vaddq_s32(temp_sum.val[1], bias_dup);
temp_sum.val[2] = vaddq_s32(temp_sum.val[2], bias_dup);
temp_sum.val[3] = vaddq_s32(temp_sum.val[3], bias_dup);
temp_sum.val[0] = vminq_s32(vmaxq_s32(temp_sum.val[0], min_dup), max_dup);
temp_sum.val[1] = vminq_s32(vmaxq_s32(temp_sum.val[1], min_dup), max_dup);
temp_sum.val[2] = vminq_s32(vmaxq_s32(temp_sum.val[2], min_dup), max_dup);
temp_sum.val[3] = vminq_s32(vmaxq_s32(temp_sum.val[3], min_dup), max_dup);
uint16x4_t narrowed_low_low =
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[0]));
uint16x4_t narrowed_high_low =
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[1]));
uint16x4_t narrowed_low_high =
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[2]));
uint16x4_t narrowed_high_high =
vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[3]));
uint16x8_t combined_low =
vcombine_u16(narrowed_low_low, narrowed_high_low);
uint16x8_t combined_high =
vcombine_u16(narrowed_low_high, narrowed_high_high);
uint8x8_t narrowed_low = vmovn_u16(combined_low);
uint8x8_t narrowed_high = vmovn_u16(combined_high);
uint8x16_t combined_output = vcombine_u8(narrowed_low, narrowed_high);
uint8_t* output_data_ptr =
output_data + Offset(output_shape, out_b, 0, 0, out_d);
vst1q_u8(output_data_ptr, combined_output);
}
#endif // USE_NEON
for (; out_d < end_depth; ++out_d) {
int acc = 0;
for (int in_h = 0; in_h < input_height; ++in_h) {
for (int in_w = 0; in_w < input_width; ++in_w) {
acc += input_data[Offset(input_shape, out_b, in_h, in_w, out_d)];
}
}
acc = MultiplyByQuantizedMultiplier(acc, multiplier, shift);
acc += bias;
acc = std::min(std::max(acc, kMinValue), kMaxValue);
output_data[Offset(output_shape, out_b, 0, 0, out_d)] =
static_cast<uint8_t>(acc);
}
}
}
struct MeanWorkerTask : cpu_backend_threadpool::Task {
MeanWorkerTask(const tflite::MeanParams& op_params,
const RuntimeShape& input_shape, const uint8_t* input_data,
int32 multiplier, int32 shift, int32 bias,
const RuntimeShape& output_shape, uint8_t* output_data,
int start_height, int end_height)
: op_params(op_params),
input_shape(input_shape),
input_data(input_data),
multiplier(multiplier),
shift(shift),
bias(bias),
output_shape(output_shape),
output_data(output_data),
start_height(start_height),
end_height(end_height) {}
void Run() override {
MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
output_shape, output_data, start_height, end_height);
}
private:
const tflite::MeanParams& op_params;
const RuntimeShape& input_shape;
const uint8_t* input_data;
int32 multiplier;
int32 shift;
int32 bias;
const RuntimeShape& output_shape;
uint8_t* output_data;
int start_height;
int end_height;
};
inline void Mean(const tflite::MeanParams& op_params,
const RuntimeShape& unextended_input_shape,
const uint8_t* input_data, int32 input_zero_point,
float input_scale, const RuntimeShape& unextended_output_shape,
uint8_t* output_data, int32 output_zero_point,
float output_scale, CpuBackendContext* cpu_backend_context) {
ruy::profiler::ScopeLabel label("Mean4D/Uint8");
// Current implementation only supports dimension equals 4 and simultaneous
// reduction over width and height.
TFLITE_CHECK_EQ(unextended_input_shape.DimensionsCount(), 4);
TFLITE_CHECK_LE(unextended_output_shape.DimensionsCount(), 4);
const RuntimeShape input_shape =
RuntimeShape::ExtendedShape(4, unextended_input_shape);
const RuntimeShape output_shape =
RuntimeShape::ExtendedShape(4, unextended_output_shape);
const int output_height = output_shape.Dims(1);
const int output_width = output_shape.Dims(2);
const int output_depth = output_shape.Dims(3);
TFLITE_CHECK_EQ(op_params.axis_count, 2);
TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
(op_params.axis[0] == 2 && op_params.axis[1] == 1));
TFLITE_CHECK_EQ(output_height, 1);
TFLITE_CHECK_EQ(output_width, 1);
const int input_height = input_shape.Dims(1);
const int input_width = input_shape.Dims(2);
const float num_elements_in_axis = input_width * input_height;
float temp = input_zero_point * input_scale / output_scale;
temp = temp > 0 ? temp + 0.5f : temp - 0.5f;
int32_t bias = output_zero_point - static_cast<int32_t>(temp);
float real_scale = input_scale / (num_elements_in_axis * output_scale);
int32 multiplier, shift;
QuantizeMultiplier(real_scale, &multiplier, &shift);
constexpr int kMinDepthPerThread = 8;
int thread_count = output_depth / kMinDepthPerThread;
thread_count = thread_count > 0 ? thread_count : 1;
const int capped_thread_count =
std::min(thread_count, cpu_backend_context->max_num_threads());
if (capped_thread_count == 1) {
MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
output_shape, output_data, 0, output_depth);
} else {
// Instead parallel for batch, we loop for the output_depth since batch
// is typical 1.
std::vector<MeanWorkerTask> tasks;
// TODO(b/131746020) don't create new heap allocations every time.
// At least we make it a single heap allocation by using reserve().
tasks.reserve(capped_thread_count);
int depth_start = 0;
for (int i = 0; i < capped_thread_count; ++i) {
// Try to distribute the tasks as even as possible.
int depth_end = depth_start +
(output_depth - depth_start) / (capped_thread_count - i);
tasks.emplace_back(op_params, input_shape, input_data, multiplier, shift,
bias, output_shape, output_data, depth_start,
depth_end);
depth_start = depth_end;
}
cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
cpu_backend_context);
}
}
template <typename T>
struct SumOp {
inline T operator()(const T& a) const { return a; }
inline T operator()(const T& a, const T& b) const { return a + b; }
static constexpr T kNeutralElement = T(0);
};
template <typename T, typename U>
struct CastSumOp {
inline U operator()(const T& a) const { return static_cast<U>(a); }
inline U operator()(const U& a, const T& b) const {
return a + static_cast<U>(b);
}
static constexpr U kNeutralElement = U(0);
};
template <typename T>
struct ProdOp {
inline T operator()(const T& a) const { return a; }
inline T operator()(const T& a, const T& b) const { return a * b; }
static constexpr T kNeutralElement = T(1);
};
template <typename T>
struct MaxOp {
inline T operator()(const T& a) const { return a; }
inline T operator()(const T& a, const T& b) const { return (a > b) ? a : b; }
static constexpr T kNeutralElement = std::numeric_limits<T>::lowest();
};
template <typename T>
struct MinOp {
inline T operator()(const T& a) const { return a; }
inline T operator()(const T& a, const T& b) const { return (a < b) ? a : b; }
static constexpr T kNeutralElement = std::numeric_limits<T>::max();
};
struct AndOp {
inline bool operator()(bool a) const { return a; }
inline bool operator()(bool a, bool b) const { return a && b; }
static constexpr bool kNeutralElement = true;
};
struct OrOp {
inline bool operator()(bool a) const { return a; }
inline bool operator()(bool a, bool b) const { return a || b; }
static constexpr bool kNeutralElement = false;
};
// When the number of axis is zero, the reduction is simply a copy.
template <typename T>
void ReduceIsCopy(const T* input_data, const int* input_dims,
const int input_num_dims, T* output_data) {
int num_elems = NumElements(input_dims, input_num_dims);
memcpy(output_data, input_data, num_elems * sizeof(T));
}
// Reduces the input over either odd or even dimensions using Op.
// One recursive call for each dimension is made.
// 'depth' is the depth of recursion.
// 'parity' indicates whether odd or even dimensions are being reduced.
// ReducerFirst is applied to the first element to be written to each output
// position.
// ReducerNext is applied to each subsequent element to be written to each
// output position.
template <typename T, typename U, typename ReducerFirst, typename ReducerNext>
inline std::pair<const T*, U*> ReduceImpl(const T* input_data,
const int* input_dims, U* output_data,
int depth, int parity, bool next,
const ReducerFirst& reducer_first,
const ReducerNext& reducer_next) {
// The output pointer is incremented conditionally depending on whether the
// odd or even dimension is being reduced.
// The input pointer is always incremented as each input is read once.
if (depth > 0) {
U* future_output = output_data;
bool update_output = (depth % 2) == parity;
for (int i = 0; i < input_dims[0]; ++i) {
if (i > 0 && !update_output) {
next = true;
}
std::tie(input_data, future_output) =
ReduceImpl(input_data, &input_dims[1], output_data, depth - 1, parity,
next, reducer_first, reducer_next);
if (update_output) {
output_data = future_output;
}
}
output_data = future_output;
} else {
// Reduce the final dimension.
if (parity) {
// Reduce the even dimension. The entire dimension is reduced into one
// value.
U res = next ? reducer_next(*output_data, *input_data++)
: reducer_first(*input_data++);
for (int i = 1; i < input_dims[0]; ++i) {
res = reducer_next(res, *input_data++);
}
*output_data++ = res;
} else {
// Reduce the odd dimension. Each input is accumulated into a separate
// output.
if (!next) {
for (int i = 0; i < input_dims[0]; ++i) {
U res = reducer_first(*input_data++);
*output_data++ = res;
}
} else {
for (int i = 0; i < input_dims[0]; ++i) {
U res = *output_data;
res = reducer_next(res, *input_data++);
*output_data++ = res;
}
}
}
}
return {input_data, output_data};
}
// A generic reduce method that can be used for reduce_sum, reduce_mean, etc.
// This method iterates through input data and reduce elements along the
// dimensions given in axis. ReducerFirst is used the first time each output
// element is written and ReducerNext is used for all subsequent writes.
template <typename In, typename Out, typename ReducerFirst,
typename ReducerNext>
inline bool Reduce(const In* input_data, const int* input_dims,
const int input_num_dims, const int* axis,
const int num_axis, Out* output_data,
const ReducerFirst& reducer_first,
const ReducerNext& reducer_next) {
const int parity = (axis[num_axis - 1] == input_num_dims - 1) ? 1 : 0;
ReduceImpl(input_data, input_dims, output_data, input_num_dims - 1, parity,
/*next=*/false, reducer_first, reducer_next);
return true;
}
// Computes the mean or sum of elements across dimensions given in axis.
// It does so in two stages, first calculates the sum of elements along the axis
// then divides it by the number of element in axis for quantized values.
template <typename T, typename U>
bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point,
float input_scale, const int* input_dims,
const int input_num_dims, T* output_data,
int32_t output_zero_point, float output_scale,
const int* output_dims, const int output_num_dims,
const int* axis, const int num_axis_dimensions,
bool keep_dims, int* normalized_dims,
int* resolved_axis, U* temp_sum, bool compute_sum) {
const int32_t kMinValue = std::numeric_limits<T>::min();
const int32_t kMaxValue = std::numeric_limits<T>::max();
ruy::profiler::ScopeLabel label(compute_sum ? "QuantizedSum"
: "QuantizedMean");
// Reset output data.
size_t num_outputs = 1;
for (int idx = 0; idx < output_num_dims; ++idx) {
size_t current = static_cast<size_t>(output_dims[idx]);
// Overflow prevention.
if (num_outputs > std::numeric_limits<size_t>::max() / current) {
return false;
}
num_outputs *= current;
}
// Return early when input shape has zero dim. This is done after initializing
// data for output tensor because there are cases that the input tensor is
// empty but output tensor is not. In that case, output tensor should be
// filled with init_value.
for (int i = 0; i < input_num_dims; ++i) {
if (input_dims[i] == 0) return true;
}
// Resolve axis.
int num_resolved_axis = 0;
int normalized_num_dims = 0;
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
resolved_axis, num_resolved_axis, input_dims,
normalized_dims, normalized_num_dims)) {
return false;
}
if (num_resolved_axis == 0) {
int count = NumElements(input_dims, input_num_dims);
for (int i = 0; i < count; ++i) {
temp_sum[i] = U(input_data[i]);
}
} else {
if (!Reduce<T, U, CastSumOp<T, U>, CastSumOp<T, U>>(
input_data, normalized_dims, normalized_num_dims, resolved_axis,
num_resolved_axis, temp_sum, CastSumOp<T, U>(),
CastSumOp<T, U>())) {
return false;
}
}
// Calculate mean by dividing output_data by num of aggregated element.
size_t num_elements_in_axis = 1;
for (int idx = 0; idx < num_resolved_axis; ++idx) {
size_t current = static_cast<size_t>(normalized_dims[resolved_axis[idx]]);
// Overflow prevention.
if (current > (std::numeric_limits<size_t>::max() / num_elements_in_axis)) {
return false;
}
num_elements_in_axis *= current;
}
if (num_elements_in_axis > 0) {
const float scale = input_scale / output_scale;
if (compute_sum) {
const float bias = -input_zero_point * scale * num_elements_in_axis;
for (size_t idx = 0; idx < num_outputs; ++idx) {
U value = static_cast<U>(TfLiteRound(temp_sum[idx] * scale + bias)) +
output_zero_point;
value = std::min(std::max(value, kMinValue), kMaxValue);
output_data[idx] = static_cast<T>(value);
}
} else {
const float bias = -input_zero_point * scale;
for (size_t idx = 0; idx < num_outputs; ++idx) {
float float_mean = static_cast<float>(temp_sum[idx]) /
static_cast<float>(num_elements_in_axis);
float result = TfLiteMin(
TfLiteRound(float_mean * scale + bias) + output_zero_point,
static_cast<float>(std::numeric_limits<T>::max()));
result = TfLiteMax(result,
static_cast<float>(std::numeric_limits<T>::min()));
output_data[idx] = static_cast<T>(result);
}
}
}
return true;
}
using ops::builtin::reduce::ReduceType;
template <typename T>
inline bool ReduceDispatcher(const T* input_data, const int* input_dims,
const int input_num_dims, const int* output_dims,
int output_num_dims, T* output_data,
const int* axis, const int64_t num_axis_dimensions,
ReduceType reduce_type) {
T init_value;
switch (reduce_type) {
case ReduceType::kProd:
init_value = ProdOp<T>::kNeutralElement;
break;
case ReduceType::kSum:
init_value = SumOp<T>::kNeutralElement;
break;
case ReduceType::kMin:
init_value = MinOp<T>::kNeutralElement;
break;
case ReduceType::kMax:
init_value = MaxOp<T>::kNeutralElement;
break;
default:
return false;
}
// Return early when input shape has zero dim. This is done after initializing
// data for output tensor because there are cases that the input tensor is
// empty but output tensor is not. In that case, output tensor should be
// filled with Op::kNeutralElement.
for (int i = 0; i < input_num_dims; ++i) {
if (input_dims[i] == 0) {
return reference_ops::InitTensorDataForReduce(
output_dims, output_num_dims, init_value, output_data);
}
}
switch (reduce_type) {
case ReduceType::kProd:
return Reduce<T, T, ProdOp<T>, ProdOp<T>>(
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
output_data, ProdOp<T>(), ProdOp<T>());
case ReduceType::kSum:
return Reduce<T, T, SumOp<T>, SumOp<T>>(
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
output_data, SumOp<T>(), SumOp<T>());
case ReduceType::kMin:
return Reduce<T, T, MinOp<T>, MinOp<T>>(
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
output_data, MinOp<T>(), MinOp<T>());
case ReduceType::kMax:
return Reduce<T, T, MaxOp<T>, MaxOp<T>>(
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
output_data, MaxOp<T>(), MaxOp<T>());
default:
return false;
}
}
template <>
inline bool ReduceDispatcher<bool>(const bool* input_data,
const int* input_dims,
const int input_num_dims,
const int* output_dims, int output_num_dims,
bool* output_data, const int* axis,
const int64_t num_axis_dimensions,
ReduceType reduce_type) {
bool init_value;
switch (reduce_type) {
case ReduceType::kAny:
init_value = OrOp::kNeutralElement;
break;
case ReduceType::kAll:
init_value = AndOp::kNeutralElement;
break;
default:
return false;
}
// Return early when input shape has zero dim. This is done after initializing
// data for output tensor because there are cases that the input tensor is
// empty but output tensor is not. In that case, output tensor should be
// filled with Op::kNeutralElement.
for (int i = 0; i < input_num_dims; ++i) {
if (input_dims[i] == 0) {
return reference_ops::InitTensorDataForReduce(
output_dims, output_num_dims, init_value, output_data);
}
}
switch (reduce_type) {
case ReduceType::kAll:
return Reduce<bool, bool, AndOp, AndOp>(
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
output_data, AndOp(), AndOp());
case ReduceType::kAny:
return Reduce<bool, bool, OrOp, OrOp>(
input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
output_data, OrOp(), OrOp());
default:
return false;
}
}
// Calculate the reduced product by rescaling each multiplication step to
// avoid an overflow.
template <typename T>
struct ReducerFirst {
explicit ReducerFirst(int input_zero_point_arg)
: input_zero_point(input_zero_point_arg) {}
int32_t operator()(T in) const { return in - input_zero_point; }
int input_zero_point;
};
template <typename T>
struct ReducerNext {
ReducerNext(int32_t input_zero_point_arg, int32_t scaling_multiplier_arg,
int32_t scaling_shift_arg)
: input_zero_point(input_zero_point_arg),
scaling_multiplier(scaling_multiplier_arg),
scaling_shift(scaling_shift_arg) {}
int32_t operator()(int32_t current, T in) const {
const int64_t result =
static_cast<int64_t>(current) * (in - input_zero_point);
return MultiplyByQuantizedMultiplier(result, scaling_multiplier,
scaling_shift);
}
int32_t input_zero_point, scaling_multiplier, scaling_shift;
};
template <typename T>
inline bool QuantizedReduceProd(
const T* input_data, int32_t input_zero_point,
const RuntimeShape& input_shape, T* output_data, int32_t output_zero_point,
const RuntimeShape& output_shape, const int* axis,
const int64_t num_axis_dimensions, int* resolved_axis, int* normalized_dims,
int32_t* temp_prod, int32_t scaling_multiplier, int scaling_shift) {
const int32_t kMinValue = std::numeric_limits<T>::min();
const int32_t kMaxValue = std::numeric_limits<T>::max();
// Resolve axis.
int num_resolved_axis = 0;
int normalized_num_dims = 0;
if (!reduce_utils::ResolveAxis(input_shape.DimensionsCount(), axis,
num_axis_dimensions, resolved_axis,
num_resolved_axis, input_shape.DimsData(),
normalized_dims, normalized_num_dims)) {
return false;
}
if (!Reduce<T, int32_t, ReducerFirst<T>, ReducerNext<T>>(
input_data, normalized_dims, normalized_num_dims, resolved_axis,
num_resolved_axis, temp_prod, ReducerFirst<T>(input_zero_point),
ReducerNext<T>(input_zero_point, scaling_multiplier,
scaling_shift))) {
return false;
}
for (int i = 0; i < output_shape.FlatSize(); i++) {
int32_t result =
MultiplyByQuantizedMultiplier(static_cast<int64_t>(temp_prod[i]),
scaling_multiplier, scaling_shift) +
output_zero_point;
result = std::min(std::max(result, kMinValue), kMaxValue);
output_data[i] = static_cast<T>(result);
}
return true;
}
template <typename T>
inline void Mean(const tflite::MeanParams& op_params,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& output_shape, T* output_data) {
return reference_ops::Mean(op_params, input_shape, input_data, output_shape,
output_data);
}
// Computes the mean of elements across dimensions given in axis.
// It does so in two stages, first calculates the sum of elements along the axis
// then divides it by the number of element in axis.
template <typename T, typename U>
inline bool MeanGeneral(const T* input_data, const int* input_dims,
const int input_num_dims, T* output_data,
const int* output_dims, const int output_num_dims,
const int* axis, const int num_axis_dimensions,
bool keep_dims, int* normalized_dims,
int* resolved_axis, U* temp_sum) {
ruy::profiler::ScopeLabel label("Mean");
// Resolve axis.
int num_resolved_axis = 0;
int normalized_num_dims = 0;
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
resolved_axis, num_resolved_axis, input_dims,
normalized_dims, normalized_num_dims)) {
return false;
}
if (num_resolved_axis == 0) {
optimized_ops::ReduceIsCopy(input_data, input_dims, input_num_dims,
output_data);
return true;
}
// Reset output data.
size_t num_outputs = 1;
for (int idx = 0; idx < output_num_dims; ++idx) {
size_t current = static_cast<size_t>(output_dims[idx]);
// Overflow prevention.
if (num_outputs > std::numeric_limits<size_t>::max() / current) {
return false;
}
num_outputs *= current;
}
if (!Reduce<T, U, CastSumOp<T, U>, CastSumOp<T, U>>(
input_data, normalized_dims, normalized_num_dims, resolved_axis,
num_resolved_axis, temp_sum, CastSumOp<T, U>(), CastSumOp<T, U>())) {
return false;
}
// Calculate mean by dividing output_data by num of aggregated element.
size_t num_elements_in_axis = 1;
for (int idx = 0; idx < num_resolved_axis; ++idx) {
size_t current = static_cast<size_t>(normalized_dims[resolved_axis[idx]]);
// Overflow prevention.
if (current > (std::numeric_limits<size_t>::max() / num_elements_in_axis)) {
return false;
}
num_elements_in_axis *= current;
}
if (num_elements_in_axis > 0) {
for (size_t idx = 0; idx < num_outputs; ++idx) {
output_data[idx] =
static_cast<T>(temp_sum[idx] / static_cast<U>(num_elements_in_axis));
}
}
return true;
}
template <typename T, typename U>
inline bool Mean(const T* input_data, const int* input_dims,
const int input_num_dims, T* output_data,
const int* output_dims, const int output_num_dims,
const int* axis, const int num_axis_dimensions, bool keep_dims,
int* normalized_dims, int* resolved_axis, U* temp_sum) {
return MeanGeneral(input_data, input_dims, input_num_dims, output_data,
output_dims, output_num_dims, axis, num_axis_dimensions,
false, normalized_dims, resolved_axis, temp_sum);
}
// Use Eigen when Mean is calculated over the last dimension only of a float
// tensor.
template <>
inline bool Mean<float, float>(const float* input_data, const int* input_dims,
const int input_num_dims, float* output_data,
const int* output_dims,
const int output_num_dims, const int* axis,
const int num_axis_dimensions, bool keep_dims,
int* normalized_dims, int* resolved_axis,
float* temp_sum) {
// Handle reduce_mean for the last dimensions.
int num_resolved_axis = 0;
int normalized_num_dims = 0;
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
resolved_axis, num_resolved_axis, input_dims,
normalized_dims, normalized_num_dims)) {
return false;
}
if (normalized_num_dims > 1 && num_resolved_axis == 1 &&
resolved_axis[0] == (normalized_num_dims - 1)) {
ruy::profiler::ScopeLabel label("MeanLastDim/Float");
int output_size = normalized_dims[0];
const int last_input_dim = normalized_dims[1];
// TODO(b/152563685): Consider use eigen to cover more general cases.
const MatrixMap<const float> in_mat(input_data, last_input_dim,
output_size);
VectorMap<float> out(output_data, output_size, 1);
out = (in_mat.array().colwise().sum()) / static_cast<float>(last_input_dim);
return true;
}
return MeanGeneral(input_data, input_dims, input_num_dims, output_data,
output_dims, output_num_dims, axis, num_axis_dimensions,
false, normalized_dims, resolved_axis, temp_sum);
}
// Computes the generic value (i.e., sum/max/min/prod) of elements across
// dimensions given in axis. It needs to pass in init_value and reducer.
template <typename T>
inline bool ReduceGeneric(const T* input_data, const int* input_dims,
const int input_num_dims, T* output_data,
const int* output_dims, const int output_num_dims,
const int* axis, const int64_t num_axis_dimensions,
int* resolved_axis, int* normalized_dims,
ReduceType reduce_type) {
int num_resolved_axis = 0;
int normalized_num_dims = 0;
if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
resolved_axis, num_resolved_axis, input_dims,
normalized_dims, normalized_num_dims)) {
return false;
}
if (num_resolved_axis == 0) {
optimized_ops::ReduceIsCopy(input_data, input_dims, input_num_dims,
output_data);
return true;
}
return ReduceDispatcher(input_data, normalized_dims, normalized_num_dims,
output_dims, output_num_dims, output_data,
resolved_axis, num_resolved_axis, reduce_type);
}
} // namespace optimized_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
tensorflow/lite/kernels/internal/optimized/reduce_utils.h
+138
View File
@@ -0,0 +1,138 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_UTILS_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_UTILS_H_
#include <stdint.h>
#include <algorithm>
#include <cstring>
namespace tflite {
namespace reduce_utils {
inline void RemoveSize1Dims(int* shape_out, int& out_num_dims, int* axis_out,
int& out_num_axis) {
for (int64_t i = 0; i < out_num_dims;) {
if (shape_out[i] == 1) {
for (int64_t j = i + 1; j < out_num_dims; ++j) {
shape_out[j - 1] = shape_out[j];
}
for (int64_t j = 0; j < out_num_axis; ++j) {
if (axis_out[j] == i) {
for (int64_t k = j + 1; k < out_num_axis; ++k) {
axis_out[k - 1] = axis_out[k];
}
out_num_axis -= 1;
break;
}
}
for (int64_t j = 0; j < out_num_axis; ++j) {
if (axis_out[j] > i) {
axis_out[j] -= 1;
}
}
--out_num_dims;
} else {
++i;
}
}
}
// This method parses the input 'axis' to remove duplicates, handle negative
// values and remove redundant dimensions. It returns a valid 'axis_out' and
// 'shape_out' contains the flattened input shape. 'out_num_dims' contains the
// reduced number of dimensions.
inline bool ResolveAxis(const int num_dims, const int* axis,
const int64_t num_axis, int* axis_out,
int& out_num_axis, const int* shape_in, int* shape_out,
int& out_num_dims) {
// Short-circuit axis resolution for scalars; the axis will go unused.
if (num_dims == 0) {
out_num_axis = 0;
out_num_dims = 0;
return true;
}
out_num_axis = 0;
out_num_dims = num_dims;
// o(n^2) is fine since out_num_axis should be really small, mostly <= 4
for (int64_t idx = 0; idx < num_axis; ++idx) {
// Handle negative index. A positive index 'p_idx' can be represented as a
// negative index 'n_idx' as: n_idx = p_idx-num_dims
// eg: For num_dims=3, [0, 1, 2] is the same as [-3, -2, -1] */
int current = axis[idx] < 0 ? (axis[idx] + num_dims) : axis[idx];
if (current < 0 || current >= num_dims) {
return false;
}
bool is_dup = false;
for (int j = 0; j < out_num_axis; ++j) {
if (axis_out[j] == current) {
is_dup = true;
break;
}
}
if (!is_dup) {
axis_out[out_num_axis] = current;
out_num_axis += 1;
}
}
// If two or more adjacent dimensions are either reduced
// over or not, then the second and subsequent dimensions may be flattened.
memcpy(shape_out, shape_in, num_dims * sizeof(int));
std::sort(&axis_out[0], &axis_out[out_num_axis]);
RemoveSize1Dims(shape_out, out_num_dims, axis_out, out_num_axis);
if (out_num_axis > 0) {
int64_t j = out_num_axis - 1;
// true if the previous index is present in axis_out.
bool previous_here = (axis_out[j] == out_num_dims - 1);
if (previous_here) {
j -= 1;
}
for (int64_t i = out_num_dims - 2; i >= 0; --i) {
// true if the current index is present in axis_out.
bool current_here = j >= 0 ? (axis_out[j] == i) : false;
if (current_here == previous_here) {
shape_out[i] *= shape_out[i + 1];
for (int64_t k = i + 1; k + 1 < out_num_dims; ++k) {
shape_out[k] = shape_out[k + 1];
}
// All axis bigger than this need to be reduced by 1.
for (int64_t k = 0; k < out_num_axis; ++k) {
if (axis_out[k] > i) {
axis_out[k] -= 1;
}
}
if (current_here) {
for (int64_t k = j + 1; k + 1 < out_num_axis; ++k) {
axis_out[k] = axis_out[k + 1];
}
out_num_axis -= 1;
}
out_num_dims -= 1;
}
if (current_here) {
j -= 1;
}
previous_here = current_here;
}
}
return true;
}
} // namespace reduce_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_UTILS_H_
tensorflow/lite/kernels/internal/optimized/reduce_utils_test.cc
+137
View File
@@ -0,0 +1,137 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/lite/kernels/internal/optimized/reduce_utils.h"
#include <gmock/gmock.h>
namespace tflite {
namespace reduce_utils {
namespace {
using ::testing::ElementsAreArray;
void TestFunction(const std::vector<int>& axis_in,
const std::vector<int>& shape_in,
const std::vector<int>& expected_axis_out,
const std::vector<int>& expected_shape_out) {
int num_dims = shape_in.size();
int expected_out_num_dims = expected_shape_out.size();
int actual_out_num_dims;
int expected_out_num_axis = expected_axis_out.size();
int actual_out_num_axis;
std::vector<int> actual_shape_out(num_dims);
std::vector<int> actual_axis_out(num_dims);
ResolveAxis(shape_in.size(), axis_in.data(), axis_in.size(),
actual_axis_out.data(), actual_out_num_axis, shape_in.data(),
actual_shape_out.data(), actual_out_num_dims);
EXPECT_EQ(expected_out_num_dims, actual_out_num_dims);
EXPECT_EQ(expected_out_num_axis, actual_out_num_axis);
EXPECT_THAT(expected_shape_out,
ElementsAreArray(actual_shape_out.data(), expected_out_num_dims));
EXPECT_THAT(expected_axis_out,
ElementsAreArray(actual_axis_out.data(), expected_out_num_axis));
}
TEST(ResolveAxisTest, Flatten_0_1_2) {
const std::vector<int> axis_in = {0, 1, 2};
const std::vector<int> shape_in = {2, 3, 4, 5};
const std::vector<int> expected_shape_out{24, 5};
const std::vector<int> expected_axis_out{0};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, Flatten_0_1_2_3) {
const std::vector<int> axis_in = {3, 2};
const std::vector<int> shape_in = {2, 3, 4, 5};
const std::vector<int> expected_shape_out{6, 20};
const std::vector<int> expected_axis_out{1};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, ZeroDims) {
const std::vector<int> axis_in = {};
const std::vector<int> shape_in = {};
const std::vector<int> expected_shape_out{};
const std::vector<int> expected_axis_out{};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, DoNothing) {
const std::vector<int> axis_in = {0};
const std::vector<int> shape_in = {4, 5};
const std::vector<int> expected_shape_out{4, 5};
const std::vector<int> expected_axis_out{0};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, NegativeAxis) {
const std::vector<int> axis_in = {-2};
const std::vector<int> shape_in = {4, 3};
const std::vector<int> expected_shape_out{4, 3};
const std::vector<int> expected_axis_out{0};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, NegativeAxisFold) {
const std::vector<int> axis_in = {-1};
const std::vector<int> shape_in = {4, 3, 5};
const std::vector<int> expected_shape_out{12, 5};
const std::vector<int> expected_axis_out{1};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, DuplicateAxis) {
const std::vector<int> axis_in = {2, 1, 2, 1, 2, 1};
const std::vector<int> shape_in = {4, 3, 2};
const std::vector<int> expected_shape_out{4, 6};
const std::vector<int> expected_axis_out{1};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, DuplicateNegativeAxis) {
const std::vector<int> axis_in = {2, -1, -2, -1, 2, 1};
const std::vector<int> shape_in = {4, 3, 2};
const std::vector<int> expected_shape_out{4, 6};
const std::vector<int> expected_axis_out{1};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, RemoveSize1Dim) {
const std::vector<int> axis_in = {0};
const std::vector<int> shape_in = {1, 4, 3, 1};
const std::vector<int> expected_shape_out{4, 3};
const std::vector<int> expected_axis_out{};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, OneSize1DimToScalar) {
const std::vector<int> axis_in = {0};
const std::vector<int> shape_in = {1};
const std::vector<int> expected_shape_out{};
const std::vector<int> expected_axis_out{};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
TEST(ResolveAxisTest, InterleavedSize1Dim) {
const std::vector<int> axis_in = {1, 3};
const std::vector<int> shape_in = {1, 2, 1, 4, 1, 7};
const std::vector<int> expected_shape_out{8, 7};
const std::vector<int> expected_axis_out{0};
TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
}
} // namespace
} // namespace reduce_utils
} // namespace tflite
tensorflow/lite/kernels/internal/optimized/resize_bilinear.h
+1736
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/kernels/internal/optimized/sparse_ops/fully_connected.h
+268
View File
@@ -0,0 +1,268 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SPARSE_OPS_FULLY_CONNECTED_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SPARSE_OPS_FULLY_CONNECTED_H_
#include <algorithm>
#include <cstdint>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/cppmath.h"
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/lite/kernels/internal/quantization_util.h"
#include "tensorflow/lite/kernels/internal/tensor_utils.h"
#include "tensorflow/lite/kernels/internal/types.h"
namespace tflite {
namespace optimized_ops {
inline void FullyConnectedSparseWeight(
const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& weights_shape, const float* weights_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data) {
ruy::profiler::ScopeLabel label("FullyConnected");
ruy::profiler::ScopeLabel inner_label("Random Sparse");
const float output_activation_min = params.float_activation_min;
const float output_activation_max = params.float_activation_max;
const int output_elements = output_shape.FlatSize();
const int output_dims_count = output_shape.DimensionsCount();
const int weights_dims_count = weights_shape.DimensionsCount();
const int batches = FlatSizeSkipDim(output_shape, output_dims_count - 1);
const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,
output_shape, output_dims_count - 1);
const int accum_depth = weights_shape.Dims(weights_dims_count - 1);
const int w0_size = sparsity.dim_metadata[0].dense_size;
const int* w1_segments = sparsity.dim_metadata[1].array_segments->data;
const int* w1_indices = sparsity.dim_metadata[1].array_indices->data;
for (int i = 0; i < output_elements; ++i) {
output_data[i] = 0.f;
}
for (int b = 0; b < batches; ++b) {
for (int idx_0 = 0; idx_0 < w0_size; ++idx_0) {
for (int pw1 = w1_segments[idx_0]; pw1 < w1_segments[idx_0 + 1]; ++pw1) {
int idx_1 = w1_indices[pw1];
output_data[b * output_depth + idx_0] +=
weights_data[pw1] * input_data[b * accum_depth + idx_1];
}
}
}
for (int b = 0; b < batches; ++b) {
for (int i = 0; i < output_depth; ++i) {
float total = output_data[b * output_depth + i];
const float bias_value = bias_data ? bias_data[i] : 0;
output_data[b * output_depth + i] = ActivationFunctionWithMinMax(
total + bias_value, output_activation_min, output_activation_max);
}
}
}
inline void FullyConnectedSparseWeight1x16Impl(
const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
const RuntimeShape& input_shape, const int8_t* input_data,
const RuntimeShape& weights_shape, const int8_t* weights_data,
const int32_t* per_channel_scale, const int32_t* per_channel_shift,
const RuntimeShape& bias_shape, const int32_t* bias_data,
const RuntimeShape& output_shape, int8_t* output_data, int thread_start,
int thread_end, const CpuBackendContext& cpu_backend_context) {
ruy::profiler::ScopeLabel label("FullyConnected");
ruy::profiler::ScopeLabel inner_label("1x16 Block Sparse");
const int input_dims_count = input_shape.DimensionsCount();
const int output_dims_count = output_shape.DimensionsCount();
const int weights_dims_count = weights_shape.DimensionsCount();
const int batches = thread_end - thread_start;
const int input_depth = MatchingDim(weights_shape, weights_dims_count - 1,
input_shape, input_dims_count - 1);
const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,
output_shape, output_dims_count - 1);
const int32_t input_offset = params.input_offset;
const int32_t output_offset = params.output_offset;
const int32_t output_multiplier = params.output_multiplier;
const int32_t output_shift = params.output_shift;
const int32_t output_activation_min = params.quantized_activation_min;
const int32_t output_activation_max = params.quantized_activation_max;
const int* w1_segments = sparsity.dim_metadata[1].array_segments->data;
const int* w1_indices = sparsity.dim_metadata[1].array_indices->data;
tensor_utils::SparseMatrixBatchVectorMultiplyAccumulate1x16(
weights_data, w1_segments, w1_indices, weights_shape.Dims(0),
weights_shape.Dims(1), input_data + thread_start * input_depth, bias_data,
batches, input_offset, output_multiplier, output_shift, per_channel_scale,
per_channel_shift, output_offset, output_activation_min,
output_activation_max, output_data + thread_start * output_depth);
}
inline void FullyConnectedSparseWeight1x4Impl(
const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& weights_shape, const float* weights_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data, int thread_start,
int thread_end, const CpuBackendContext& cpu_backend_context) {
ruy::profiler::ScopeLabel label("FullyConnected");
ruy::profiler::ScopeLabel inner_label("1x4 Block Sparse");
const float output_activation_min = params.float_activation_min;
const float output_activation_max = params.float_activation_max;
const int input_dims_count = input_shape.DimensionsCount();
const int output_dims_count = output_shape.DimensionsCount();
const int weights_dims_count = weights_shape.DimensionsCount();
const int batches = thread_end - thread_start;
const int input_depth = MatchingDim(weights_shape, weights_dims_count - 1,
input_shape, input_dims_count - 1);
const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,
output_shape, output_dims_count - 1);
const int* w1_segments = sparsity.dim_metadata[1].array_segments->data;
const int* w1_indices = sparsity.dim_metadata[1].array_indices->data;
tensor_utils::SparseMatrixBatchVectorMultiplyAccumulate1x4(
weights_data, w1_segments, w1_indices, weights_shape.Dims(0),
weights_shape.Dims(1), input_data + thread_start * input_depth, batches,
output_data + thread_start * output_depth);
ruy::profiler::ScopeLabel activation_label("activation function");
for (int b = thread_start; b < thread_end; ++b) {
for (int i = 0; i < output_depth; ++i) {
float total = output_data[b * output_depth + i];
const float bias_value = bias_data ? bias_data[i] : 0;
output_data[b * output_depth + i] = ActivationFunctionWithMinMax(
total + bias_value, output_activation_min, output_activation_max);
}
}
}
struct FullyConnectedSparseWeight1x4Task : cpu_backend_threadpool::Task {
FullyConnectedSparseWeight1x4Task(
const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& weights_shape, const float* weights_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data, int thread_start,
int thread_end, const CpuBackendContext& cpu_backend_context_x)
: sparsity(sparsity),
params(params),
input_shape(input_shape),
input_data(input_data),
weights_shape(weights_shape),
weights_data(weights_data),
bias_shape(bias_shape),
bias_data(bias_data),
output_shape(output_shape),
output_data(output_data),
thread_start(thread_start),
thread_end(thread_end),
cpu_backend_context(cpu_backend_context_x) {}
void Run() override {
FullyConnectedSparseWeight1x4Impl(
sparsity, params, input_shape, input_data, weights_shape, weights_data,
bias_shape, bias_data, output_shape, output_data, thread_start,
thread_end, cpu_backend_context);
}
private:
const TfLiteSparsity& sparsity;
const FullyConnectedParams& params;
const RuntimeShape& input_shape;
const float* input_data;
const RuntimeShape& weights_shape;
const float* weights_data;
const RuntimeShape& bias_shape;
const float* bias_data;
const RuntimeShape& output_shape;
float* output_data;
int thread_start;
int thread_end;
const CpuBackendContext& cpu_backend_context;
};
inline void FullyConnectedSparseWeight1x16(
const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
const RuntimeShape& input_shape, const int8_t* input_data,
const RuntimeShape& weights_shape, const int8_t* weights_data,
const int32_t* per_channel_scale, const int32_t* per_channel_shift,
const RuntimeShape& bias_shape, const int32_t* bias_data,
const RuntimeShape& output_shape, int8_t* output_data,
CpuBackendContext* cpu_backend_context) {
const int output_elements = output_shape.FlatSize();
memset(output_data, 0, output_elements * sizeof(int8_t));
const int batches =
FlatSizeSkipDim(output_shape, output_shape.DimensionsCount() - 1);
// TODO(b/220851507): Add multi-thread support for quantized sparse kernel.
return FullyConnectedSparseWeight1x16Impl(
sparsity, params, input_shape, input_data, weights_shape, weights_data,
per_channel_scale, per_channel_shift, bias_shape, bias_data, output_shape,
output_data, 0, batches, *cpu_backend_context);
}
// The multi-threaded kernel slices the workload along the batch dimension. If
// there's not enough batches of data, the number of threads used is equal to
// the batch size. We can improve this later with slicing along the row
// dimension of the weight.
inline void FullyConnectedSparseWeight1x4(
const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
const RuntimeShape& input_shape, const float* input_data,
const RuntimeShape& weights_shape, const float* weights_data,
const RuntimeShape& bias_shape, const float* bias_data,
const RuntimeShape& output_shape, float* output_data,
CpuBackendContext* cpu_backend_context) {
const int output_elements = output_shape.FlatSize();
memset(output_data, 0, output_elements * sizeof(float));
const int max_threads = cpu_backend_context->max_num_threads();
const int batches =
FlatSizeSkipDim(output_shape, output_shape.DimensionsCount() - 1);
const int thread_count = std::max(1, std::min(batches, max_threads));
if (thread_count == 1) {
return FullyConnectedSparseWeight1x4Impl(
sparsity, params, input_shape, input_data, weights_shape, weights_data,
bias_shape, bias_data, output_shape, output_data, 0, batches,
*cpu_backend_context);
}
std::vector<FullyConnectedSparseWeight1x4Task> tasks;
tasks.reserve(thread_count);
int thread_start = 0;
for (int i = 0; i < thread_count; ++i) {
// This makes sure the workload is relatively balanced when batches is not a
// multiple of thread_count. The first mod(batches, thread_count) tasks need
// to process one more batch than the rest.
int thread_end = thread_start + batches / thread_count;
if (i < batches % thread_count) thread_end++;
tasks.emplace_back(sparsity, params, input_shape, input_data, weights_shape,
weights_data, bias_shape, bias_data, output_shape,
output_data, thread_start, thread_end,
*cpu_backend_context);
thread_start = thread_end;
}
cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
cpu_backend_context);
}
} // namespace optimized_ops
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SPARSE_OPS_FULLY_CONNECTED_H_
tensorflow/lite/kernels/internal/optimized/sse_check.h
+28
View File
@@ -0,0 +1,28 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_CHECK_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_CHECK_H_
#if defined(__SSSE3__)
// SSSE 3 available: Use the SSE code.
#define SSE_OR_PORTABLE(funcname, ...) Sse##funcname(__VA_ARGS__)
#else
// No SSSE 3 available: Use Portable code
#define SSE_OR_PORTABLE(funcname, ...) Portable##funcname(__VA_ARGS__)
#endif
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_CHECK_H_
tensorflow/lite/kernels/internal/optimized/sse_tensor_utils.cc
+676
View File
@@ -0,0 +1,676 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#include "tensorflow/lite/kernels/internal/optimized/sse_tensor_utils_impl.h"
#ifdef __SSSE3__
#include <emmintrin.h> // SSE2
#include <tmmintrin.h> // SSSE3
#ifdef __SSE4_1__
#include <smmintrin.h> // SSE4.1
#endif
#ifdef __AVX2__
#include <immintrin.h>
#include "absl/base/prefetch.h"
#endif
#include <cstdint>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/cpu_backend_gemm.h"
#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
namespace tflite {
namespace tensor_utils {
namespace {
#if defined(__SSE2__)
// Note: this part is copied from XNNPACK/src/xnnpack/intrinsics-polyfill.h
// w.r.t the defition of '_mm_loadu_si32' intrinsic.
// GCC any, Clang pre-8, Android NDK Clang pre-8.0.7, Apple Clang pre-11, and
// ICC pre-16
#if (defined(__GNUC__) && !defined(__clang__) && \
!defined(__INTEL_COMPILER)) || \
(defined(__clang__) && !defined(__apple_build_version__) && \
(__clang_major__ < 8)) || \
(defined(__clang__) && defined(__ANDROID__) && (__clang_major__ == 8) && \
(__clang_minor__ == 0) && (__clang_patchlevel__ < 7)) || \
(defined(__clang__) && defined(__apple_build_version__) && \
(__apple_build_version__ < 11000000)) || \
(defined(__INTEL_COMPILER) && (__INTEL_COMPILER < 1600))
static inline __m128i _mm_loadu_si32(const void* address) {
return _mm_cvtsi32_si128(*((const int*)address));
}
#endif // GCC any, Clang pre-8, Android NDK Clang pre-8.0.7, Apple Clang pre-11
// and ICC pre-16
#endif // __SSE2__
// Dot product of four int8 vectors of 4 elements packed into a XMM register.
// Result is four int32 scalars packed into a XMM register.
// int8x4x4 · int8x4x4 => int32x4
static inline __m128i DotProdInt8x4x4(__m128i a_8x16, __m128i b_8x16) {
// Transfer sign from 'a' to 'b', as _mm_maddubs_epi16 treats 'a' unsigned.
b_8x16 = _mm_sign_epi8(b_8x16, a_8x16);
a_8x16 = _mm_abs_epi8(a_8x16);
// sumprod[i] = a[2*i]*b[2*i] + a[2*i+1]*b[2*i+1] (i = 0..7)
__m128i sumprod_16x8 = _mm_maddubs_epi16(a_8x16, b_8x16);
// sumprod[i] = sumprod[2*i]*1 + sumprod[2*i+1]*1 (i = 0..3)
return _mm_madd_epi16(sumprod_16x8, _mm_set1_epi16(1));
}
// Horizontally add 4 int32 values stored in a single XMM register to int32_t.
static inline int32_t ReduceInt32x4(__m128i acc) {
// Shuffle to contain high half of acc (both in high and low halfs).
__m128i shuffle = _mm_unpackhi_epi64(acc, acc);
// Add shuffle and acc; low half is sums of twos (high half is ignored).
acc = _mm_add_epi32(acc, shuffle);
// Shuffle the two elements in low half (ignore high half).
shuffle = _mm_shuffle_epi32(acc, _MM_SHUFFLE(2, 3, 0, 1));
// Add shuffle and acc; lowest element is sum of all 4 input.
acc = _mm_add_epi32(acc, shuffle);
// Return lowest element as int32_t.
return _mm_cvtsi128_si32(acc);
}
#ifdef __AVX2__
// Horizontally add 4 float values stored in a single XMM register to float.
static inline float ReduceFloat32x4(__m128 acc) {
__m128 shuffle = _mm_movehdup_ps(acc);
acc = _mm_add_ps(acc, shuffle);
shuffle = _mm_movehl_ps(shuffle, acc);
acc = _mm_add_ss(acc, shuffle);
return _mm_cvtss_f32(acc);
}
// Horizontally add 8 float values stored in a single XMM register to float.
static inline float ReduceFloat32x8(__m256 acc) {
__m128 low = _mm256_extractf128_ps(acc, 0);
__m128 high = _mm256_extractf128_ps(acc, 1);
return ReduceFloat32x4(_mm_add_ps(low, high));
}
// Dot product of four int8 vectors of 4 elements packed into a YMM register.
// Result is eight int32 scalars packed into a YMM register.
// int8x4x8 · int8x4x8 => int32x8
static inline __m256i DotProdInt8x4x8(__m256i a_16x16, __m256i b_16x16) {
// Transfer sign from 'a' to 'b', as _mm256_maddubs_epi16 treats 'a' unsigned.
b_16x16 = _mm256_sign_epi8(b_16x16, a_16x16);
a_16x16 = _mm256_abs_epi8(a_16x16);
// sumprod[i] = a[2*i]*b[2*i] + a[2*i+1]*b[2*i+1] (i = 0..15)
__m256i sumprod_16x16 = _mm256_maddubs_epi16(a_16x16, b_16x16);
// sumprod[i] = sumprod[2*i]*1 + sumprod[2*i+1]*1 (i = 0..7)
return _mm256_madd_epi16(sumprod_16x16, _mm256_set1_epi16(1));
}
#endif // __AVX2__
// Horizontally add each of 4 XMM registers with 4 int32 values, pack result
// into a single XMM register. Similar to ReduceInt32x4, but with 4x inputs.
static inline __m128i ReduceInt32x4x4(__m128i a, __m128i b, __m128i c,
__m128i d) {
// Assuming x = [x0, x1, x2, x3]
const __m128i a_b_lo_half = _mm_unpacklo_epi32(a, b); // [a0, b0, a1, b1]
const __m128i a_b_hi_half = _mm_unpackhi_epi32(a, b); // [a2, b2, a3, b3]
const __m128i a_plus_b =
_mm_add_epi32(a_b_lo_half, a_b_hi_half); // [a0+a2, b0+b2, a1+a3, b1+b3]
const __m128i c_d_lo_half = _mm_unpacklo_epi32(c, d); // [c0, d0, c1, d1]
const __m128i c_d_hi_half = _mm_unpackhi_epi32(c, d); // [c2, d2, c3, d3]
const __m128i c_plus_d =
_mm_add_epi32(c_d_lo_half, c_d_hi_half); // [c0+c2, d0+d2, c1+c3, d1+d3]
const __m128i all_evns =
_mm_unpacklo_epi64(a_plus_b, c_plus_d); // [a02, b02, c02, d02]
const __m128i all_odds =
_mm_unpackhi_epi64(a_plus_b, c_plus_d); // [a13, b13, c13, d13]
return _mm_add_epi32(all_evns, all_odds); // [a0123, b0123, c0123, d0123]
}
// Returns the ith element of a XMM register holding float numbers.
template <int i>
float GetFloatVectorElement(__m128 v) {
static_assert(i >= 0 && i < 4, "The index must be 0 <= i < 4.");
// Note, _mm_extract_ps returns int, so we can't use it here.
// These lines will be optimized to extractps anyway.
v = _mm_shuffle_ps(v, v, _MM_SHUFFLE(i, i, i, i));
return _mm_cvtss_f32(v);
}
} // namespace
#ifdef __AVX2__
constexpr int kFloatValuesPerAvx2Vector = 8;
template <int PerVectorSize>
inline int RoundDownVectors(int size) {
return size & ~(PerVectorSize - 1);
}
void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
const float* __restrict__ matrix, int m_rows, int m_cols,
const float* __restrict__ vector, int n_batch, float* __restrict__ result) {
// If v_size is not divisible by the vector size, then we need to process the
// final few elements sequentially. postamble_start shows the start index
// where this should happen.
const int postamble_start =
RoundDownVectors<kFloatValuesPerAvx2Vector>(m_cols);
for (int b = 0; b < n_batch; ++b) {
float* result_in_batch = result + b * m_rows;
const float* vector_in_batch = vector + b * m_cols;
const float* matrix_row = matrix;
// Main matrix by vector multiplication loop
for (int r = 0; r < m_rows; ++r) {
__m256 acc_32x8 = _mm256_setzero_ps();
int c = 0;
for (; c < postamble_start; c += kFloatValuesPerAvx2Vector) {
// Load 8 float values from vector and matrix row.
__m256 vector_f32x8 = _mm256_loadu_ps(vector_in_batch + c);
__m256 matrix_f32x8 = _mm256_loadu_ps(matrix_row + c);
// Multiply the vector and matrix row and add to accumulator.
__m256 res = _mm256_mul_ps(vector_f32x8, matrix_f32x8);
acc_32x8 = _mm256_add_ps(acc_32x8, res);
}
// Add the 8 intermediate sum values to get the final dot-prod value for
// this column.
float sum = ReduceFloat32x8(acc_32x8);
for (; (c < m_cols); c++) {
sum += matrix_row[c] * vector_in_batch[c];
}
*result_in_batch += sum;
++result_in_batch;
matrix_row += m_cols;
}
}
}
void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, const int32_t* row_sums) {
for (std::intptr_t batch = 0; batch < n_batch; ++batch) {
const float batch_scaling_factor = scaling_factors[batch];
const int32_t batch_offset = input_offset ? input_offset[batch] : 0;
// Compute dot-product for every column.
for (std::intptr_t row = 0; row < m_rows; ++row) {
// Get the address of the first element of the row.
const int8_t* __restrict__ row_ptr = matrix + row * m_cols;
const float row_scale =
per_channel_scale ? per_channel_scale[row] * batch_scaling_factor
: batch_scaling_factor;
const int32_t row_offset =
row_sums && batch_offset ? batch_offset * row_sums[row] : 0;
// Initialize the dot product sum for the row to 0.
__m256i dotprod_32x8 = _mm256_setzero_si256();
std::intptr_t col = 0;
constexpr int prefetch_distance = 704;
// For every block of 32x 8-bit inputs.
while (col < (m_cols & ~31)) {
absl::PrefetchToLocalCache(vectors + col + prefetch_distance);
absl::PrefetchToLocalCache(row_ptr + col + prefetch_distance);
const __m256i vec_16x16 =
_mm256_loadu_si256(reinterpret_cast<const __m256i*>(vectors + col));
const __m256i row_16x16 =
_mm256_loadu_si256(reinterpret_cast<const __m256i*>(row_ptr + col));
// dotprod += vec · row
dotprod_32x8 = _mm256_add_epi32(dotprod_32x8,
DotProdInt8x4x8(vec_16x16, row_16x16));
col += 32;
}
// Sum lower and upper halves of 32x8 vector into 32x4 vector
__m128i low = _mm256_extracti128_si256(dotprod_32x8, 0);
__m128i high = _mm256_extracti128_si256(dotprod_32x8, 1);
__m128i dotprod_32x4 = _mm_add_epi32(low, high);
// Postamble for 16x 8-bit inputs.
if (col < (m_cols & ~15)) {
const __m128i vec_16x8 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(vectors + col));
const __m128i row_16x8 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(row_ptr + col));
// dotprod += vec · row
dotprod_32x4 =
_mm_add_epi32(dotprod_32x4, DotProdInt8x4x4(vec_16x8, row_16x8));
col += 16;
}
// Postamble for 8x 8-bit inputs.
if (col < (m_cols & ~7)) {
const __m128i vec_16x8 = _mm_cvtepi8_epi16(
_mm_loadl_epi64(reinterpret_cast<const __m128i*>(vectors + col)));
const __m128i row_16x8 = _mm_cvtepi8_epi16(
_mm_loadl_epi64(reinterpret_cast<const __m128i*>(row_ptr + col)));
// dotprod += vec · row
dotprod_32x4 =
_mm_add_epi32(dotprod_32x4, _mm_madd_epi16(vec_16x8, row_16x8));
col += 8;
}
// Postamble for 4x 8-bit inputs.
if (col < (m_cols & ~3)) {
const __m128i vec_32x4 = _mm_cvtepi8_epi32(
_mm_loadu_si32(reinterpret_cast<const __m128i*>(vectors + col)));
const __m128i row_32x4 = _mm_cvtepi8_epi32(
_mm_loadu_si32(reinterpret_cast<const __m128i*>(row_ptr + col)));
// dotprod += vec · row
dotprod_32x4 =
_mm_add_epi32(dotprod_32x4, _mm_mullo_epi32(vec_32x4, row_32x4));
col += 4;
}
// Horizontally add the 4 intermediate sum values to get the final
// dot-prod value for this row.
int32_t sum = ReduceInt32x4(dotprod_32x4);
#pragma clang loop unroll(disable) vectorize(disable)
// Postamble loop for <4x remaining 8-bit inputs.
for (; col < m_cols; ++col) {
sum += row_ptr[col] * vectors[col];
} // for col
if (row_offset) {
sum -= row_offset;
}
*result += sum * row_scale;
++result;
} // for row
vectors += m_cols;
} // for batch
}
#endif // __AVX2__
void SseMatrixBatchVectorMultiplyAccumulateImpl(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, const int32_t* row_sums) {
#ifdef __AVX2__
Avx2MatrixBatchVectorMultiplyAccumulateImpl(
matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
per_channel_scale, input_offset, row_sums);
return;
#else
for (std::intptr_t batch = 0; batch < n_batch; ++batch) {
const float batch_scaling_factor = scaling_factors[batch];
const int32_t batch_offset = input_offset ? input_offset[batch] : 0;
// Compute dot-product for every column.
for (std::intptr_t row = 0; row < m_rows; ++row) {
// Get the address of the first element of the row.
const int8_t* __restrict__ row_ptr = matrix + row * m_cols;
const float row_scale =
per_channel_scale ? per_channel_scale[row] * batch_scaling_factor
: batch_scaling_factor;
const int32_t row_offset =
row_sums && batch_offset ? batch_offset * row_sums[row] : 0;
// Initialize the dot product sum for the row to 0.
__m128i dotprod_32x4 = _mm_setzero_si128();
std::intptr_t col = 0;
// For every block of 16x 8-bit inputs.
while (col < (m_cols & ~15)) {
const __m128i vec_8x16 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(vectors + col));
const __m128i row_8x16 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(row_ptr + col));
// dotprod += vec · row
dotprod_32x4 =
_mm_add_epi32(dotprod_32x4, DotProdInt8x4x4(vec_8x16, row_8x16));
col += 16;
}
#ifdef __SSE4_1__
// Postamble for 8x 8-bit inputs.
if (col < (m_cols & ~7)) {
const __m128i vec_16x8 = _mm_cvtepi8_epi16(
_mm_loadl_epi64(reinterpret_cast<const __m128i*>(vectors + col)));
const __m128i row_16x8 = _mm_cvtepi8_epi16(
_mm_loadl_epi64(reinterpret_cast<const __m128i*>(row_ptr + col)));
// dotprod += vec · row
dotprod_32x4 =
_mm_add_epi32(dotprod_32x4, _mm_madd_epi16(vec_16x8, row_16x8));
col += 8;
}
// Postamble for 4x 8-bit inputs.
if (col < (m_cols & ~3)) {
const __m128i vec_32x4 = _mm_cvtepi8_epi32(
_mm_loadu_si32(reinterpret_cast<const __m128i*>(vectors + col)));
const __m128i row_32x4 = _mm_cvtepi8_epi32(
_mm_loadu_si32(reinterpret_cast<const __m128i*>(row_ptr + col)));
// dotprod += vec · row
dotprod_32x4 =
_mm_add_epi32(dotprod_32x4, _mm_mullo_epi32(vec_32x4, row_32x4));
col += 4;
}
#endif
// Horizontally add the 4 intermediate sum values to get the final
// dot-prod value for this row.
int32_t sum = ReduceInt32x4(dotprod_32x4);
#if defined(__SSE4_1__) && defined(__clang__)
// SSE 4.1: Don't try to unroll and vectorize this, already done above.
#pragma clang loop unroll(disable) vectorize(disable)
#endif
// Postamble loop for <4x (<16x without SSE 4.1) remaining 8-bit inputs.
for (; col < m_cols; ++col) {
sum += row_ptr[col] * vectors[col];
} // for col
if (row_offset) {
sum -= row_offset;
}
*result += sum * row_scale;
++result;
} // for row
vectors += m_cols;
} // for batch
#endif // ifdef __AVX2__
}
void SseCpuBackendGemm(const int8_t* input, const int32_t* bias,
const int8_t* input_to_gate_weights, int32_t n_batch,
int32_t n_input, int32_t n_output, int32_t output_zp,
int32_t* scratch, CpuBackendContext* context) {
using ::tflite::cpu_backend_gemm::Gemm;
using ::tflite::cpu_backend_gemm::GemmParams;
using ::tflite::cpu_backend_gemm::MatrixParams;
MatrixParams<int8_t> lhs_params;
lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
lhs_params.rows = n_output;
lhs_params.cols = n_input;
lhs_params.cache_policy = cpu_backend_gemm::CachePolicy::kCacheIfLargeSpeedup;
MatrixParams<int8_t> rhs_params;
rhs_params.order = cpu_backend_gemm::Order::kColMajor;
rhs_params.rows = n_input;
rhs_params.cols = n_batch;
MatrixParams<int32_t> dst_params;
dst_params.order = cpu_backend_gemm::Order::kColMajor;
dst_params.rows = n_output;
dst_params.cols = n_batch;
GemmParams<int32, int32> gemm_params;
if (bias) {
gemm_params.bias = bias;
}
cpu_backend_gemm::Gemm(lhs_params, input_to_gate_weights, rhs_params, input,
dst_params, scratch, gemm_params, context);
}
void SseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result) {
SseMatrixBatchVectorMultiplyAccumulateImpl(
matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
/*per_channel_scale=*/nullptr, /*input_offset=*/nullptr,
/*row_sums=*/nullptr);
}
void SseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch, int32_t* scratch,
float* __restrict__ result, CpuBackendContext* context) {
// TODO(b/183178387): Use a proper query to detect AVX/optimized paths.
if (m_rows % 4 == 0 && !context->PreferGemmlowpOnX86()) {
const int32_t* bias = static_cast<const int32_t*>(nullptr);
SseCpuBackendGemm(vectors, bias, matrix, n_batch, m_cols, m_rows,
/*output_zp=*/0, scratch, context);
{
ruy::profiler::ScopeLabel label("HybridMultiplyScalingFactor");
// Multiply by float scaling factors and write to result
const int total_size = n_batch * m_rows;
int i = 0;
for (; i <= total_size - 8; i += 8, result += 8) {
const float batch_scaling_factor0 = scaling_factors[i / m_rows];
const float batch_scaling_factor1 = scaling_factors[(i + 4) / m_rows];
const __m128 scaling_factor0 = _mm_set1_ps(batch_scaling_factor0);
const __m128 scaling_factor1 = _mm_set1_ps(batch_scaling_factor1);
const __m128i scratch_val0 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(scratch + i));
const __m128i scratch_val1 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(scratch + i + 4));
const __m128 float_val0 = _mm_cvtepi32_ps(scratch_val0);
const __m128 float_val1 = _mm_cvtepi32_ps(scratch_val1);
const __m128 prod0 = _mm_mul_ps(float_val0, scaling_factor0);
const __m128 result0 = _mm_add_ps(_mm_load1_ps(result), prod0);
const __m128 prod1 = _mm_mul_ps(float_val1, scaling_factor1);
const __m128 result1 = _mm_add_ps(_mm_load1_ps(result + 4), prod1);
_mm_store_ps(result, result0);
_mm_store_ps(result + 4, result1);
}
scratch += i;
for (; i < total_size; i++) {
const float batch_scaling_factor = scaling_factors[i / m_rows];
int32_t x = *(scratch++);
*result += x * batch_scaling_factor;
++result;
}
}
return;
}
SseMatrixBatchVectorMultiplyAccumulateImpl(
matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
/*per_channel_scale=*/nullptr, /*input_offset=*/nullptr,
/*row_sums=*/nullptr);
}
void SseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
bool* compute_row_sums, CpuBackendContext* context) {
if ((input_offset != nullptr) && (!compute_row_sums || *compute_row_sums)) {
SseReductionSumVector(matrix, row_sums, m_rows, m_cols);
if (compute_row_sums) {
*compute_row_sums = false;
}
}
SseMatrixBatchVectorMultiplyAccumulateImpl(
matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
per_channel_scale, input_offset, row_sums);
}
namespace {
// Implements sparse-matrix - vector multiply-accumulate.
inline void SseSparseMatrixVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
const int m_rows, const int m_cols, const int8_t* __restrict__ vector,
const float batch_scaling_factor, float* __restrict__ result,
const float* per_channel_scale) {
static const std::intptr_t kBlockSize = 16;
TFLITE_DCHECK_EQ(m_cols % kBlockSize, 0);
const uint8_t* __restrict__ ledger_ptr = ledger;
for (std::intptr_t row = 0; row < m_rows; ++row) {
// Initialize the dot product sum for the row to 0.
__m128i dotprod_32x4 = _mm_setzero_si128();
std::intptr_t num_nonzero_blocks = *ledger_ptr++;
for (std::intptr_t i = 0; i < num_nonzero_blocks; i++) {
const std::intptr_t col_index = *ledger_ptr++ * kBlockSize;
const __m128i vec_8x16 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(vector + col_index));
const __m128i row_8x16 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(matrix));
// dotprod += vec · row
dotprod_32x4 =
_mm_add_epi32(dotprod_32x4, DotProdInt8x4x4(vec_8x16, row_8x16));
matrix += kBlockSize;
} // for col
// Horizontally add the 4 intermediate sum values to get the final
// dot-prod value for this row.
int32_t dotprod = ReduceInt32x4(dotprod_32x4);
const float total_scaling_factor =
per_channel_scale ? per_channel_scale[row] * batch_scaling_factor
: batch_scaling_factor;
result[row] += dotprod * total_scaling_factor;
} // for row
}
// Implements sparse-matrix - batch-of-4-vectors multiply-accumulate.
// The stride between vectors and results must be equal to m_cols.
// Parameter 'batch' is the index of the first batch, must be a multiple of 4.
inline void SseSparseMatrix4VectorsMultiplyAccumulate(
const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
const int m_rows, const int m_cols,
const int8_t* __restrict__ const vectors,
const __m128 batch_scaling_factors_fx4, float* __restrict__ const results,
const float* per_channel_scale) {
static const std::intptr_t kBlockSize = 16;
TFLITE_DCHECK_EQ(m_cols % kBlockSize, 0);
const int8_t* __restrict__ vector0 = vectors + 0 * m_cols;
const int8_t* __restrict__ vector1 = vectors + 1 * m_cols;
const int8_t* __restrict__ vector2 = vectors + 2 * m_cols;
const int8_t* __restrict__ vector3 = vectors + 3 * m_cols;
float* __restrict__ result0 = results + 0 * m_rows;
float* __restrict__ result1 = results + 1 * m_rows;
float* __restrict__ result2 = results + 2 * m_rows;
float* __restrict__ result3 = results + 3 * m_rows;
for (std::intptr_t row = 0; row < m_rows; ++row) {
// Initialize the dot product sum for the row to 0.
__m128i dp0_32x4 = _mm_setzero_si128();
__m128i dp1_32x4 = _mm_setzero_si128();
__m128i dp2_32x4 = _mm_setzero_si128();
__m128i dp3_32x4 = _mm_setzero_si128();
std::intptr_t num_nonzero_blocks = *ledger++;
for (std::intptr_t i = 0; i < num_nonzero_blocks; i++) {
const std::intptr_t col_index = *ledger++ * kBlockSize;
// vecN are for different batches
const __m128i vec0_8x16 = _mm_loadu_si128(
reinterpret_cast<const __m128i*>(vector0 + col_index));
const __m128i vec1_8x16 = _mm_loadu_si128(
reinterpret_cast<const __m128i*>(vector1 + col_index));
const __m128i vec2_8x16 = _mm_loadu_si128(
reinterpret_cast<const __m128i*>(vector2 + col_index));
const __m128i vec3_8x16 = _mm_loadu_si128(
reinterpret_cast<const __m128i*>(vector3 + col_index));
const __m128i row_8x16 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(matrix));
// dp += vec · row
// dpN are for different batches
dp0_32x4 = _mm_add_epi32(dp0_32x4, DotProdInt8x4x4(vec0_8x16, row_8x16));
dp1_32x4 = _mm_add_epi32(dp1_32x4, DotProdInt8x4x4(vec1_8x16, row_8x16));
dp2_32x4 = _mm_add_epi32(dp2_32x4, DotProdInt8x4x4(vec2_8x16, row_8x16));
dp3_32x4 = _mm_add_epi32(dp3_32x4, DotProdInt8x4x4(vec3_8x16, row_8x16));
matrix += kBlockSize;
} // for col
// Horizontally add the 4 intermediate values.
const __m128i dp_32x4 =
ReduceInt32x4x4(dp0_32x4, dp1_32x4, dp2_32x4, dp3_32x4);
// Convert to float
const __m128 dp_fx4 = _mm_cvtepi32_ps(dp_32x4);
// Load the results (This is an Accumulate function..)
__m128 result_fx4 =
_mm_set_ps(result3[row], result2[row], result1[row], result0[row]);
const __m128 total_scaling_factors_fx4 =
per_channel_scale ? _mm_mul_ps(batch_scaling_factors_fx4,
_mm_set1_ps(per_channel_scale[row]))
: batch_scaling_factors_fx4;
// result += dp .* scaling
result_fx4 =
_mm_add_ps(result_fx4, _mm_mul_ps(dp_fx4, total_scaling_factors_fx4));
// Save the results
result0[row] = GetFloatVectorElement<0>(result_fx4);
result1[row] = GetFloatVectorElement<1>(result_fx4);
result2[row] = GetFloatVectorElement<2>(result_fx4);
result3[row] = GetFloatVectorElement<3>(result_fx4);
} // for row
}
} // namespace
void SseSparseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
const int m_rows, const int m_cols, const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ results, const float* per_channel_scale) {
int batch = 0;
const int kBatchSize4 = 4;
const int n_batch_rounddown_to_batchsize_4 = n_batch & ~(kBatchSize4 - 1);
while (batch < n_batch_rounddown_to_batchsize_4) {
const __m128 scaling_factors_fx4 = _mm_loadu_ps(scaling_factors + batch);
SseSparseMatrix4VectorsMultiplyAccumulate(matrix, ledger, m_rows, m_cols,
vectors, scaling_factors_fx4,
results, per_channel_scale);
batch += kBatchSize4;
vectors += kBatchSize4 * m_cols;
results += kBatchSize4 * m_rows;
} // for batch
while (batch < n_batch) {
SseSparseMatrixVectorMultiplyAccumulate(matrix, ledger, m_rows, m_cols,
vectors, scaling_factors[batch],
results, per_channel_scale);
++batch;
vectors += m_cols;
results += m_rows;
} // for batch
}
void SseReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
const int output_size, const int reduction_size) {
static constexpr std::intptr_t kBlockSize = 16;
for (std::intptr_t row = 0; row < output_size; ++row) {
const int8_t* __restrict__ row_ptr = input_vector + row * reduction_size;
__m128i row_sum_16x8 = _mm_setzero_si128();
std::intptr_t col = 0;
for (; col < (reduction_size & ~(kBlockSize - 1)); col += kBlockSize) {
const __m128i row_8x16 =
_mm_loadu_si128(reinterpret_cast<const __m128i*>(row_ptr + col));
const __m128i row_16x8 = _mm_maddubs_epi16(_mm_set1_epi8(1), row_8x16);
row_sum_16x8 = _mm_add_epi16(row_sum_16x8, row_16x8);
} // for col
#ifdef __SSE4_1__
// Postamble for 8x 8-bit inputs.
if (col < (reduction_size & ~7)) {
// _mm_loadu_si64 not supported in gcc versions < 9, breaks kokoro build.
const __m128i row_16x8 = _mm_cvtepi8_epi16(
_mm_loadl_epi64(reinterpret_cast<const __m128i*>(row_ptr + col)));
// dotprod += vec · row
row_sum_16x8 = _mm_add_epi16(row_sum_16x8, row_16x8);
col += 8;
}
#endif
const __m128i row_sum_32x4 =
_mm_madd_epi16(row_sum_16x8, _mm_set1_epi16(1));
int32_t row_sum = ReduceInt32x4(row_sum_32x4);
#if defined(__SSE4_1__) && defined(__clang__)
// SSE 4.1: Don't try to unroll and vectorize this, already done above.
#pragma clang loop unroll(disable) vectorize(disable)
#endif
for (; col < reduction_size; col++) {
row_sum += row_ptr[col];
}
output_vector[row] = row_sum;
}
}
} // namespace tensor_utils
} // namespace tflite
#endif // __SSSE3__
tensorflow/lite/kernels/internal/optimized/sse_tensor_utils.h
+347
View File
@@ -0,0 +1,347 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_H_
// Note: This file is a copy-paste version of neon_tensor_utils.h, only
// difference is in MatrixBatchVectorMultiplyAccumulate and
// SparseMatrixBatchVectorMultiplyAccumulate (other functions do not have SSE
// implementation yet).
// Note: Most of the functions below use NEON_OR_PORTABLE, through the Intel
// NEON_2_SSE translator library. If a native SSE version of a function is
// implemented, replace the appropriate one to SSE_OR_PORTABLE.
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
#include "tensorflow/lite/kernels/internal/optimized/neon_tensor_utils_impl.h"
#include "tensorflow/lite/kernels/internal/optimized/sse_check.h"
#include "tensorflow/lite/kernels/internal/optimized/sse_tensor_utils_impl.h"
#include "tensorflow/lite/kernels/internal/reference/portable_tensor_utils_impl.h"
namespace tflite {
namespace tensor_utils {
void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
int m_cols, const float* vector,
int n_batch, float* result) {
#if defined(__AVX2__)
Avx2MatrixBatchVectorMultiplyAccumulateImpl(matrix, m_rows, m_cols, vector,
n_batch, result);
#else
NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vector, n_batch, result);
#endif
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result) {
SSE_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vectors, scaling_factors, n_batch, result);
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors, const float* scaling_factors,
int n_batch, float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
bool* compute_row_sums, CpuBackendContext* context) {
SSE_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vectors, scaling_factors, n_batch, result, per_channel_scale,
input_offset, scratch, row_sums, compute_row_sums, context);
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
int32_t* __restrict__ scratch, float* __restrict__ result,
CpuBackendContext* __restrict__ context) {
SSE_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
vectors, scaling_factors, n_batch, scratch, result, context);
}
void SparseMatrixBatchVectorMultiplyAccumulate1x4(
const float* __restrict__ matrix, const int32_t* __restrict__ segments,
const int32_t* __restrict__ indices, int m_rows, int m_cols,
const float* __restrict__ vector, int n_batch, float* __restrict__ result) {
NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x4, matrix,
segments, indices, m_rows, m_cols, vector, n_batch, result);
}
void SparseMatrixBatchVectorMultiplyAccumulate1x16(
const int8_t* __restrict__ matrix, const int32_t* __restrict__ segments,
const int32_t* __restrict__ indices, int m_rows, int m_cols,
const int8_t* __restrict__ vector, const int32_t* __restrict__ bias_vector,
int n_batch, const int32_t input_offset, const int32_t output_multiplier,
const int32_t output_shift, const int32_t* per_channel_scale,
const int32_t* per_channel_shift, const int32_t output_offset,
const int32_t output_activation_min, const int32_t output_activation_max,
int8_t* __restrict__ result) {
NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x16, matrix,
segments, indices, m_rows, m_cols, vector, bias_vector,
n_batch, input_offset, output_multiplier, output_shift,
per_channel_scale, per_channel_shift, output_offset,
output_activation_min, output_activation_max, result);
}
void SparseMatrixBatchVectorMultiplyAccumulate(
const float* __restrict__ matrix, const uint8_t* __restrict__ ledger,
int m_rows, int m_cols, const float* __restrict__ vector, int n_batch,
float* __restrict__ result) {
NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
m_rows, m_cols, vector, n_batch, result);
}
void SparseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
const int m_rows, const int m_cols, const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result, const float* per_channel_scale) {
SSE_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
m_rows, m_cols, vectors, scaling_factors, n_batch, result,
per_channel_scale);
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* input, const int32_t* input_zeropoint_times_weights,
const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
int32_t* scratch, int16_t* output, CpuBackendContext* context) {
PortableMatrixBatchVectorMultiplyAccumulate(
input, input_zeropoint_times_weights, input_to_gate_weights, multiplier,
shift, n_batch, n_input, n_output, output_zp, scratch, output, context);
}
void MatrixBatchVectorMultiplyAccumulate(
const int8_t* input, const int32_t* input_zeropoint_times_weights,
const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
int32_t* scratch, int8_t* output, CpuBackendContext* context) {
PortableMatrixBatchVectorMultiplyAccumulate(
input, input_zeropoint_times_weights, input_to_gate_weights, multiplier,
shift, n_batch, n_input, n_output, output_zp, scratch, output, context);
}
void MatrixBatchVectorMultiply(const int8_t* input, int32_t input_zeropoint,
const int8_t* input_to_gate_weights,
int32_t input_to_gate_effective_scale_a,
int32_t input_to_gate_effective_scale_b,
int32_t n_batch, int32_t n_input, int32_t n_cell,
int8_t* gate_output, int8_t gate_output_zp) {
PortableMatrixBatchVectorMultiply(
input, input_zeropoint, input_to_gate_weights,
input_to_gate_effective_scale_a, input_to_gate_effective_scale_b, n_batch,
n_input, n_cell, gate_output, gate_output_zp);
}
void MatrixBatchVectorMultiply(const int16_t* hidden,
const int8_t* hidden_to_output_weights,
int32_t proj_effective_scale_a,
int32_t proj_effective_scale_b,
const int32_t* gate_bias, int32_t n_batch,
int32_t n_hidden, int32_t n_output,
int32_t output_zp, int8_t* proj_output) {
PortableMatrixBatchVectorMultiply(hidden, hidden_to_output_weights,
proj_effective_scale_a,
proj_effective_scale_b, gate_bias, n_batch,
n_hidden, n_output, output_zp, proj_output);
}
void MatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar,
int32_t n_row, int32_t n_col,
int32_t* output) {
PortableMatrixScalarMultiplyAccumulate(matrix, scalar, n_row, n_col, output);
}
void ApplyLayerNorm(const int16_t* input, const int16_t* layer_norm_weights,
const int32_t* bias, int32_t layer_norm_scale_a,
int32_t layer_norm_scale_b, int32_t variance_limit,
int n_batch, int n_input, int16_t* output) {
PortableApplyLayerNorm(input, layer_norm_weights, bias, layer_norm_scale_a,
layer_norm_scale_b, variance_limit, n_batch, n_input,
output);
}
void ApplyLayerNormFloat(const int16_t* input,
const int16_t* layer_norm_weights,
int32_t layer_norm_scale_a, int32_t layer_norm_scale_b,
const int32_t* bias, int n_batch, int n_input,
int16_t* output) {
PortableApplyLayerNormFloat(input, layer_norm_weights, layer_norm_scale_a,
layer_norm_scale_b, bias, n_batch, n_input,
output);
}
void ApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input,
int16_t* output) {
PortableApplySigmoid(input, n_batch, n_input, output);
}
void ApplySigmoidFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
int16_t* output) {
PortableApplySigmoidFloat(input, n_batch, n_input, output);
}
void ApplyTanh(int32_t intger_bits, const int16_t* input, int32_t n_batch,
int32_t n_input, int16_t* output) {
PortableApplyTanh(intger_bits, input, n_batch, n_input, output);
}
void ApplyTanhFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
int32_t integer_bits, int16_t* output) {
PortableApplyTanhFloat(input, n_batch, n_input, integer_bits, output);
}
void CwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
int n_input, int shift, int16_t* output) {
PortableCwiseMul(input_1, input_2, n_batch, n_input, shift, output);
}
void CwiseMul(const int16_t* input_1, const int16_t* input_2,
int32_t multiplier, int32_t shift, int32_t n_batch,
int32_t n_input, int32_t output_zp, int8_t* output) {
PortableCwiseMul(input_1, input_2, multiplier, shift, n_batch, n_input,
output_zp, output);
}
void CwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch,
int n_input, int16_t* output) {
PortableCwiseAdd(input_1, input_2, n_batch, n_input, output);
}
void CwiseClipping(float* vector, const int v_size,
const float clipping_value) {
PortableCwiseClipping(vector, v_size, clipping_value);
}
void CwiseClipping(int16_t* vector, const int v_size,
const int16_t clipping_value) {
PortableCwiseClipping(vector, v_size, clipping_value);
}
void CwiseClipping(int8_t* vector, const int v_size,
const int8_t clipping_value) {
PortableCwiseClipping(vector, v_size, clipping_value);
}
void BatchVectorBatchVectorDotProduct(const int16_t* vector1,
const int16_t* vector2, int v_size,
int n_batch, int32_t* result) {
PortableBatchVectorBatchVectorDotProduct(vector1, vector2, v_size, n_batch,
result);
}
void VectorBatchVectorCwiseProductAccumulate(const int16_t* vector, int v_size,
const int16_t* batch_vector,
int n_batch, int32_t multiplier,
int shift, int16_t* result) {
NEON_OR_PORTABLE(VectorBatchVectorCwiseProductAccumulate, vector, v_size,
batch_vector, n_batch, multiplier, shift, result);
}
float VectorVectorDotProduct(const float* vector1, const float* vector2,
int v_size) {
return NEON_OR_PORTABLE(VectorVectorDotProduct, vector1, vector2, v_size);
}
void Sub1Vector(const float* vector, int v_size, float* result) {
NEON_OR_PORTABLE(Sub1Vector, vector, v_size, result);
}
void Sub1Vector(const int16_t* vector, int v_size, int16_t* result) {
PortableSub1Vector(vector, v_size, result);
}
// Check if all entries of a vector are zero for float.
bool IsZeroVector(const float* vector, int v_size) {
return NEON_OR_PORTABLE(IsZeroVector, vector, v_size);
}
// Check if all entries of a vector are zero for int8.
bool IsZeroVector(const int8_t* vector, int v_size) {
return PortableIsZeroVector(vector, v_size);
}
void VectorScalarMultiply(const int8_t* vector, int v_size, float scale,
float* result) {
NEON_OR_PORTABLE(VectorScalarMultiply, vector, v_size, scale, result);
}
void SymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float* min_value,
float* max_value, float* scaling_factor) {
NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
min_value, max_value, scaling_factor);
}
void SymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float min_value,
float max_value, float* scaling_factor) {
NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
min_value, max_value, scaling_factor);
}
void AsymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float* scaling_factor,
int32_t* offset) {
NEON_OR_PORTABLE(AsymmetricQuantizeFloats, values, size, quantized_values,
scaling_factor, offset);
}
void ReductionSumVector(const float* input_vector, float* output_vector,
int output_size, int reduction_size) {
NEON_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
reduction_size);
}
void ReductionSumVector(const int32_t* input_vector, int32_t* output_vector,
int output_size, int reduction_size) {
PortableReductionSumVector(input_vector, output_vector, output_size,
reduction_size);
}
void ReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
int output_size, int reduction_size) {
SSE_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
reduction_size);
}
void MeanStddevNormalization(const float* __restrict__ input_vector,
float* __restrict__ output_vector, int v_size,
int n_batch) {
PortableMeanStddevNormalization(input_vector, output_vector, v_size, n_batch);
}
void TwoGateSaturatingAdd(const int8_t* input, int8_t input_zp,
const int8_t* recurrent, int8_t recurrent_zp,
int32_t input_effective_scale_a,
int32_t input_effective_scale_b,
int32_t recurrent_effective_scale_a,
int32_t recurrent_effective_scale_b, int32_t n_batch,
int32_t n_cell, int16_t* output) {
PortableTwoGateSaturatingAdd(
input, input_zp, recurrent, recurrent_zp, input_effective_scale_a,
input_effective_scale_b, recurrent_effective_scale_a,
recurrent_effective_scale_b, n_batch, n_cell, output);
}
} // namespace tensor_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_H_
tensorflow/lite/kernels/internal/optimized/sse_tensor_utils_impl.h
+87
View File
@@ -0,0 +1,87 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_IMPL_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_IMPL_H_
#include <cstdint>
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#if defined(_MSC_VER)
#define __restrict__ __restrict
#endif
namespace tflite {
namespace tensor_utils {
#if defined(__AVX2__)
// Matrix multiplication for float values.
void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
const float* __restrict__ matrix, int m_rows, int m_cols,
const float* __restrict__ vector, int n_batch, float* __restrict__ result);
// Matrix multiplication for quantized values using asymmetric quantization.
void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
bool* compute_row_sums, CpuBackendContext* context);
#endif // defined(__AVX2__)
#ifdef __SSSE3__
// Matrix multiplication for quantized values using symmetric quantization.
void SseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result);
// Matrix multiplication for quantized values using symmetric quantization
// with additional scratch memory for GEMM operation prior to scaling.
void SseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch, int32_t* scratch,
float* __restrict__ result, CpuBackendContext* context);
// Matrix multiplication for quantized values using asymmetric quantization.
void SseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result, const float* per_channel_scale,
const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
bool* compute_row_sums, CpuBackendContext* context);
// Matrix multiplication for quantized values using symmetric quantization.
// Sparse version.
void SseSparseMatrixBatchVectorMultiplyAccumulate(
const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
const int m_rows, const int m_cols, const int8_t* __restrict__ vectors,
const float* __restrict__ scaling_factors, int n_batch,
float* __restrict__ result, const float* per_channel_scale);
void SseReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
const int output_size, const int reduction_size);
#endif // __SSSE3__
} // namespace tensor_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_IMPL_H_