.. _cpu-threading-torchscript-inference:
CPU threading and TorchScript inference
=================================================
PyTorch allows using multiple CPU threads during TorchScript model inference.
The following figure shows different levels of parallelism one would find in a
typical application:
.. image:: cpu_threading_torchscript_inference.svg
:width: 75%
One or more inference threads execute a model's forward pass on the given inputs.
Each inference thread invokes a JIT interpreter that executes the ops
of a model inline, one by one. A model can utilize a ``fork`` TorchScript
primitive to launch an asynchronous task. Forking several operations at once
results in a task that is executed in parallel. The ``fork`` operator returns a
``future`` object which can be used to synchronize on later, for example:
.. code-block:: python
@torch.jit.script
def compute_z(x):
return torch.mm(x, self.w_z)
@torch.jit.script
def forward(x):
# launch compute_z asynchronously:
fut = torch.jit._fork(compute_z, x)
# execute the next operation in parallel to compute_z:
y = torch.mm(x, self.w_y)
# wait for the result of compute_z:
z = torch.jit._wait(fut)
return y + z
PyTorch uses a single thread pool for the inter-op parallelism, this thread pool
is shared by all inference tasks that are forked within the application process.
In addition to the inter-op parallelism, PyTorch can also utilize multiple threads
within the ops (`intra-op parallelism`). This can be useful in many cases,
including element-wise ops on large tensors, convolutions, GEMMs, embedding
lookups and others.
Build options
-------------
PyTorch uses an internal ATen library to implement ops. In addition to that,
PyTorch can also be built with support of external libraries, such as MKL_ and MKL-DNN_,
to speed up computations on CPU.
ATen, MKL and MKL-DNN support intra-op parallelism and depend on the
following parallelization libraries to implement it:
* OpenMP_ - a standard (and a library, usually shipped with a compiler), widely used in external libraries;
* TBB_ - a newer parallelization library optimized for task-based parallelism and concurrent environments.
OpenMP historically has been used by a large number of libraries. It is known
for a relative ease of use and support for loop-based parallelism and other primitives.
At the same time OpenMP is not known for a good interoperability with other threading
libraries used by the application. In particular, OpenMP does not guarantee that a single per-process intra-op thread
pool is going to be used in the application. On the contrary, two different inter-op
threads will likely use different OpenMP thread pools for intra-op work.
This might result in a large number of threads used by the application.
TBB is used to a lesser extent in external libraries, but, at the same time,
is optimized for the concurrent environments. PyTorch's TBB backend guarantees that
there's a separate, single, per-process intra-op thread pool used by all of the
ops running in the application.
Depending of the use case, one might find one or another parallelization
library a better choice in their application.
PyTorch allows selecting of the parallelization backend used by ATen and other
libraries at the build time with the following build options:
+------------+-----------------------+-----------------------------+----------------------------------------+
| Library | Build Option | Values | Notes |
+============+=======================+=============================+========================================+
| ATen | ``ATEN_THREADING`` | ``OMP`` (default), ``TBB`` | |
+------------+-----------------------+-----------------------------+----------------------------------------+
| MKL | ``MKL_THREADING`` | (same) | To enable MKL use ``BLAS=MKL`` |
+------------+-----------------------+-----------------------------+----------------------------------------+
| MKL-DNN | ``MKLDNN_THREADING`` | (same) | To enable MKL-DNN use ``USE_MKLDNN=1`` |
+------------+-----------------------+-----------------------------+----------------------------------------+
It is strongly recommended not to mix OpenMP and TBB within one build.
Any of the ``TBB`` values above require ``USE_TBB=1`` build setting (default: OFF).
A separate setting ``USE_OPENMP=1`` (default: ON) is required for OpenMP parallelism.
Runtime API
-----------
The following API is used to control thread settings:
+------------------------+-----------------------------------------------------------+---------------------------------------------------------+
| Type of parallelism | Settings | Notes |
+========================+===========================================================+=========================================================+
| Inter-op parallelism | ``at::set_num_interop_threads``, | ``set*`` functions can only be called once and only |
| | ``at::get_num_interop_threads`` (C++) | during the startup, before the actual operators running;|
| | | |
| | ``set_num_interop_threads``, | Default number of threads: number of CPU cores. |
| | ``get_num_interop_threads`` (Python, :mod:`torch` module) | |
+------------------------+-----------------------------------------------------------+ |
| Intra-op parallelism | ``at::set_num_threads``, | |
| | ``at::get_num_threads`` (C++) | |
| | ``set_num_threads``, | |
| | ``get_num_threads`` (Python, :mod:`torch` module) | |
| | | |
| | Environment variables: | |
| | ``OMP_NUM_THREADS`` and ``MKL_NUM_THREADS`` | |
+------------------------+-----------------------------------------------------------+---------------------------------------------------------+
For the intra-op parallelism settings, ``at::set_num_threads``, ``torch.set_num_threads`` always take precedence
over environment variables, ``MKL_NUM_THREADS`` variable takes precedence over ``OMP_NUM_THREADS``.
.. note::
``parallel_info`` utility prints information about thread settings and can be used for debugging.
Similar output can be also obtained in Python with ``torch.__config__.parallel_info()`` call.
.. _OpenMP: https://www.openmp.org/
.. _TBB: https://github.com/intel/tbb
.. _MKL: https://software.intel.com/en-us/mkl
.. _MKL-DNN: https://github.com/intel/mkl-dnn