Skip to content

Add HiFi-optimized Softmax kernels for Xtensa #3729

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
unmeshna017 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add HiFi-optimized Softmax kernels for Xtensa #3729

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
792f766
Add HiFi-optimized Softmax kernels for Xtensa
unmeshna017 Sep 16, 2026
690bd50
Include NNLib headers for HIFI_IQ target in xtensa.h
unmeshna017 Sep 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion tensorflow/lite/micro/kernels/softmax.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ inline TFLMRegistration Register_SOFTMAX_INT8_INT16() {
}
#endif

#if defined(CMSIS_NN)
#if defined(CMSIS_NN) || defined(XTENSA)
// Returns a TFLMRegistration struct for kernel variant that only supports
// int8 input/output and uses the latency optimized implementations.
TFLMRegistration Register_SOFTMAX_INT8();
Expand All @@ -62,6 +62,10 @@ inline TFLMRegistration Register_SOFTMAX_INT8() { return Register_SOFTMAX(); }
inline TFLMRegistration Register_SOFTMAX_INT16() { return Register_SOFTMAX(); }
#endif

#if defined(XTENSA)
TFLMRegistration Register_SOFTMAX_FLOAT32();
#endif

} // namespace tflite

#endif // TENSORFLOW_LITE_MICRO_KERNELS_SOFTMAX_H_
147 changes: 141 additions & 6 deletions tensorflow/lite/micro/kernels/xtensa/softmax.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,10 @@ limitations under the License.
#include "tensorflow/lite/micro/kernels/xtensa/xtensa.h"
#include "tensorflow/lite/micro/kernels/xtensa/xtensa_softmax.h"
#include "tensorflow/lite/micro/micro_log.h"

namespace tflite {
namespace {

#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
TfLiteStatus EvalHifiInt8(const XtensaSoftmaxOpData* op_data,
const TfLiteEvalTensor* input,
TfLiteEvalTensor* output, TfLiteContext* context) {
Expand All @@ -47,16 +46,46 @@ TfLiteStatus EvalHifiInt8(const XtensaSoftmaxOpData* op_data,

void* p_scratch = static_cast<void*>(
context->GetScratchBuffer(context, op_data->scratch_tensor_index));
int lut_flag = (outer_size * depth) > 256? 1 : 0 ;
if(lut_flag)
{
int err = xa_nn_vec_batch_softmax_lut_asym8s_asym8s(
output_data, input_data, op_data->softmax_lut,
depth, outer_size, p_scratch);
TF_LITE_ENSURE(context, err == 0);
}
else
{
for (int i = 0; i < outer_size; ++i) {
int err = xa_nn_vec_softmax_asym8s_asym8s(
&output_data[i * depth], &input_data[i * depth],
op_data->params.diff_min, op_data->params.input_left_shift,
op_data->params.input_multiplier, depth, p_scratch);
TF_LITE_ENSURE(context, err == 0);
}
}
return kTfLiteOk;
}
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)

#if defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
TfLiteStatus EvalHifiInt16(const XtensaSoftmaxOpData* op_data,
const TfLiteEvalTensor* input,
TfLiteEvalTensor* output, TfLiteContext* context) {
const RuntimeShape& input_shape = tflite::micro::GetTensorShape(input);
const int16_t* input_data = tflite::micro::GetTensorData<int16_t>(input);
const RuntimeShape& output_shape = tflite::micro::GetTensorShape(output);
int16_t* output_data = tflite::micro::GetTensorData<int16_t>(output);
const int trailing_dim = input_shape.DimensionsCount() - 1;
const int outer_size =
MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
const int depth =
MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);
int err = xa_nn_vec_batch_softmax_sym16s_16(output_data, input_data, op_data->params.input_left_shift, op_data->params.input_multiplier, depth, outer_size);
TF_LITE_ENSURE(context, err == 0);
return kTfLiteOk;
}
#endif // defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, 0);
Expand All @@ -68,7 +97,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {

TFLITE_DCHECK(node->user_data != nullptr);

#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
XtensaSoftmaxOpData op_data =
*static_cast<XtensaSoftmaxOpData*>(node->user_data);
SoftmaxParams params = op_data.params;
Expand All @@ -77,7 +106,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
#endif

if (input->type == kTfLiteInt8 && output->type == kTfLiteInt8) {
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
return EvalHifiInt8(static_cast<XtensaSoftmaxOpData*>(node->user_data),
input, output, context);
#elif defined(VISION_P6)
Expand All @@ -91,16 +120,21 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
tflite::micro::GetTensorShape(output),
tflite::micro::GetTensorData<int8_t>(output));
return kTfLiteOk;
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
}

if (input->type == kTfLiteInt16 && output->type == kTfLiteInt16) {
#if defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
return EvalHifiInt16(static_cast<XtensaSoftmaxOpData*>(node->user_data),
input, output, context);
#else
tflite::reference_ops::SoftmaxInt16(
params, tflite::micro::GetTensorShape(input),
tflite::micro::GetTensorData<int16_t>(input),
tflite::micro::GetTensorShape(output),
tflite::micro::GetTensorData<int16_t>(output));
return kTfLiteOk;
#endif // defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
}

if (input->type == kTfLiteFloat32) {
Expand All @@ -116,11 +150,112 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
return kTfLiteError;
}

TfLiteStatus EvalInt8(TfLiteContext* context, TfLiteNode* node) {
const TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, 0);
TfLiteEvalTensor* output = tflite::micro::GetEvalOutput(context, node, 0);

TFLITE_DCHECK(node->user_data != nullptr);

#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) || defined(VISION_P6))
SoftmaxParams params = *static_cast<SoftmaxParams*>(node->user_data);
#endif

if (input->type == kTfLiteInt8 && output->type == kTfLiteInt8) {
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
return EvalHifiInt8(static_cast<XtensaSoftmaxOpData*>(node->user_data),
input, output, context);
#elif defined(VISION_P6)
return SoftmaxEvalVision(
context, node, *(static_cast<XtensaSoftmaxOpData*>(node->user_data)),
input, output);
#else
tflite::reference_ops::Softmax(
params, tflite::micro::GetTensorShape(input),
tflite::micro::GetTensorData<int8_t>(input),
tflite::micro::GetTensorShape(output),
tflite::micro::GetTensorData<int8_t>(output));
return kTfLiteOk;
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
}

MicroPrintf("Type %s (%d) not supported.", TfLiteTypeGetName(input->type),
input->type);
return kTfLiteError;
}

