Softmax forward execute src_data: input data buffer of src dst_data: output data buffer of dst
| 63 | // src_data: input data buffer of src |
| 64 | // dst_data: output data buffer of dst |
| 65 | void Execute(const T* src_data, T* dst_data, |
| 66 | std::shared_ptr<stream> fwd_cpu_stream) { |
| 67 | #ifdef DNNL_AARCH64_USE_ACL |
| 68 | mutex_lock lock(primitive_execution_mu_); |
| 69 | #endif |
| 70 | #ifdef ENABLE_DNNL_THREADPOOL |
| 71 | context_.src_mem->set_data_handle( |
| 72 | static_cast<void*>(const_cast<T*>(src_data)), *fwd_cpu_stream); |
| 73 | context_.dst_mem->set_data_handle(static_cast<void*>(dst_data), |
| 74 | *fwd_cpu_stream); |
| 75 | #else |
| 76 | context_.src_mem->set_data_handle( |
| 77 | static_cast<void*>(const_cast<T*>(src_data))); |
| 78 | context_.dst_mem->set_data_handle(static_cast<void*>(dst_data)); |
| 79 | #endif // ENABLE_DNNL_THREADPOOL |
| 80 | DCHECK_EQ(context_.fwd_primitives.size(), context_.fwd_net_args.size()); |
| 81 | execute_primitives(context_.fwd_primitives, fwd_cpu_stream, |
| 82 | context_.fwd_net_args); |
| 83 | |
| 84 | // After execution, set data handle back. |
| 85 | context_.src_mem->set_data_handle(DummyData); |
| 86 | context_.dst_mem->set_data_handle(DummyData); |
| 87 | } |
| 88 | |
| 89 | std::shared_ptr<dnnl::softmax_forward::primitive_desc> GetSoftmaxFwdPd() { |
| 90 | return context_.fwd_pd; |
no test coverage detected