Skip to content

Conversation

@pramods-cad
Copy link
Contributor

Enabling batching in LSTM on HiFi for cases where batch > 1 and time iteration = 1 irrespective of time_major flag to improve optimization.
bug=fixes #3243

@pramods-cad pramods-cad requested a review from a team as a code owner October 27, 2025 15:30
Comment on lines +669 to +673
const auto& size_info = op_data.size_info;
if(size_info.batch_size > 1 && size_info.time_steps == 1) {
num_batches = size_info.batch_size;
}

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe move this up to the definition of num_batches and do the logic there? I think all those values should be const int as well.

Comment on lines 811 to +813
} else {
// batch first, unable to size the input data. single batch inference
for (int b = 0; b < size_info.batch_size; b++) {
for (int t = 0; t < size_info.time_steps; t++) {
lstm_internal::LstmStep<ActivationType, WeightType, CellType, BiasType>(
step_info, op_data, kernel_content, buffers);
// prepare for the next time step
step_info.UpdateTime();
if(size_info.batch_size > 1 && size_info.time_steps == 1) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

} else if (...) {?

@gonnet gonnet added the ci:run label Oct 27, 2025
@TFLM-bot TFLM-bot removed the ci:run label Oct 27, 2025
@veblush veblush added the ci:run label Oct 29, 2025
@TFLM-bot TFLM-bot removed the ci:run label Oct 29, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Optimization improvement in LSTM on HiFi for batch > 1 cases.

4 participants