RAII guard that pins libtorch's intra-op thread count for the current scope and restores the previous value on destruction. More...

#include <NEML2Utils.h>

Public Member Functions
	ScopedNumThreads (int n)

	~ScopedNumThreads ()

	ScopedNumThreads (const ScopedNumThreads &)=delete

ScopedNumThreads &	operator= (const ScopedNumThreads &)=delete

Private Attributes
const int	_saved

const bool	_changed

Detailed Description

RAII guard that pins libtorch's intra-op thread count for the current scope and restores the previous value on destruction.

NEML2 models are evaluated with MPI as the parallelism axis (each rank owns a slice of the element batch), and the AOTI-compiled graphs get essentially no speedup from libtorch's intra-op threads. Left at libtorch's default (the core count), those per-rank thread pools oversubscribe the node under MPI and degrade strong scaling – so NEML2 pins the count around each evaluation. We save and restore rather than set-once because other parts of MOOSE/libtorch may change the global value, and we do not want to disturb it outside NEML2's evaluation.

Intra-op only: at::set_num_interop_threads is one-shot (it throws once the interop pool has been used) and so cannot be scoped; the interop count is left to the application-level setup.

Definition at line 214 of file NEML2Utils.h.

Constructor & Destructor Documentation

◆ ScopedNumThreads() [1/2]

NEML2Utils::ScopedNumThreads::ScopedNumThreads ( int n )

inlineexplicit

Definition at line 217 of file NEML2Utils.h.

                                   : _saved(at::get_num_threads()), _changed(n > 0 && n != _saved)
  {
    if (_changed)
      at::set_num_threads(n);
  }

◆ ~ScopedNumThreads()

NEML2Utils::ScopedNumThreads::~ScopedNumThreads ( )

inline

Definition at line 222 of file NEML2Utils.h.

  {
    if (_changed)
      at::set_num_threads(_saved);
  }

◆ ScopedNumThreads() [2/2]

NEML2Utils::ScopedNumThreads::ScopedNumThreads ( const ScopedNumThreads & )

delete

Member Function Documentation

◆ operator=()

ScopedNumThreads & NEML2Utils::ScopedNumThreads::operator= ( const ScopedNumThreads & )

delete

Member Data Documentation

◆ _changed

const bool NEML2Utils::ScopedNumThreads::_changed

private

Definition at line 232 of file NEML2Utils.h.

Referenced by ScopedNumThreads(), and ~ScopedNumThreads().

◆ _saved

const int NEML2Utils::ScopedNumThreads::_saved

private

Definition at line 231 of file NEML2Utils.h.

Referenced by ~ScopedNumThreads().

The documentation for this class was generated from the following file:

include/neml2/utils/NEML2Utils.h

Public Member Functions

Private Attributes