Skip to content

Commit 0830c99

Browse files
veblushmansnils
veblush
authored and
mansnils
committed
Fixed LSTM
1 parent b26042b commit 0830c99

3 files changed

Lines changed: 93 additions & 24 deletions

File tree

‎Include/arm_nn_types.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,7 @@ typedef struct
272272
void *temp1;
273273
void *temp2;
274274
void *cell_state;
275+
void *hidden_state;
275276
} cmsis_nn_lstm_context;
276277

277278
/**

‎Source/LSTMFunctions/arm_lstm_unidirectional_s16.c‎

Lines changed: 46 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,13 @@ arm_cmsis_nn_status arm_lstm_unidirectional_s16(const int16_t *input,
5252
cmsis_nn_lstm_context *buffers)
5353
{
5454

55-
int16_t *hidden_in = NULL;
56-
memset(buffers->cell_state, 0, params->batch_size * params->hidden_size * sizeof(int16_t));
55+
int16_t *hidden_in = (int16_t *)buffers->hidden_state;
56+
57+
if (buffers->hidden_state == NULL)
58+
{
59+
memset(buffers->cell_state, 0, params->batch_size * params->hidden_size * sizeof(int16_t));
60+
}
61+
5762
if (params->time_major)
5863
{
5964
// First dimension is time, input/output for each time step is stored continously in memory
@@ -69,22 +74,51 @@ arm_cmsis_nn_status arm_lstm_unidirectional_s16(const int16_t *input,
6974
// Output is used as recurrent input/hidden state for the next timestep.
7075
hidden_in = &hidden_out[0];
7176
}
77+
78+
if (buffers->hidden_state != NULL && params->time_steps > 0)
79+
{
80+
memcpy(buffers->hidden_state, hidden_in, params->batch_size * params->hidden_size * sizeof(int16_t));
81+
}
7282
}
7383
else
7484
{
75-
// First dimension is time, add batch_offset to jump in memory for each batch
76-
for (int t = 0; t < params->time_steps; t++)
85+
// Batch major: [batch, time, size]
86+
// arm_nn_lstm_step_s16 expects data_in and hidden_in to have the same batch_offset.
87+
// Since the initial hidden_state is contiguous, we must process one batch at a time.
88+
cmsis_nn_lstm_params step_params = *params;
89+
step_params.batch_size = 1;
90+
91+
for (int b = 0; b < params->batch_size; b++)
7792
{
78-
const int16_t *data_in = input + (t * params->input_size);
79-
int16_t *hidden_out = output + (t * params->hidden_size);
80-
arm_cmsis_nn_status status =
81-
arm_nn_lstm_step_s16(data_in, hidden_in, hidden_out, params, buffers, params->time_steps);
82-
if (status != ARM_CMSIS_NN_SUCCESS)
93+
int16_t *step_hidden_in = (buffers->hidden_state != NULL)
94+
? ((int16_t *)buffers->hidden_state + b * params->hidden_size)
95+
: NULL;
96+
97+
cmsis_nn_lstm_context step_buffers = *buffers;
98+
step_buffers.cell_state = (int16_t *)buffers->cell_state + b * params->hidden_size;
99+
100+
for (int t = 0; t < params->time_steps; t++)
83101
{
84-
return status;
102+
const int16_t *data_in = input + (b * params->time_steps + t) * params->input_size;
103+
int16_t *hidden_out = output + (b * params->time_steps + t) * params->hidden_size;
104+
105+
arm_cmsis_nn_status status = arm_nn_lstm_step_s16(
106+
data_in, step_hidden_in, hidden_out, &step_params, &step_buffers, 1);
107+
108+
if (status != ARM_CMSIS_NN_SUCCESS)
109+
{
110+
return status;
111+
}
112+
113+
step_hidden_in = hidden_out;
114+
}
115+
116+
if (buffers->hidden_state != NULL && params->time_steps > 0)
117+
{
118+
memcpy((int16_t *)buffers->hidden_state + b * params->hidden_size,
119+
step_hidden_in,
120+
params->hidden_size * sizeof(int16_t));
85121
}
86-
// Output is used as recurrent input/hidden state for the next timestep.
87-
hidden_in = &hidden_out[0];
88122
}
89123
}
90124
return ARM_CMSIS_NN_SUCCESS;

‎Source/LSTMFunctions/arm_lstm_unidirectional_s8.c‎

Lines changed: 46 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,13 @@ arm_cmsis_nn_status arm_lstm_unidirectional_s8(const int8_t *input,
5252
cmsis_nn_lstm_context *buffers)
5353
{
5454

55-
int8_t *hidden_in = NULL;
56-
memset(buffers->cell_state, 0, params->batch_size * params->hidden_size * sizeof(int16_t));
55+
int8_t *hidden_in = (int8_t *)buffers->hidden_state;
56+
57+
if (buffers->hidden_state == NULL)
58+
{
59+
memset(buffers->cell_state, 0, params->batch_size * params->hidden_size * sizeof(int16_t));
60+
}
61+
5762
if (params->time_major)
5863
{
5964
// First dimension is time, input/output for each time step is stored continously in memory
@@ -69,22 +74,51 @@ arm_cmsis_nn_status arm_lstm_unidirectional_s8(const int8_t *input,
6974
// Output is used as recurrent input/hidden state for the next timestep.
7075
hidden_in = &hidden_out[0];
7176
}
77+
78+
if (buffers->hidden_state != NULL && params->time_steps > 0)
79+
{
80+
memcpy(buffers->hidden_state, hidden_in, params->batch_size * params->hidden_size * sizeof(int8_t));
81+
}
7282
}
7383
else
7484
{
75-
// First dimension is time, add batch_offset to jump in memory for each batch
76-
for (int t = 0; t < params->time_steps; t++)
85+
// Batch major: [batch, time, size]
86+
// arm_nn_lstm_step_s8 expects data_in and hidden_in to have the same batch_offset.
87+
// Since the initial hidden_state is contiguous, we must process one batch at a time.
88+
cmsis_nn_lstm_params step_params = *params;
89+
step_params.batch_size = 1;
90+
91+
for (int b = 0; b < params->batch_size; b++)
7792
{
78-
const int8_t *data_in = input + (t * params->input_size);
79-
int8_t *hidden_out = output + (t * params->hidden_size);
80-
arm_cmsis_nn_status status =
81-
arm_nn_lstm_step_s8(data_in, hidden_in, hidden_out, params, buffers, params->time_steps);
82-
if (status != ARM_CMSIS_NN_SUCCESS)
93+
int8_t *step_hidden_in = (buffers->hidden_state != NULL)
94+
? ((int8_t *)buffers->hidden_state + b * params->hidden_size)
95+
: NULL;
96+
97+
cmsis_nn_lstm_context step_buffers = *buffers;
98+
step_buffers.cell_state = (int16_t *)buffers->cell_state + b * params->hidden_size;
99+
100+
for (int t = 0; t < params->time_steps; t++)
83101
{
84-
return status;
102+
const int8_t *data_in = input + (b * params->time_steps + t) * params->input_size;
103+
int8_t *hidden_out = output + (b * params->time_steps + t) * params->hidden_size;
104+
105+
arm_cmsis_nn_status status = arm_nn_lstm_step_s8(
106+
data_in, step_hidden_in, hidden_out, &step_params, &step_buffers, 1);
107+
108+
if (status != ARM_CMSIS_NN_SUCCESS)
109+
{
110+
return status;
111+
}
112+
113+
step_hidden_in = hidden_out;
114+
}
115+
116+
if (buffers->hidden_state != NULL && params->time_steps > 0)
117+
{
118+
memcpy((int8_t *)buffers->hidden_state + b * params->hidden_size,
119+
step_hidden_in,
120+
params->hidden_size * sizeof(int8_t));
85121
}
86-
// Output is used as recurrent input/hidden state for the next timestep.
87-
hidden_in = &hidden_out[0];
88122
}
89123
}
90124
return ARM_CMSIS_NN_SUCCESS;

0 commit comments

Comments
 (0)