TfLiteStatus EvalInt16(TfLiteContext* context, TfLiteNode* node) {
const TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, 0);
TfLiteEvalTensor* output = tflite::micro::GetEvalOutput(context, node, 0);

TFLITE_DCHECK(node->user_data != nullptr);

#if !(defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ) || defined(VISION_P6))
SoftmaxParams params = *static_cast<SoftmaxParams*>(node->user_data);
#endif

if (input->type == kTfLiteInt16 && output->type == kTfLiteInt16) {
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
return EvalHifiInt16(static_cast<XtensaSoftmaxOpData*>(node->user_data),
input, output, context);
#else
tflite::reference_ops::SoftmaxInt16(
params, tflite::micro::GetTensorShape(input),
tflite::micro::GetTensorData<int16_t>(input),
tflite::micro::GetTensorShape(output),
tflite::micro::GetTensorData<int16_t>(output));
return kTfLiteOk;
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
}

MicroPrintf("Type %s (%d) not supported.", TfLiteTypeGetName(input->type),
input->type);
return kTfLiteError;
}

TfLiteStatus EvalFloat32(TfLiteContext* context, TfLiteNode* node) {
const TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, 0);
TfLiteEvalTensor* output = tflite::micro::GetEvalOutput(context, node, 0);

TFLITE_DCHECK(node->user_data != nullptr);

#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
XtensaSoftmaxOpData op_data =
*static_cast<XtensaSoftmaxOpData*>(node->user_data);
SoftmaxParams params = op_data.params;
#else
SoftmaxParams params = *static_cast<SoftmaxParams*>(node->user_data);
#endif

if (input->type == kTfLiteFloat32) {
tflite::reference_ops::Softmax(params, tflite::micro::GetTensorShape(input),
tflite::micro::GetTensorData<float>(input),
tflite::micro::GetTensorShape(output),
tflite::micro::GetTensorData<float>(output));
return kTfLiteOk;
}

MicroPrintf("Type %s (%d) not supported.", TfLiteTypeGetName(input->type),
input->type);
return kTfLiteError;
}
} // namespace

TFLMRegistration Register_SOFTMAX() {
return tflite::micro::RegisterOp(XtensaInitSoftmax, XtensaPrepareSoftmax,
Eval);
}

TFLMRegistration Register_SOFTMAX_INT8() {
return tflite::micro::RegisterOp(XtensaInitSoftmax, XtensaPrepareSoftmax,
EvalInt8);
}

