mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
48ee0984ac
This PR adds a C function to check if all torch function is disabled. Recall that there are three torch function enablement states: * All disabled * Torch Function Subclass disabled * All enabled The API before this change provides two functions: * `_is_torch_function_enabled` - returns True iff the current TF state is All enabled * `_is_torch_function_mode_enabled` - returns True iff the state is not All disabled and the torch function mode stack is non-empty. The crux of why a new API is needed is the following: If dynamo enters a frame with the torch function mode stack empty, `_is_torch_function_enabled` == False, it is impossible to determine if after a new mode is pushed whether we should enter the mode or not. This is because we don't know if the enablement state is All disabled or only subclass disabled. Adding this API to check if All disabled is True allows us to disambiguate this case. In the next PR, Dynamo InstructionTranslator will have clearer flags than the underlying C API: * A flag to indicate if subclasses are disabled (ie All disabled or Subclass Disabled is the current state) * A flag to indicate if modes are disabled (ie if All disabled is the current state) * A symbolic stack which can be checked if any modes are present Pull Request resolved: https://github.com/pytorch/pytorch/pull/133136 Approved by: https://github.com/bdhirsh ghstack dependencies: #133130, #133729, #133131, #133132, #133133, #133134
37 lines
1.2 KiB
C++
37 lines
1.2 KiB
C++
#pragma once
|
|
|
|
#include <c10/core/SafePyObject.h>
|
|
#include <c10/macros/Macros.h>
|
|
|
|
namespace at::impl {
|
|
|
|
enum TorchFunctionDisabledState { ENABLED, SUBCLASSES_DISABLED, ALL_DISABLED };
|
|
|
|
struct TORCH_API PythonTorchFunctionTLS {
|
|
static void set_disabled_state(TorchFunctionDisabledState disabled_state_);
|
|
static TorchFunctionDisabledState get_disabled_state();
|
|
|
|
static void push_onto_stack(std::shared_ptr<SafePyObject> mode);
|
|
static const std::shared_ptr<SafePyObject> pop_stack();
|
|
static const std::shared_ptr<SafePyObject>& get_stack_at(int64_t idx);
|
|
static int64_t stack_len();
|
|
|
|
static const PythonTorchFunctionTLS& get_state();
|
|
static void set_state(const PythonTorchFunctionTLS& state);
|
|
|
|
private:
|
|
// The mode TLS is split into
|
|
// - disabled_state, which says which part of torch function are disabled
|
|
// - stack_, which is a vector of modes representing the stack of user
|
|
// defined modes
|
|
TorchFunctionDisabledState disabled_state_ =
|
|
TorchFunctionDisabledState::ENABLED;
|
|
std::vector<std::shared_ptr<c10::SafePyObject>> stack_;
|
|
};
|
|
|
|
TORCH_API bool torch_function_mode_enabled();
|
|
|
|
TORCH_API bool torch_function_all_disabled();
|
|
|
|
} // namespace at::impl
|