From 5e3c68cc0d50a28dee9bbfacfc0facbf141add0c Mon Sep 17 00:00:00 2001 From: unmeshna Date: Sun, 16 Aug 2026 23:58:35 -0700 Subject: [PATCH 1/2] Add Xtensa reference/trivial data-movement kernels Adds Xtensa kernels for fill and circular_buffer (via xa_nn_memset/xa_nn_memmove) plus reference expand_dims, pack, unpack, split, and split_v. Each shares a filename with its reference kernel and is auto-substituted by specialize_files.py, so no xtensa.inc changes are required. --- .../micro/kernels/xtensa/circular_buffer.cc | 127 ++++++++++++++ .../lite/micro/kernels/xtensa/expand_dims.cc | 160 ++++++++++++++++++ tensorflow/lite/micro/kernels/xtensa/fill.cc | 153 +++++++++++++++++ tensorflow/lite/micro/kernels/xtensa/pack.cc | 119 +++++++++++++ tensorflow/lite/micro/kernels/xtensa/split.cc | 125 ++++++++++++++ .../lite/micro/kernels/xtensa/split_v.cc | 127 ++++++++++++++ .../lite/micro/kernels/xtensa/unpack.cc | 111 ++++++++++++ 7 files changed, 922 insertions(+) create mode 100644 tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc create mode 100644 tensorflow/lite/micro/kernels/xtensa/expand_dims.cc create mode 100644 tensorflow/lite/micro/kernels/xtensa/fill.cc create mode 100644 tensorflow/lite/micro/kernels/xtensa/pack.cc create mode 100644 tensorflow/lite/micro/kernels/xtensa/split.cc create mode 100644 tensorflow/lite/micro/kernels/xtensa/split_v.cc create mode 100644 tensorflow/lite/micro/kernels/xtensa/unpack.cc diff --git a/tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc b/tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc new file mode 100644 index 00000000000..b0440705ff0 --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc @@ -0,0 +1,127 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/micro/kernels/circular_buffer.h" + +#include "tensorflow/lite/c/builtin_op_data.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/compatibility.h" +#include "tensorflow/lite/kernels/internal/quantization_util.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" +#include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/kernels/op_macros.h" +#include "tensorflow/lite/micro/flatbuffer_utils.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" +#include "tensorflow/lite/micro/kernels/xtensa/xtensa.h" + +/* + * The circular buffer custom operator is used to implement strided streaming + * convolutions on TFLite Micro. Each time this operator is invoked, it checks + * whether or not to run, based on a predetermined stride in time. If the op + * runs, it inserts the input into the end of the output buffer and shifts the + * output values towards the start of the buffer. It discards the oldest value + * in the output buffer. + * + * Input: [, , , ] + * + * After shifting: + * Output: [, , , ] + * + * We make some assumptions in this custom operator: + * - Input shape must be [1, 1, 1, depth] + * - Output shape must be [1, num_slots, 1, depth] + * - Input and output types must match. + * - Input and output quantization params must be identical. + */ +namespace tflite { + +void* CircularBufferInit(TfLiteContext* context, const char* buffer, + size_t length) { + TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr); + OpDataCircularBuffer* op_data = static_cast( + context->AllocatePersistentBuffer(context, sizeof(OpDataCircularBuffer))); + + if (buffer != nullptr && length > 0) { + const uint8_t* buffer_t = reinterpret_cast(buffer); + tflite::FlexbufferWrapper wrapper(buffer_t, length); + op_data->cycles_max = wrapper.ElementAsInt32(kCircularBufferCyclesMaxIndex); + } else { + op_data->cycles_max = 0; + } + + return op_data; +} + +// Shifts buffer over by the output depth, and write new input to end of buffer. +// num_slots is the number of samples stored in the output buffer. +// depth is the size of each sample. +void EvalInt8(const int8_t* input, int num_slots, int depth, int8_t* output) { + memmove(output, &output[depth], (num_slots - 1) * depth); + memcpy(&output[(num_slots - 1) * depth], input, depth); +} + +TfLiteStatus CircularBufferEval(TfLiteContext* context, TfLiteNode* node) { + const TfLiteEvalTensor* input = + tflite::micro::GetEvalInput(context, node, kCircularBufferInputTensor); + TfLiteEvalTensor* output = + tflite::micro::GetEvalOutput(context, node, kCircularBufferOutputTensor); + + TFLITE_DCHECK(node->user_data != nullptr); + OpDataCircularBuffer* data = + reinterpret_cast(node->user_data); + + int num_slots = output->dims->data[1]; + int depth = output->dims->data[2] * output->dims->data[3]; + + if (input->type == kTfLiteInt8) { +#if defined(HIFI5) || defined(HIFI4) || defined(HIFI_IQ) + const int8_t* xa_input; + int8_t* xa_output; + int err; + xa_input = tflite::micro::GetTensorData(input); + xa_output =tflite::micro::GetTensorData(output); + err = xa_nn_memmove_8_8(xa_output, &xa_output[depth], (num_slots-1)*depth); + TF_LITE_ENSURE(context, (err==0) ); + memcpy(&xa_output[(num_slots - 1) * depth], xa_input, depth); +#else + EvalInt8(tflite::micro::GetTensorData(input), num_slots, depth, + tflite::micro::GetTensorData(output)); +#endif // defined(HIFI5) || defined(HIFI4) + } else { + TF_LITE_KERNEL_LOG(context, "Type %s (%d) not supported.", + TfLiteTypeGetName(input->type), input->type); + return kTfLiteError; + } + + if (--data->cycles_until_run != 0) { + // Signal the interpreter to end current run if the delay before op invoke + // has not been reached. + // TODO(b/149795762): Add kTfLiteAbort to TfLiteStatus enum. + return static_cast(kTfLiteAbort); + } + + data->cycles_until_run = data->cycles_max; + + return kTfLiteOk; +} + +TFLMRegistration* Register_CIRCULAR_BUFFER() { + static TFLMRegistration r = tflite::micro::RegisterOp(CircularBufferInit, CircularBufferPrepare, CircularBufferEval); + return &r; +} + +} // namespace tflite diff --git a/tensorflow/lite/micro/kernels/xtensa/expand_dims.cc b/tensorflow/lite/micro/kernels/xtensa/expand_dims.cc new file mode 100644 index 00000000000..8e8033badbb --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/expand_dims.cc @@ -0,0 +1,160 @@ +/* Copyright 2021 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 + +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" +#include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" +#include "tensorflow/lite/micro/micro_log.h" +#include "tensorflow/lite/micro/micro_utils.h" + +namespace tflite { +namespace { + +constexpr int kInputTensor = 0; +constexpr int kAxisTensor = 1; +constexpr int kOutputTensor = 0; + +TfLiteStatus GetAxisValueFromTensor(TfLiteContext* context, + const TfLiteTensor* axis, + int32_t* axis_value) { + const int axis_dims = (tflite::GetTensorShape(axis)).DimensionsCount(); + if (axis_dims > 1) { + MicroPrintf("Axis has only one element for Expand_Dims.", axis_dims); + return kTfLiteError; + } + + if (kTfLiteInt32 == (axis->type)) { + const int32_t* axis_ptr = tflite::GetTensorData(axis); + *axis_value = axis_ptr[0]; + return kTfLiteOk; + } else { + MicroPrintf("Axis type %s (%d) not supported by Expand_Dims.", + TfLiteTypeGetName(axis->type), axis->type); + return kTfLiteError; + } +} + +// Verifies that the output tensor's dimension shape is equivalent to inserting +// a dimension of length 1 at the dimension index axis of input's shape as +// defined in https://www.tensorflow.org/api_docs/python/tf/expand_dims. +TfLiteStatus VerifyTensorDim(TfLiteContext* context, const TfLiteTensor* input, + const TfLiteTensor* axis_tensor, + const TfLiteTensor* output) { + int32_t axis_value = 0; + TF_LITE_ENSURE_OK(context, + GetAxisValueFromTensor(context, axis_tensor, &axis_value)); + + tflite::RuntimeShape input_shape = tflite::GetTensorShape(input); + if (axis_value < 0) { + axis_value = input_shape.DimensionsCount() + 1 + axis_value; + } + TF_LITE_ENSURE(context, axis_value <= input_shape.DimensionsCount()); + + // TFLM only supports fixed dimension tensor and assumes that the output shape + // is fully specified in the model. As such, TFLM directly use the pointer to + // the dimension array in the model buffer. + tflite::RuntimeShape output_shape = tflite::GetTensorShape(output); + + TF_LITE_ENSURE(context, output_shape.DimensionsCount() == + input_shape.DimensionsCount() + 1); + for (int i = 0; i < output_shape.DimensionsCount(); ++i) { + if (i < axis_value) { + TF_LITE_ENSURE(context, output_shape.Dims(i) == input_shape.Dims(i)); + } else if (i == axis_value) { + TF_LITE_ENSURE(context, output_shape.Dims(i) == 1); + } else { + TF_LITE_ENSURE(context, output_shape.Dims(i) == input_shape.Dims(i - 1)); + } + } + return kTfLiteOk; +} + +TfLiteStatus ExpandDimsPrepare(TfLiteContext* context, TfLiteNode* node) { + MicroContext* micro_context = GetMicroContext(context); + + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + TfLiteTensor* input = + micro_context->AllocateTempInputTensor(node, kInputTensor); + TF_LITE_ENSURE(context, input != nullptr); + TfLiteTensor* axis = + micro_context->AllocateTempInputTensor(node, kAxisTensor); + TF_LITE_ENSURE(context, axis != nullptr); + TfLiteTensor* output = + micro_context->AllocateTempOutputTensor(node, kOutputTensor); + TF_LITE_ENSURE(context, output != nullptr); + output->type = input->type; + if (IsDynamicTensor(axis)) { + MicroPrintf("DynamicTensor is not yet supported by Expand_Dims."); + return kTfLiteError; + } + TF_LITE_ENSURE_OK(context, VerifyTensorDim(context, input, axis, output)); + + micro_context->DeallocateTempTfLiteTensor(input); + micro_context->DeallocateTempTfLiteTensor(axis); + micro_context->DeallocateTempTfLiteTensor(output); + return kTfLiteOk; +} + +template +void memCopyN(T* out, const T* in, const int num_elements) { +#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + memcpy(out, in, num_elements * sizeof(T)); +#else // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + for (int i = 0; i < num_elements; ++i) { + out[i] = in[i]; + } +#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) +} + +TfLiteStatus ExpandDimsEval(TfLiteContext* context, TfLiteNode* node) { + const TfLiteEvalTensor* input = + tflite::micro::GetEvalInput(context, node, kInputTensor); + TfLiteEvalTensor* output = + tflite::micro::GetEvalOutput(context, node, kOutputTensor); + const int flat_size = ElementCount(*input->dims); + + switch (input->type) { + case kTfLiteFloat32: { + memCopyN(tflite::micro::GetTensorData(output), + tflite::micro::GetTensorData(input), flat_size); + } break; + case kTfLiteInt16: { + memCopyN(tflite::micro::GetTensorData(output), + tflite::micro::GetTensorData(input), flat_size); + } break; + case kTfLiteInt8: { + memCopyN(tflite::micro::GetTensorData(output), + tflite::micro::GetTensorData(input), flat_size); + } break; + default: + MicroPrintf( + "Expand_Dims only currently supports int8, int16 and float32, got " + "%d.", + input->type); + return kTfLiteError; + } + return kTfLiteOk; +} +} // namespace + +TFLMRegistration Register_EXPAND_DIMS() { + return tflite::micro::RegisterOp(nullptr, ExpandDimsPrepare, ExpandDimsEval); +} + +} // namespace tflite diff --git a/tensorflow/lite/micro/kernels/xtensa/fill.cc b/tensorflow/lite/micro/kernels/xtensa/fill.cc new file mode 100644 index 00000000000..c7c2e976f5c --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/fill.cc @@ -0,0 +1,153 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/kernels/internal/reference/fill.h" + +#include + +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" +#include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" +#include "tensorflow/lite/micro/kernels/xtensa/xtensa.h" + +namespace tflite { + +namespace { + +template +TfLiteStatus EnsureEqImpl(TfLiteContext* context, const TfLiteIntArray* array, + const TfLiteTensor* tensor) { + for (int i = 0; i < array->size; ++i) { + TF_LITE_ENSURE_EQ(context, array->data[i], GetTensorData(tensor)[i]); + } + return kTfLiteOk; +} + +// Ensure the equality of an int array and a tensor, which must be +// one-dimensional and of an integer type. +TfLiteStatus EnsureEq(TfLiteContext* context, const TfLiteIntArray* array, + const TfLiteTensor* tensor) { + TF_LITE_ENSURE_EQ(context, NumDimensions(tensor), 1); + const auto tensor_len = tensor->dims->data[0]; + TF_LITE_ENSURE_EQ(context, array->size, tensor_len); + + switch (tensor->type) { + case kTfLiteInt8: + return EnsureEqImpl(context, array, tensor); + case kTfLiteInt16: + return EnsureEqImpl(context, array, tensor); + case kTfLiteInt32: + return EnsureEqImpl(context, array, tensor); + case kTfLiteInt64: + return EnsureEqImpl(context, array, tensor); + default: + TF_LITE_KERNEL_LOG(context, + "cannot compare int array to tensor of type %d.", + tensor->type); + return kTfLiteError; + } +} + +constexpr int kDimsTensor = 0; +constexpr int kValueTensor = 1; +constexpr int kOutputTensor = 0; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + MicroContext* micro_context = GetMicroContext(context); + + // Ensure inputs and outputs exist. + TfLiteTensor* dims = + micro_context->AllocateTempInputTensor(node, kDimsTensor); + TF_LITE_ENSURE(context, dims != nullptr); + TfLiteTensor* value = + micro_context->AllocateTempInputTensor(node, kValueTensor); + TF_LITE_ENSURE(context, value != nullptr); + TfLiteTensor* output = + micro_context->AllocateTempOutputTensor(node, kOutputTensor); + TF_LITE_ENSURE(context, output != nullptr); + + // The value tensor must be a scalar. + TF_LITE_ENSURE_EQ(context, NumDimensions(value), 0); + + // The value type and output type must match. + TF_LITE_ENSURE_EQ(context, value->type, output->type); + + // The dimension of the output tensor is known in model already. + TFLITE_DCHECK(output->dims != nullptr); + + if (dims->data.data != nullptr) { + // When the dims tensor is specified in model already (i.e. is not an + // activation tensor), the dims tensor must match the output tensor shape. + // As a byproduct, ensures the dims tensor is of an integer type. + TF_LITE_ENSURE_OK(context, EnsureEq(context, output->dims, dims)); + } + + micro_context->DeallocateTempTfLiteTensor(dims); + micro_context->DeallocateTempTfLiteTensor(value); + micro_context->DeallocateTempTfLiteTensor(output); + return kTfLiteOk; +} + +template +void FillImpl(const TfLiteEvalTensor* value, TfLiteEvalTensor* output) { + reference_ops::Fill( + micro::GetTensorShape(value), micro::GetTensorData(value), + micro::GetTensorShape(output), micro::GetTensorData(output)); +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + const TfLiteEvalTensor* value = + micro::GetEvalInput(context, node, kValueTensor); + TfLiteEvalTensor* output = micro::GetEvalOutput(context, node, kOutputTensor); + + switch (value->type) { + case kTfLiteFloat32: + { +#if defined(INCLUDE_FLOAT_OPT) + float memsetValue = *(float *)micro::GetTensorData(value); + int numElem = micro::GetTensorShape(output).FlatSize(); + float *dst = micro::GetTensorData(output); + int err; + err = xa_nn_memset_f32_f32(dst,memsetValue, numElem); + TF_LITE_ENSURE(context, (err==0) ); +#else + FillImpl(value, output); +#endif // defined(INCLUDE_FLOAT_OPT) + break; + } + case kTfLiteInt32: + FillImpl(value, output); + break; + case kTfLiteInt8: + FillImpl(value, output); + break; + default: + TF_LITE_KERNEL_LOG( + context, "Fill only currently supports float32 for input 1, got %d.", + TfLiteTypeGetName(value->type)); + return kTfLiteError; + } + + return kTfLiteOk; +} + +} // namespace + +TFLMRegistration Register_FILL() { + return tflite::micro::RegisterOp(nullptr, Prepare, Eval); +} + +} // namespace tflite diff --git a/tensorflow/lite/micro/kernels/xtensa/pack.cc b/tensorflow/lite/micro/kernels/xtensa/pack.cc new file mode 100644 index 00000000000..33b9b624374 --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/pack.cc @@ -0,0 +1,119 @@ +/* Copyright 2021 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/c/builtin_op_data.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" + +namespace tflite { + +namespace { + +constexpr int kOutputTensor = 0; + +template +TfLiteStatus PackImpl(TfLiteContext* context, TfLiteNode* node, + TfLiteEvalTensor* output, int values_count, int axis) { + const TfLiteEvalTensor* input0 = + tflite::micro::GetEvalInput(context, node, 0); + + const int dimensions = output->dims->size; + const TfLiteIntArray* input_dims = input0->dims; + const TfLiteIntArray* output_dims = output->dims; + + if (axis < 0) { + axis += dimensions; + } + + int outer_size = 1; + for (int i = 0; i < axis; ++i) { + outer_size *= output_dims->data[i]; + } + int copy_size = 1; + for (int i = axis + 1; i < dimensions; ++i) { + copy_size *= output_dims->data[i]; + } + int input_size = 1; + for (int i = 0; i < input_dims->size; ++i) { + input_size *= input_dims->data[i]; + } + TFLITE_DCHECK_EQ(input_size, copy_size * outer_size); + + T* output_data = tflite::micro::GetTensorData(output); + + for (int i = 0; i < values_count; ++i) { + const TfLiteEvalTensor* t = tflite::micro::GetEvalInput(context, node, i); + const T* input_data = tflite::micro::GetTensorData(t); + for (int k = 0; k < outer_size; ++k) { + const T* input_ptr = input_data + copy_size * k; + int loc = k * values_count * copy_size + i * copy_size; + T* output_ptr = output_data + loc; +#if defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + memcpy(output_ptr, input_ptr, copy_size*sizeof(T)); +#else + for (int j = 0; j < copy_size; ++j) output_ptr[j] = input_ptr[j]; +#endif // defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) + } + } + + return kTfLiteOk; +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + const TfLitePackParams* data = + reinterpret_cast(node->builtin_data); + + TfLiteEvalTensor* output = + tflite::micro::GetEvalOutput(context, node, kOutputTensor); + + switch (output->type) { + case kTfLiteFloat32: { + return PackImpl(context, node, output, data->values_count, + data->axis); + } + case kTfLiteInt8: { + return PackImpl(context, node, output, data->values_count, + data->axis); + } + case kTfLiteInt16: { + return PackImpl(context, node, output, data->values_count, + data->axis); + } + case kTfLiteInt32: { + return PackImpl(context, node, output, data->values_count, + data->axis); + } + case kTfLiteInt64: { + return PackImpl(context, node, output, data->values_count, + data->axis); + } + default: { + TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by pack.", + TfLiteTypeGetName(output->type)); + return kTfLiteError; + } + } + + return kTfLiteOk; +} + +} // namespace + +TFLMRegistration Register_PACK() { + return tflite::micro::RegisterOp(nullptr, nullptr, Eval); +} + +} // namespace tflite diff --git a/tensorflow/lite/micro/kernels/xtensa/split.cc b/tensorflow/lite/micro/kernels/xtensa/split.cc new file mode 100644 index 00000000000..61e6b04b785 --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/split.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/lite/c/builtin_op_data.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" +#include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" +#include "tensorflow/lite/micro/micro_log.h" + +namespace tflite { + +namespace { + +template +TfLiteStatus SplitImpl(TfLiteContext* context, TfLiteNode* node, + const TfLiteEvalTensor* input, int axis_value) { + const int output_count = NumOutputs(node); + const TfLiteIntArray* input_dims = input->dims; + const TfLiteEvalTensor* output0 = + tflite::micro::GetEvalOutput(context, node, 0); + const TfLiteIntArray* output_dims = output0->dims; + + const int split_dimensions = input_dims->size; + int axis = axis_value < 0 ? axis_value + split_dimensions : axis_value; + + TFLITE_DCHECK_LT(axis, split_dimensions); + TFLITE_DCHECK_EQ(output_dims->size, split_dimensions); + + int64_t split_size = output_dims->data[axis] * output_count; + + TFLITE_DCHECK_EQ(split_size, input_dims->data[axis]); + int64_t outer_size = 1; + for (int i = 0; i < axis; ++i) { + outer_size *= input_dims->data[i]; + } + + int64_t base_inner_size = 1; + for (int i = axis + 1; i < split_dimensions; ++i) { + base_inner_size *= input_dims->data[i]; + } + + const T* input_ptr = tflite::micro::GetTensorData(input); + for (int k = 0; k < outer_size; ++k) { + for (int i = 0; i < output_count; ++i) { + TfLiteEvalTensor* t = tflite::micro::GetEvalOutput(context, node, i); + T* output_data = tflite::micro::GetTensorData(t); + const int copy_size = output_dims->data[axis] * base_inner_size; + T* output_ptr = output_data + k * copy_size; + memcpy(output_ptr, input_ptr, (copy_size * sizeof(T))); + input_ptr += copy_size; + } + } + + return kTfLiteOk; +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + MicroContext* micro_context = GetMicroContext(context); + TfLiteTensor* axis = micro_context->AllocateTempInputTensor(node, 0); + TF_LITE_ENSURE(context, axis != nullptr); + + // Dynamic output tensors are needed if axis tensor is not constant. + // But Micro doesn't support dynamic memory allocation, so we only support + // constant axis tensor for now. + TF_LITE_ENSURE_MSG(context, IsConstantTensor(axis), + "Non constant axis tensor not supported"); + + micro_context->DeallocateTempTfLiteTensor(axis); + return kTfLiteOk; +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + const TfLiteEvalTensor* axis = tflite::micro::GetEvalInput(context, node, 0); + const TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, 1); + + int axis_value = tflite::micro::GetTensorData(axis)[0]; + if (axis_value < 0) { + axis_value += input->dims->size; + } + + TF_LITE_ENSURE(context, axis_value >= 0); + TF_LITE_ENSURE(context, axis_value < input->dims->size); + + switch (input->type) { + case kTfLiteFloat32: { + return SplitImpl(context, node, input, axis_value); + } + case kTfLiteInt8: { + return SplitImpl(context, node, input, axis_value); + } + case kTfLiteInt16: { + return SplitImpl(context, node, input, axis_value); + } + case kTfLiteInt32: { + return SplitImpl(context, node, input, axis_value); + } + default: + MicroPrintf("Type %s currently not supported.", + TfLiteTypeGetName(input->type)); + return kTfLiteError; + } + + return kTfLiteOk; +} + +} // namespace + +TFLMRegistration Register_SPLIT() { + return tflite::micro::RegisterOp(nullptr, Prepare, Eval); +} + +} // namespace tflite diff --git a/tensorflow/lite/micro/kernels/xtensa/split_v.cc b/tensorflow/lite/micro/kernels/xtensa/split_v.cc new file mode 100644 index 00000000000..11af32ac004 --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/split_v.cc @@ -0,0 +1,127 @@ +/* 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 "tensorflow/lite/c/builtin_op_data.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" +#include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/kernels/op_macros.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" +#include "tensorflow/lite/micro/micro_log.h" + +namespace tflite { + +namespace { + +template +TfLiteStatus SplitImpl(TfLiteContext* context, TfLiteNode* node, + const TfLiteEvalTensor* input, int axis_value) { + const TfLiteIntArray* input_dims = input->dims; + const TfLiteEvalTensor* output0 = + tflite::micro::GetEvalOutput(context, node, 0); + + const int split_dimensions = input_dims->size; + + TFLITE_DCHECK_LT(axis_value, split_dimensions); + TFLITE_DCHECK_EQ(output0->dims->size, split_dimensions); + + int64_t split_size = 0; + const int output_count = NumOutputs(node); + for (int i = 0; i < output_count; i++) { + split_size += + tflite::micro::GetEvalOutput(context, node, i)->dims->data[axis_value]; + } + TFLITE_DCHECK_EQ(split_size, input_dims->data[axis_value]); + int64_t outer_size = 1; + for (int i = 0; i < axis_value; ++i) { + outer_size *= input_dims->data[i]; + } + + int64_t base_inner_size = 1; + for (int i = axis_value + 1; i < split_dimensions; ++i) { + base_inner_size *= input_dims->data[i]; + } + + const T* input_ptr = tflite::micro::GetTensorData(input); + for (int k = 0; k < outer_size; ++k) { + for (int i = 0; i < output_count; ++i) { + TfLiteEvalTensor* output_tensor = + tflite::micro::GetEvalOutput(context, node, i); + T* output_data = tflite::micro::GetTensorData(output_tensor); + const int copy_size = + output_tensor->dims->data[axis_value] * base_inner_size; + T* output_ptr = output_data + k * copy_size; + memcpy(output_ptr, input_ptr, (copy_size * sizeof(T))); + input_ptr += copy_size; + } + } + + return kTfLiteOk; +} + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 3); + + MicroContext* micro_context = GetMicroContext(context); + // Dynamic output tensors are needed if axis tensor is not constant. + // But Micro doesn't support dynamic memory allocation, so we only support + // constant axis tensor for now. + TfLiteTensor* axis = micro_context->AllocateTempInputTensor(node, 2); + TF_LITE_ENSURE_MSG(context, IsConstantTensor(axis), + "Non constant axis tensor not supported"); + micro_context->DeallocateTempTfLiteTensor(axis); + return kTfLiteOk; +} + +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + const TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, 0); + const TfLiteEvalTensor* axis = tflite::micro::GetEvalInput(context, node, 2); + + int axis_value = tflite::micro::GetTensorData(axis)[0]; + if (axis_value < 0) { + axis_value += input->dims->size; + } + + TF_LITE_ENSURE(context, axis_value >= 0); + TF_LITE_ENSURE(context, axis_value < input->dims->size); + + switch (input->type) { + case kTfLiteFloat32: { + return SplitImpl(context, node, input, axis_value); + } + case kTfLiteInt8: { + return SplitImpl(context, node, input, axis_value); + } + case kTfLiteInt16: { + return SplitImpl(context, node, input, axis_value); + } + case kTfLiteInt32: { + return SplitImpl(context, node, input, axis_value); + } + default: + MicroPrintf("Type %s currently not supported.", + TfLiteTypeGetName(input->type)); + return kTfLiteError; + } + return kTfLiteOk; +} + +} // namespace + +TFLMRegistration Register_SPLIT_V() { + return tflite::micro::RegisterOp(nullptr, Prepare, Eval); +} + +} // namespace tflite diff --git a/tensorflow/lite/micro/kernels/xtensa/unpack.cc b/tensorflow/lite/micro/kernels/xtensa/unpack.cc new file mode 100644 index 00000000000..b30bec519da --- /dev/null +++ b/tensorflow/lite/micro/kernels/xtensa/unpack.cc @@ -0,0 +1,111 @@ +/* Copyright 2025 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/c/builtin_op_data.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" +#include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/micro/kernels/kernel_util.h" +#include "tensorflow/lite/micro/micro_log.h" + +namespace tflite { + +namespace { + +constexpr int kInputTensor = 0; + +template +TfLiteStatus UnpackImpl(TfLiteContext* context, TfLiteNode* node, + const TfLiteEvalTensor* input, int output_count, + int axis) { + const TfLiteEvalTensor* output0 = + tflite::micro::GetEvalOutput(context, node, 0); + const TfLiteIntArray* input_dims = input->dims; + const TfLiteIntArray* output_dims = output0->dims; + const int dimensions = input_dims->size; + + if (axis < 0) { + axis += input->dims->size; + } + + TFLITE_DCHECK_LT(axis, dimensions); + + int outer_size = 1; + for (int i = 0; i < axis; ++i) { + outer_size *= input_dims->data[i]; + } + int copy_size = 1; + for (int i = axis + 1; i < dimensions; ++i) { + copy_size *= input_dims->data[i]; + } + int output_size = 1; + for (int i = 0; i < output_dims->size; ++i) { + output_size *= output_dims->data[i]; + } + TFLITE_DCHECK_EQ(output_size, copy_size * outer_size); + + const T* input_data = tflite::micro::GetTensorData(input); + + for (int i = 0; i < output_count; ++i) { + TfLiteEvalTensor* t = tflite::micro::GetEvalOutput(context, node, i); + T* output_data = tflite::micro::GetTensorData(t); + for (int k = 0; k < outer_size; ++k) { + T* output_ptr = output_data + copy_size * k; + int loc = k * output_count * copy_size + i * copy_size; + const T* input_ptr = input_data + loc; + memcpy(output_ptr, input_ptr, (copy_size * sizeof(T))); + } + } + + return kTfLiteOk; +} + +TfLiteStatus UnpackEval(TfLiteContext* context, TfLiteNode* node) { + TfLiteUnpackParams* data = + reinterpret_cast(node->builtin_data); + + const TfLiteEvalTensor* input = + tflite::micro::GetEvalInput(context, node, kInputTensor); + + switch (input->type) { + case kTfLiteFloat32: { + return UnpackImpl(context, node, input, data->num, data->axis); + } + case kTfLiteInt32: { + return UnpackImpl(context, node, input, data->num, data->axis); + } + case kTfLiteInt8: { + return UnpackImpl(context, node, input, data->num, data->axis); + } + case kTfLiteInt16: { + return UnpackImpl(context, node, input, data->num, data->axis); + } + default: { + MicroPrintf("Type '%s' is not supported by unpack.", + TfLiteTypeGetName(input->type)); + return kTfLiteError; + } + } + + return kTfLiteOk; +} + +} // namespace + +TFLMRegistration Register_UNPACK() { + return tflite::micro::RegisterOp(nullptr, nullptr, UnpackEval); +} + +} // namespace tflite From 56a1957e43bf81a197ec134a60f861cd25d1c590 Mon Sep 17 00:00:00 2001 From: unmeshna Date: Mon, 17 Aug 2026 01:40:39 -0700 Subject: [PATCH 2/2] Sync xtensa fill/circular_buffer with upstream reference kernels fill: require constant dims tensor (IsConstantTensor) matching master. circular_buffer: add CircularBufferReset reset callback to registration. --- .../lite/micro/kernels/xtensa/circular_buffer.cc | 13 ++++++++++++- tensorflow/lite/micro/kernels/xtensa/fill.cc | 11 +++++------ 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc b/tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc index b0440705ff0..f0b940392d0 100644 --- a/tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc +++ b/tensorflow/lite/micro/kernels/xtensa/circular_buffer.cc @@ -119,8 +119,19 @@ TfLiteStatus CircularBufferEval(TfLiteContext* context, TfLiteNode* node) { return kTfLiteOk; } +// Restores period counter (cycles_until_run) on reset. Buffer memory cleanup +// is not needed here: because the output tensor is a variable tensor. +// ResetVariableTensors() automatically zero-points its memory upon reset. +void CircularBufferReset(TfLiteContext* context, void* buffer) { + TFLITE_DCHECK(buffer != nullptr); + OpDataCircularBuffer* data = static_cast(buffer); + data->cycles_until_run = data->cycles_max; +} + TFLMRegistration* Register_CIRCULAR_BUFFER() { - static TFLMRegistration r = tflite::micro::RegisterOp(CircularBufferInit, CircularBufferPrepare, CircularBufferEval); + static TFLMRegistration r = tflite::micro::RegisterOp( + CircularBufferInit, CircularBufferPrepare, CircularBufferEval, + /*free=*/nullptr, CircularBufferReset); return &r; } diff --git a/tensorflow/lite/micro/kernels/xtensa/fill.cc b/tensorflow/lite/micro/kernels/xtensa/fill.cc index c7c2e976f5c..37f874c3e54 100644 --- a/tensorflow/lite/micro/kernels/xtensa/fill.cc +++ b/tensorflow/lite/micro/kernels/xtensa/fill.cc @@ -88,12 +88,11 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // The dimension of the output tensor is known in model already. TFLITE_DCHECK(output->dims != nullptr); - if (dims->data.data != nullptr) { - // When the dims tensor is specified in model already (i.e. is not an - // activation tensor), the dims tensor must match the output tensor shape. - // As a byproduct, ensures the dims tensor is of an integer type. - TF_LITE_ENSURE_OK(context, EnsureEq(context, output->dims, dims)); - } + TF_LITE_ENSURE_MSG(context, IsConstantTensor(dims), + "Non-constant >dims< tensor is not supported"); + // The dims tensor must match the output tensor shape. + // As a byproduct, ensures the dims tensor is of an integer type. + TF_LITE_ENSURE_OK(context, EnsureEq(context, output->dims, dims)); micro_context->DeallocateTempTfLiteTensor(dims); micro_context->DeallocateTempTfLiteTensor(value);