TFLMRegistration Register_SOFTMAX_INT16() {
return tflite::micro::RegisterOp(XtensaInitSoftmax, XtensaPrepareSoftmax,
EvalInt16);
}
TFLMRegistration Register_SOFTMAX_FLOAT32() {
return tflite::micro::RegisterOp(XtensaInitSoftmax, XtensaPrepareSoftmax,
EvalFloat32);
}
} // namespace tflite
37 changes: 28 additions & 9 deletions tensorflow/lite/micro/kernels/xtensa/softmax_int8_int16.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ limitations under the License.
namespace tflite {
namespace {

#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
TfLiteStatus PrepareHifi(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_OK(context, SoftmaxPrepare(context, node));

Expand All @@ -46,14 +46,33 @@ TfLiteStatus PrepareHifi(TfLiteContext* context, TfLiteNode* node) {
MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);

if (input->type == kTfLiteInt8) {
int required_scratch =
get_softmax_scratch_size(PREC_ASYM8S, PREC_ASYM8S, depth);
int required_scratch;
const int outer_size =
MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
int lut_flag = (outer_size * depth) > 256? 1 : 0 ;
int depth_aligned = (depth+31)& ~31;
if(lut_flag)
{
required_scratch = 4 * 4 * depth_aligned;
}
else{
required_scratch = 4 * depth_aligned;
}
TF_LITE_ENSURE(context, required_scratch > 0);

auto* data = static_cast<XtensaSoftmaxOpData*>(node->user_data);
TF_LITE_ENSURE_OK(
context, context->RequestScratchBufferInArena(
context, required_scratch, &(data->scratch_tensor_index)));
if(lut_flag)
{
void* raw_lut = context->AllocatePersistentBuffer(
context, sizeof(int32_t) * 256);
TF_LITE_ENSURE(context, raw_lut != nullptr);
data->softmax_lut = reinterpret_cast<int32_t*>(raw_lut);
xa_nn_init_lut_asym8s_softmax(data->softmax_lut, data->params.diff_min,
data->params.input_multiplier,data->params.input_left_shift);
}
}

micro_context->DeallocateTempTfLiteTensor(input);
Expand Down Expand Up @@ -86,13 +105,13 @@ TfLiteStatus EvalHifi(const XtensaSoftmaxOpData* op_data,
}
return kTfLiteOk;
}
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)

} // namespace

void* XtensaInitSoftmax(TfLiteContext* context, const char* buffer,
size_t length) {
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
return context->AllocatePersistentBuffer(context,
sizeof(XtensaSoftmaxOpData));
Expand All @@ -105,11 +124,11 @@ void* XtensaInitSoftmax(TfLiteContext* context, const char* buffer,
sizeof(XtensaSoftmaxOpData));
#else
return SoftmaxInit(context, buffer, length);
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
}

TfLiteStatus XtensaPrepareSoftmax(TfLiteContext* context, TfLiteNode* node) {
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
return PrepareHifi(context, node);
#else
TF_LITE_ENSURE_OK(context, SoftmaxPrepare(context, node));
Expand All @@ -127,7 +146,7 @@ TfLiteStatus XtensaEvalSoftmaxInt8Int16(TfLiteContext* context,
TFLITE_DCHECK(node->user_data != nullptr);

if (input->type == kTfLiteInt8 && output->type == kTfLiteInt16) {
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
return EvalHifi(static_cast<XtensaSoftmaxOpData*>(node->user_data), input,
output, context);
#else
Expand All @@ -138,7 +157,7 @@ TfLiteStatus XtensaEvalSoftmaxInt8Int16(TfLiteContext* context,
tflite::micro::GetTensorShape(output),
tflite::micro::GetTensorData<int16_t>(output));
return kTfLiteOk;
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
} else {
MicroPrintf("Type %s (%d) not supported.", TfLiteTypeGetName(input->type),
input->type);
Expand Down
2 changes: 1 addition & 1 deletion tensorflow/lite/micro/kernels/xtensa/xtensa.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ limitations under the License.
#include "tensorflow/lite/micro/kernels/xtensa/fixedpoint_utils_hifimini.h"
#endif // defined(HIFMINI)

#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
#include "include/nnlib/xa_nnlib_api.h"
#include "include/nnlib/xa_nnlib_standards.h"

Expand Down
5 changes: 3 additions & 2 deletions tensorflow/lite/micro/kernels/xtensa/xtensa_softmax.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,13 @@ limitations under the License.

namespace tflite {

#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#if defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)
struct XtensaSoftmaxOpData {
SoftmaxParams params;
int scratch_tensor_index;
int32_t* softmax_lut;
};
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5)
#endif // defined(HIFI3) || defined(HIFI4) || defined(HIFI5) || defined(HIFI_IQ)

#if defined(VISION_P6)
struct XtensaSoftmaxOpData {
Expand Down
Loading