Sync from upstream TF. (#3625)

This commit is contained in:
TFLM-bot 2026-07-13 11:31:17 -07:00 committed by GitHub
parent fddd3707a3
commit 35b20c5c74
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 204 additions and 15 deletions

View file

@ -1007,19 +1007,19 @@ template <int N>
struct NdArrayDesc {
// The "extent" of each dimension. Indices along dimension d must be in the
// half-open interval [0, extents[d]).
int extents[N];
int64_t extents[N];
// The number of *elements* (not bytes) between consecutive indices of each
// dimension.
int strides[N];
int64_t strides[N];
};
// DO NOT USE THIS FUNCTION FOR NEW FUNCTIONALITY BEYOND IMPLEMENTING
// BROADCASTING.
//
// Same as Offset(), except takes as NdArrayDesc<N> instead of Dims<N>.
inline int SubscriptToIndex(const NdArrayDesc<4>& desc, int i0, int i1, int i2,
int i3) {
inline int64_t SubscriptToIndex(const NdArrayDesc<4>& desc, int i0, int i1,
int i2, int i3) {
TFLITE_DCHECK(i0 >= 0 && i0 < desc.extents[0]);
TFLITE_DCHECK(i1 >= 0 && i1 < desc.extents[1]);
TFLITE_DCHECK(i2 >= 0 && i2 < desc.extents[2]);
@ -1028,13 +1028,13 @@ inline int SubscriptToIndex(const NdArrayDesc<4>& desc, int i0, int i1, int i2,
i3 * desc.strides[3];
}
inline int SubscriptToIndex(const NdArrayDesc<5>& desc, int indexes[5]) {
inline int64_t SubscriptToIndex(const NdArrayDesc<5>& desc, int indexes[5]) {
return indexes[0] * desc.strides[0] + indexes[1] * desc.strides[1] +
indexes[2] * desc.strides[2] + indexes[3] * desc.strides[3] +
indexes[4] * desc.strides[4];
}
inline int SubscriptToIndex(const NdArrayDesc<8>& desc, int indexes[8]) {
inline int64_t SubscriptToIndex(const NdArrayDesc<8>& desc, int indexes[8]) {
return indexes[0] * desc.strides[0] + indexes[1] * desc.strides[1] +
indexes[2] * desc.strides[2] + indexes[3] * desc.strides[3] +
indexes[4] * desc.strides[4] + indexes[5] * desc.strides[5] +
@ -1102,7 +1102,7 @@ inline void NdArrayDescsForElementwiseBroadcast(const Dims<N>& input0_dims,
template <int N>
TFLITE_NOINLINE void CopyDimsToDesc(const RuntimeShape& input_shape,
NdArrayDesc<N>* desc_out) {
int desc_stride = 1;
int64_t desc_stride = 1;
for (int i = N - 1; i >= 0; --i) {
desc_out->extents[i] = input_shape.Dims(i);
desc_out->strides[i] = desc_stride;
@ -1283,8 +1283,9 @@ inline int LegacyHowManyThreads(int max_num_threads, int rows, int cols,
static constexpr std::uint64_t min_cubic_size_per_thread = 64 * 1024;
// We can only multiply two out of three sizes without risking overflow
const std::uint64_t cubic_size =
std::uint64_t(rows) * std::uint64_t(cols) * std::uint64_t(depth);
const std::uint64_t cubic_size = static_cast<std::uint64_t>(rows) *
static_cast<std::uint64_t>(cols) *
static_cast<std::uint64_t>(depth);
thread_count = std::min(
thread_count, static_cast<int>(cubic_size / min_cubic_size_per_thread));

View file

@ -29,6 +29,53 @@ namespace tflite {
namespace reference_ops {
template <typename T>
inline void Slice(const std::vector<int>& begins,
const RuntimeShape& input_shape,
const RuntimeShape& output_shape,
SequentialTensorWriter<T>* writer) {
const int dims = input_shape.DimensionsCount();
std::vector<int> input_strides(dims);
std::vector<int> output_strides(dims);
if (dims == 0) {
writer->Write(0);
return;
}
input_strides[dims - 1] = 1;
output_strides[dims - 1] = 1;
for (int i = dims - 2; i >= 0; --i) {
input_strides[i] = input_strides[i + 1] * input_shape.Dims(i + 1);
output_strides[i] = output_strides[i + 1] * output_shape.Dims(i + 1);
}
for (int output_index = 0; output_index < output_shape.FlatSize();
++output_index) {
int remaining_index = output_index;
int input_index = 0;
for (int dim = 0; dim < dims; ++dim) {
const int coordinate = remaining_index / output_strides[dim];
remaining_index %= output_strides[dim];
input_index += (begins[dim] + coordinate) * input_strides[dim];
}
writer->Write(input_index);
}
}
template <typename T>
inline void Slice(const std::vector<int>& begins,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& output_shape, T* output_data) {
SequentialTensorWriter<T> writer(input_data, output_data);
return Slice(begins, input_shape, output_shape, &writer);
}
template <typename T>
inline void Slice(const std::vector<int>& begins,
const RuntimeShape& input_shape, const TfLiteTensor* input,
const RuntimeShape& output_shape, TfLiteTensor* output) {
SequentialTensorWriter<T> writer(input, output);
return Slice(begins, input_shape, output_shape, &writer);
}
template <typename T>
inline void Slice(const tflite::SliceParams& op_params,
const RuntimeShape& input_shape,
@ -102,6 +149,27 @@ inline void SliceInt4(const tflite::SliceParams& op_params,
GetTensorData<int8_t>(output));
}
inline void SliceInt4(const std::vector<int>& begins,
const RuntimeShape& input_shape,
const TfLiteTensor* input,
const RuntimeShape& output_shape, TfLiteTensor* output) {
const int num_input_elements = input_shape.FlatSize();
std::vector<int8_t> unpacked_input(num_input_elements);
tensor_utils::UnpackPackedIntToInt8(GetTensorData<int8_t>(input),
num_input_elements, 4,
unpacked_input.data());
const int num_output_elements = output_shape.FlatSize();
std::vector<int8_t> unpacked_output(num_output_elements);
reference_ops::Slice<int8_t>(begins, input_shape, unpacked_input.data(),
output_shape, unpacked_output.data());
tensor_utils::PackInt8IntoDenseInt(unpacked_output.data(),
num_output_elements, 4,
GetTensorData<int8_t>(output));
}
} // namespace reference_ops
} // namespace tflite

View file

@ -15,6 +15,9 @@ limitations under the License.
#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_STRIDED_SLICE_H_
#define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_STRIDED_SLICE_H_
#include <functional>
#include <vector>
#include "ruy/profiler/instrumentation.h" // from @ruy
#include "tensorflow/lite/kernels/internal/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
@ -26,6 +29,123 @@ namespace tflite {
namespace reference_ops {
struct DynamicStridedSliceParams {
std::vector<int32_t> start_indices;
std::vector<int32_t> stop_indices;
std::vector<int32_t> strides;
uint32_t begin_mask = 0;
uint32_t end_mask = 0;
uint32_t shrink_axis_mask = 0;
bool offset = false;
};
inline bool AxisMask(uint32_t mask, int axis) {
return mask & (uint32_t{1} << axis);
}
inline int StartForAxis(const DynamicStridedSliceParams& params,
const RuntimeShape& input_shape, int axis) {
const int axis_size = input_shape.Dims(axis);
int start = params.start_indices[axis];
const int stride = params.strides[axis];
if (start < 0) {
start += axis_size;
}
if (stride > 0) {
start = strided_slice::Clamp(start, 0, axis_size);
} else {
start = strided_slice::Clamp(start, -1, axis_size - 1);
}
if (AxisMask(params.begin_mask, axis)) {
start = stride > 0 ? 0 : axis_size - 1;
}
return start;
}
inline int EndForAxis(const DynamicStridedSliceParams& params,
const RuntimeShape& input_shape, int axis, int start) {
const bool shrink_axis = AxisMask(params.shrink_axis_mask, axis);
const int axis_size = input_shape.Dims(axis);
if (shrink_axis) {
return start >= axis_size ? start : start + 1;
}
int end = params.stop_indices[axis];
if (params.offset) {
end += start;
}
const int stride = params.strides[axis];
if (end < 0) {
end += axis_size;
}
if (stride > 0) {
end = strided_slice::Clamp(end, 0, axis_size);
} else {
end = strided_slice::Clamp(end, -1, axis_size - 1);
}
if (AxisMask(params.end_mask, axis)) {
end = stride > 0 ? axis_size : -1;
}
return end;
}
template <typename T>
inline void StridedSlice(const DynamicStridedSliceParams& op_params,
const RuntimeShape& input_shape,
const RuntimeShape& output_shape,
SequentialTensorWriter<T>* writer) {
ruy::profiler::ScopeLabel label("StridedSlice");
const int dims = input_shape.DimensionsCount();
std::vector<int> starts(dims);
std::vector<int> stops(dims);
std::vector<int> input_strides(dims);
if (dims == 0) {
writer->Write(0);
return;
}
input_strides[dims - 1] = 1;
for (int i = dims - 2; i >= 0; --i) {
input_strides[i] = input_strides[i + 1] * input_shape.Dims(i + 1);
}
for (int axis = 0; axis < dims; ++axis) {
starts[axis] = StartForAxis(op_params, input_shape, axis);
stops[axis] = EndForAxis(op_params, input_shape, axis, starts[axis]);
}
auto loop_condition = [](int index, int stop, int stride) {
return stride > 0 ? index < stop : index > stop;
};
std::function<void(int, int)> write_slice = [&](int axis, int input_index) {
if (axis == dims) {
writer->Write(input_index);
return;
}
for (int offset = starts[axis];
loop_condition(offset, stops[axis], op_params.strides[axis]);
offset += op_params.strides[axis]) {
write_slice(axis + 1, input_index + offset * input_strides[axis]);
}
};
write_slice(/*axis=*/0, /*input_index=*/0);
}
template <typename T>
inline void StridedSlice(const DynamicStridedSliceParams& op_params,
const RuntimeShape& input_shape, const T* input_data,
const RuntimeShape& output_shape, T* output_data) {
SequentialTensorWriter<T> writer(input_data, output_data);
StridedSlice<T>(op_params, input_shape, output_shape, &writer);
}
template <typename T>
inline void StridedSlice(const DynamicStridedSliceParams& op_params,
const RuntimeShape& input_shape,
const TfLiteTensor* input,
const RuntimeShape& output_shape,
TfLiteTensor* output) {
SequentialTensorWriter<T> writer(input, output);
StridedSlice<T>(op_params, input_shape, output_shape, &writer);
}
template <typename T>
inline void StridedSlice(const tflite::StridedSliceParams& op_params,
const RuntimeShape& unextended_input_shape,

View file

@ -18,9 +18,9 @@ import os
import subprocess
import sys
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema # pylint:disable=g-direct-tensorflow-import
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import flatbuffer_utils
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema # pylint:disable=g-direct-tensorflow-import
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import flatbuffer_utils
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test

View file

@ -18,7 +18,7 @@ All functions that can be commonly used by various tests.
"""
import flatbuffers
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
TFLITE_SCHEMA_VERSION = 3

View file

@ -16,8 +16,8 @@
import os
import re
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import visualize
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import test_utils
from tflite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite_micro.tensorflow.lite.tools import visualize
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test