TorchDispatchModeTLS.cpp

#include <ATen/core/TorchDispatchModeTLS.h>
#include <c10/core/SafePyObject.h>
#include <c10/core/DispatchKeySet.h>

namespace at { namespace impl {

thread_local std::shared_ptr<SafePyObject> torchDispatchModeState;

void TorchDispatchModeTLS::set_state(std::shared_ptr<SafePyObject> state) {
  if (state) {
    c10::impl::tls_set_dispatch_key_included(DispatchKey::Python, true);
    c10::impl::tls_set_dispatch_key_included(DispatchKey::PythonTLSSnapshot, true);
  } else {
    TorchDispatchModeTLS::reset_state();
  }
  torchDispatchModeState = std::move(state);
}

const std::shared_ptr<SafePyObject>& TorchDispatchModeTLS::get_state() {
  return torchDispatchModeState;
}

void TorchDispatchModeTLS::reset_state() {
  torchDispatchModeState.reset();
  c10::impl::tls_set_dispatch_key_included(DispatchKey::Python, false);
  c10::impl::tls_set_dispatch_key_included(DispatchKey::PythonTLSSnapshot, false);
}

bool dispatch_mode_enabled() {
  return static_cast<bool>(at::impl::TorchDispatchModeTLS::get_state());
}

bool tensor_has_dispatch(const at::Tensor& t) {
  DispatchKeySet key_set({DispatchKey::Python, DispatchKey::PythonTLSSnapshot});
  return t.key_set().has_any(key_set);
}

bool tensorlist_has_dispatch(const at::TensorList& li) {
  for (const auto& t: li) {
    if (tensor_has_dispatch(t)) {
      return true;
    }
  }
  return false;
}

bool tensorlist_has_dispatch(const c10::List<c10::optional<at::Tensor>>& li) {
  for (auto i : c10::irange(li.size())) {
    auto t = li.get(i);
    if (t && tensor_has_dispatch(*t)) {
      return true;
    }
  }
  return false;
}

} // namespace impl
} // namespace at