| 543 | |
| 544 | template <typename TensorDataType> |
| 545 | void fp_compute_impl( |
| 546 | gru_layer<TensorDataType, data_layout::DATA_PARALLEL, El::Device::CPU>& l) |
| 547 | { |
| 548 | |
| 549 | // Matrices |
| 550 | using LocalMat = El::Matrix<TensorDataType, El::Device::CPU>; |
| 551 | const auto& input_sequence = |
| 552 | dynamic_cast<const LocalMat&>(l.get_local_prev_activations(0)); |
| 553 | const auto& init_hidden = |
| 554 | dynamic_cast<const LocalMat&>(l.get_local_prev_activations(1)); |
| 555 | auto& output_sequence = dynamic_cast<LocalMat&>(l.get_local_activations()); |
| 556 | |
| 557 | // Dimensions |
| 558 | const int local_mini_batch_size = input_sequence.Width(); |
| 559 | const int sequence_length = l.get_input_dims(0)[0]; |
| 560 | const int input_size = l.get_input_size(0) / sequence_length; |
| 561 | const int hidden_size = l.m_hidden_size; |
| 562 | const int num_layers = l.m_num_layers; |
| 563 | |
| 564 | // Return immediately if there is no local data. |
| 565 | if (local_mini_batch_size <= 0) { |
| 566 | return; |
| 567 | } |
| 568 | |
| 569 | // oneDNN objects |
| 570 | if (l.m_onednn_cpu_objects == nullptr) { |
| 571 | LBANN_ERROR(l.get_type(), |
| 572 | " layer \"", |
| 573 | l.get_name(), |
| 574 | "\" ", |
| 575 | "attempted to run oneDNN CPU implementation ", |
| 576 | "before initializing oneDNN objects"); |
| 577 | } |
| 578 | constexpr auto Device = El::Device::CPU; |
| 579 | using Backend = onednn_backend<Device>; |
| 580 | using Memory = ::dnnl::memory; |
| 581 | const auto data_type = Backend::template data_type<TensorDataType>(); |
| 582 | auto& onednn_objects = *l.m_onednn_cpu_objects; |
| 583 | auto sync_info = force(El::MakeMultiSync(get_sync_info(output_sequence), |
| 584 | get_sync_info(input_sequence), |
| 585 | get_sync_info(init_hidden))); |
| 586 | auto& engine = onednn::get_device_engine<Device>(); |
| 587 | auto stream = onednn::get_stream<Device>(engine, sync_info); |
| 588 | |
| 589 | // Configure input and output tensor descriptors |
| 590 | onednn_objects.input_sequence_desc.set( |
| 591 | data_type, |
| 592 | {sequence_length, local_mini_batch_size, input_size}, |
| 593 | {input_size, El::To<int>(input_sequence.LDim()), 1}); |
| 594 | onednn_objects.input_sequence_desc.get().set_data_handle( |
| 595 | const_cast<TensorDataType*>(input_sequence.LockedBuffer()), |
| 596 | stream); |
| 597 | onednn_objects.init_hidden_desc.set( |
| 598 | data_type, |
| 599 | {num_layers, /*num_directions=*/1, local_mini_batch_size, hidden_size}, |
| 600 | {hidden_size, 1, El::To<int>(init_hidden.LDim()), 1}); |
| 601 | onednn_objects.init_hidden_desc.get().set_data_handle( |
| 602 | const_cast<TensorDataType*>(init_hidden.LockedBuffer()), |
no test coverage detected