mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Drop std::count_if() in *EmbedLayerNorm Ops. (#8161)
* Drop std::count_if() in *EmbedLayerNorm Ops. Profiling has shown that summing up the vector using the std function can be 2x slower than just a simple plain vector sum loop. * try and revert sumodule commits * ensure mask is 1.
This commit is contained in:
parent
523db6ef44
commit
821492f6f5
2 changed files with 16 additions and 10 deletions
|
|
@ -133,14 +133,17 @@ Status EmbedLayerNorm<T>::Compute(OpKernelContext* context) const {
|
|||
|
||||
// Calculate mask
|
||||
if (nullptr != mask) {
|
||||
// TODO: Consider summing the values in the mask and measure performance.
|
||||
const int32_t* mask_data = mask->template Data<int32_t>();
|
||||
int32_t* mask_index_data = mask_index->template MutableData<int32_t>();
|
||||
for (int b = 0; b < batch_size; b++) {
|
||||
mask_index_data[b] =
|
||||
static_cast<int32_t>(std::count_if(mask_data + (static_cast<int64_t>(b) * sequence_length),
|
||||
mask_data + (static_cast<int64_t>(b) * sequence_length) + sequence_length,
|
||||
[](int v) { return v == 1; }));
|
||||
int32_t cur_sum = 0;
|
||||
const int32_t* cur_mask_data = mask_data + (static_cast<int64_t>(b) * sequence_length);
|
||||
for (int s = 0; s < sequence_length; ++s) {
|
||||
if (cur_mask_data[s] == 1) {
|
||||
cur_sum += cur_mask_data[s];
|
||||
}
|
||||
}
|
||||
mask_index_data[b] = cur_sum;
|
||||
}
|
||||
} else {
|
||||
memset(mask_index->template MutableData<int32_t>(), 0, batch_size * sizeof(int32_t));
|
||||
|
|
|
|||
|
|
@ -174,14 +174,17 @@ Status ComputeInternal(OpKernelContext* context, float epsilon) {
|
|||
|
||||
// Calculate mask
|
||||
if (nullptr != mask) {
|
||||
// TODO: Consider summing the values in the mask and measure performance.
|
||||
const int32_t* mask_data = mask->template Data<int32_t>();
|
||||
int32_t* mask_index_data = mask_index->template MutableData<int32_t>();
|
||||
for (int b = 0; b < batch_size; b++) {
|
||||
mask_index_data[b] =
|
||||
static_cast<int32_t>(std::count_if(mask_data + (static_cast<int64_t>(b) * sequence_length),
|
||||
mask_data + (static_cast<int64_t>(b) * sequence_length) + sequence_length,
|
||||
[](int v) { return v == 1; }));
|
||||
int32_t cur_sum = 0;
|
||||
const int32_t* cur_mask_data = mask_data + (static_cast<int64_t>(b) * sequence_length);
|
||||
for (int s = 0; s < sequence_length; ++s) {
|
||||
if (cur_mask_data[s] == 1) {
|
||||
cur_sum += cur_mask_data[s];
|
||||
}
|
||||
}
|
||||
mask_index_data[b] = cur_sum;
|
||||
}
|
||||
} else {
|
||||
memset(mask_index->template MutableData<int32_t>(), 0, batch_size * sizeof(int32_t));
|
||||
|
|
|
|||
Loading…
Reference in a new issue