whisper : fix DTW memory access (#2012)

* Fix DTW memory access

* Memory fix - Apply changes from denersc
This commit is contained in:
Brad Murray 2024-04-09 11:38:19 -04:00 committed by GitHub
parent c15b4cda7d
commit 5275074d37
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1148,26 +1148,31 @@ static bool aheads_masks_init(
} }
// Set data on mask tensors // Set data on mask tensors
// Since this must be backend agnostic, we get tensor data with // Since this must be backend agnostic, we write our desired values on mask_data,
// ggml_backend_tensor_get, copy our desired values and send it back // and send it to backend with ggml_backend_tensor_set.
// to backend with ggml_backend_tensor_set // Each mask in N_HEADS*N_ALIGNMENT_HEADS, one per text layer containing alignment
// heads. Each row of the mask "marks" one alignment head. E.g. if some text layer
// has a total of 10 heads and of those, heads 0,5,6 are alignment heads, the mask
// should read:
// 1 0 0 0 0 0 0 0 0 0
// 0 0 0 0 0 1 0 0 0 0
// 0 0 0 0 0 0 1 0 0 0
std::vector<float> mask_data; std::vector<float> mask_data;
for (int64_t il = 0; il < n_text_layer; ++il) { for (int64_t il = 0; il < n_text_layer; ++il) {
if (aheads_masks.m[il] != nullptr) { if (aheads_masks.m[il] != nullptr) {
auto aheads = get_alignment_heads_by_layer(cparams, il, n_text_layer, n_head); auto aheads = get_alignment_heads_by_layer(cparams, il, n_text_layer, n_head);
size_t data_size = aheads_masks.m[il]->ne[0] * aheads_masks.m[il]->ne[1] * sizeof(float); size_t data_size = aheads_masks.m[il]->ne[0] * aheads_masks.m[il]->ne[1];
size_t data_size_bytes = data_size * sizeof(float);
mask_data.resize(data_size); mask_data.resize(data_size);
ggml_backend_tensor_get(aheads_masks.m[il], mask_data.data(), 0, data_size);
memset(mask_data.data(), 0, data_size);
std::fill(mask_data.begin(), mask_data.end(), 0);
for (size_t ih = 0; ih < aheads.size(); ++ih) { for (size_t ih = 0; ih < aheads.size(); ++ih) {
size_t pos = (aheads[ih] + (ih * aheads_masks.m[il]->ne[0] * aheads[ih])); size_t pos = (aheads[ih] + (ih * aheads_masks.m[il]->ne[0]));
float v = 1.0f; mask_data[pos] = 1.0f;
memcpy(mask_data.data() + pos, &v, sizeof(float));
} }
ggml_backend_tensor_set(aheads_masks.m[il], mask_data.data(), 0, data_size); ggml_backend_tensor_set(aheads_masks.m[il], mask_data.data(), 0, data_size_bytes);
} }
} }