| 459 | virtual void Compute_Scalar(OpKernelContext* context) = 0; |
| 460 | |
| 461 | void Compute(OpKernelContext* context) override { |
| 462 | try { |
| 463 | const size_t src_index = 0; // index of src input tensor |
| 464 | const size_t dst_index = 0; // index of dst output tensor |
| 465 | const Tensor& src_tensor = MklGetInput(context, src_index); |
| 466 | MklDnnShape dnn_shape_src; |
| 467 | GetMklShape(context, src_index, &dnn_shape_src); |
| 468 | if (src_tensor.dims() == 0) { |
| 469 | Compute_Scalar(context); |
| 470 | return; |
| 471 | } |
| 472 | MklDnnShape dnn_shape_dst; |
| 473 | TensorShape tf_shape_dst; |
| 474 | Tensor* dst_tensor = nullptr; |
| 475 | // Nothing to compute, return. |
| 476 | if (src_tensor.shape().num_elements() == 0) { |
| 477 | dnn_shape_dst.SetMklTensor(false); |
| 478 | tf_shape_dst = MklGetInput(context, src_index).shape(); |
| 479 | AllocateOutputSetMklShape(context, dst_index, &dst_tensor, tf_shape_dst, |
| 480 | dnn_shape_dst); |
| 481 | return; |
| 482 | } |
| 483 | // Set DNN primitive - src |
| 484 | MklDnnData<T> src(&cpu_engine); |
| 485 | memory::dims src_dims; |
| 486 | memory::desc src_md({}, MEMORY_DATA_TYPE_UNDEF, MEMORY_FORMAT_UNDEF); |
| 487 | if (dnn_shape_src.IsMklTensor()) { |
| 488 | src_md = dnn_shape_src.GetMklLayout(); |
| 489 | src_dims = dnn_shape_src.GetSizesAsMklDnnDims(); |
| 490 | } else { |
| 491 | src_dims = TFShapeToMklDnnDims(src_tensor.shape()); |
| 492 | auto src_strides = CalculateTFStrides(src_dims); |
| 493 | // Create blocked memory descriptor |
| 494 | src_md = MklDnnData<T>::CreateBlockedMemDesc(src_dims, src_strides); |
| 495 | } |
| 496 | // Try to get an eltwise forward primitive from caching pool |
| 497 | MklEltwiseFwdParams<T> fwdParams(src_dims, src_md, alg_kind_, alpha_, |
| 498 | beta_); |
| 499 | MklDnnThreadPool eigen_tp(context); |
| 500 | MklEltwiseFwdPrimitive<T>* eltwise_fwd = |
| 501 | MklEltwiseFwdPrimitiveFactory<T>::Get(fwdParams); |
| 502 | auto eltwise_fwd_pd = eltwise_fwd->GetEltwiseFwdPd(); |
| 503 | std::shared_ptr<stream> fwd_cpu_stream; |
| 504 | fwd_cpu_stream.reset(CreateStream(&eigen_tp, eltwise_fwd->GetEngine())); |
| 505 | // Check if src needs to be reordered |
| 506 | const T* src_data = src_tensor.flat<T>().data(); |
| 507 | if (IS_SRC_REORDER_NEEDED(src_md, eltwise_fwd_pd, eltwise_fwd)) { |
| 508 | src.SetUsrMem(src_md, &src_tensor); |
| 509 | src.CheckReorderToOpMem( |
| 510 | MEMORY_PD_WITHOUT_DATA(eltwise_fwd_pd->PRIMITIVE_DESC_SRC, |
| 511 | cpu_engine), |
| 512 | context); |
| 513 | src_data = const_cast<T*>( |
| 514 | reinterpret_cast<T*>(src.GetOpMem().get_data_handle())); |
| 515 | } |
| 516 | // Allocate dst tensor, always set it as OneDNN layout |
| 517 | if (dnn_shape_src.IsMklTensor()) { |
| 518 | dnn_shape_dst.SetMklTensor(true); |
nothing calls this directly
no test coverage detected