PyTorch Ops to oneDNN Functions Mapping

PyTorch uses ops that are registered to corresponding Math Kernel Library (MKL) functions in oneDNN. The available implementations are defined in this YAML file native_functions.yaml in the aten library of PyTorch. By doing a search for the keyword “mkldnn”, all the mappings can be found.

This is summarized in the following table:

PyTorch	Op oneDNN Function
`add.Tensor`	`mkldnn_add`
`add_.Tensor`	`mkldnn_add_`
`add.out`	`mkldnn_add_out`
`copy_`	`copy_mkldnn_`
`empty.memory_format`	`empty_mkldnn`
`mkldnn_linear`	`mkldnn_linear`
`mkldnn_linear_backward_input`	`mkldnn_linear_backward_input`
`mkldnn_linear_backward_weights`	`mkldnn_linear_backward_weights`
`mkldnn_linear_backward`	`mkldnn_linear_backward`
`mkldnn_max_pool2d`	`mkldnn_max_pool2d`
`mkldnn_max_pool2d_backward`	`mkldnn_max_pool2d_backward`
`mkldnn_max_pool3d`	`mkldnn_max_pool3d`
`mkldnn_max_pool3d_backward`	`mkldnn_max_pool3d_backward`
`mkldnn_convolution`	`mkldnn_convolution`
`mkldnn_rnn_layer`	`mkldnn_rnn_layer`
`mkldnn_rnn_layer_backward`	`mkldnn_rnn_layer_backward`
`mul.Tensor`	`mkldnn_mul`
`mul_.Tensor`	`mkldnn_mul_`
`mul.out`	`mkldnn_mul_out`
`native_batch_norm`	`mkldnn_batch_norm`
`_native_batch_norm_legit`	`_mkldnn_batch_norm_legit`
`_native_batch_norm_legit.no_stats`	`_mkldnn_batch_norm_legit_no_stats`
`native_batch_norm_backward`	`mkldnn_batch_norm_backward`
`_mkldnn_reshape`	`mkldnn_reshape`
`relu`	`mkldnn_relu`
`relu_`	`mkldnn_relu_`
`_prelu_kernel`	`mkldnn_prelu_backward`
`gelu`	`mkldnn_gelu`
`gelu_backward`	`mkldnn_gelu_backward`
`sigmoid`	`mkldnn_sigmoid`
`sigmoid_`	`mkldnn_sigmoid_`
`_softmax`	`mkldnn_softmax`
`tanh`	`mkldnn_tanh`
`tanh_`	`mkldnn_tanh_`
`threshold_backward`	`mkldnn_relu_backward`
`_mkldnn_transpose`	`mkldnn_transpose`
`_mkldnn_transpose_`	`mkldnn_transpose_`
`clone`	`mkldnn_clone`
`zero_`	`mkldnn_zero_`
`_to_dense`	`mkldnn_to_dense`
`to_mkldnn`	`dense_to_mkldnn`
`mkldnn_reorder_conv2d_weight`	`mkldnn_reorder_conv2d_weight`
`mkldnn_reorder_conv3d_weight`	`mkldnn_reorder_conv3d_weight`
`view`	`mkldnn_view`
`adaptive_avg_pool2d.out`	`mkldnn_adaptive_avg_pool2d_out_stub`
`mkldnn_adaptive_avg_pool2d`	`mkldnn_adaptive_avg_pool2d`
`mkldnn_adaptive_avg_pool2d.out`	`mkldnn_adaptive_avg_pool2d_out`
`mkldnn_adaptive_avg_pool2d_backward`	`mkldnn_adaptive_avg_pool2d_backward`
`avg_pool2d.out`	`mkldnn_avg_pool2d_out`
`avg_pool2d`	`mkldnn_avg_pool2d`
`avg_pool2d_backward.grad_input`	`mkldnn_avg_pool2d_backward_out`
`avg_pool2d_backward`	`mkldnn_avg_pool2d_backward`
`avg_pool3d.out`	`mkldnn_avg_pool3d_out`
`avg_pool3d`	`mkldnn_avg_pool3d`
`avg_pool3d_backward.grad_input`	`mkldnn_avg_pool3d_backward_out`
`avg_pool3d_backward`	`mkldnn_avg_pool3d_backward`

I would love to contribute to pytorch!

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

PyTorch Ops to oneDNN Functions Mapping

Clone this wiki locally