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:
Nick Kreeger 2021-06-28 08:36:02 -05:00 committed by GitHub
parent 523db6ef44
commit 821492f6f5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 10 deletions

View file

@ -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));

View file

@ -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));