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_