If the input tensor has ref count as 1, it is forwarded to the desired output port and the function returns true. In that case, it also allocates the serialized MklDnnShape object. Otherwise, the function returns false.
| 902 | // output port and the function returns true. In that case, it also allocates |
| 903 | // the serialized MklDnnShape object. Otherwise, the function returns false. |
| 904 | inline bool ForwardMklTensorInToOutWithMklShape(OpKernelContext* context, |
| 905 | int idx_in, int idx_out, |
| 906 | Tensor** output, |
| 907 | const MklDnnShape& mkl_shape, |
| 908 | bool always_forward = true) { |
| 909 | int num_inputs = context->num_inputs(); |
| 910 | int num_outputs = context->num_outputs(); |
| 911 | int idx_data_in = GetTensorDataIndex(idx_in, num_inputs); |
| 912 | int idx_data_out = GetTensorDataIndex(idx_out, num_outputs); |
| 913 | bool is_forwarded = false; |
| 914 | const Tensor& input_tensor = context->input(idx_data_in); |
| 915 | const auto output_shape = input_tensor.shape(); |
| 916 | if (always_forward) { |
| 917 | if (IsRefType(context->input_dtype(idx_data_in))) { |
| 918 | context->forward_ref_input_to_ref_output(idx_data_in, idx_data_out); |
| 919 | } else { |
| 920 | context->set_output(idx_data_out, input_tensor); |
| 921 | } |
| 922 | } else { |
| 923 | is_forwarded = context->forward_input_to_output_with_shape( |
| 924 | idx_data_in, idx_data_out, output_shape, output); |
| 925 | } |
| 926 | if (is_forwarded || always_forward) { |
| 927 | AllocateOutputSetMklShape(context, idx_out, mkl_shape); |
| 928 | return true; |
| 929 | } |
| 930 | return false; |
| 931 | } |
| 932 | |
| 933 | // Forward the OneDNN shape ONLY (used in elementwise and other ops where |
| 934 | // we call the eigen implementation and OneDNN shape is not used) |
no test coverage detected