mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 21:49:24 +08:00
e1c4548441
## `VariableTracker::build()` hides the Builders
### The problem
In the current code, creating a `VariableTracker` involves choosing one of two `Builder` classes and either calling a method, or calling a constructor that creates an object that you immediately call, [like this](083c9149b7/torch/_dynamo/variables/functions.py (L761-L768)).
Variations on this code are repeated in many places.
More, the `Builder` classes have a lot of dependencies, so they have to be loaded late in the whole import process to avoid circular imports, so they end up being repeatedly imported at local scope.
### The solution
In this commit, the import from `builder` and the logic of choosing and calling the Builder class are hidden in a single static factory method, `VariableTracker.build()`, easier to reason about and to import.
This commit net lowers the total lines of code by over 150 lines by removing repetitive logic and unnecessary local imports.
**CHANGES:** Originally the name of the static method was `VariableTracker.create()` but a static method on a derived class, `LazyVariableTracker.create()` now exists with a different signature that's irreconcilable, so the new static method was renamed to `VariableTracker.build()`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/135714
Approved by: https://github.com/jansel
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
# mypy: ignore-errors
|
|
|
|
from inspect import getattr_static
|
|
from typing import TYPE_CHECKING
|
|
|
|
from ..bytecode_transformation import create_call_function
|
|
from ..exc import Unsupported
|
|
from ..source import AttrSource
|
|
from .base import VariableTracker
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from torch._dynamo.symbolic_convert import InstructionTranslator
|
|
|
|
PARAM_NAMES = "query key value attn_mask dropout is_causal enable_gqa".split()
|
|
|
|
|
|
class SDPAParamsVariable(VariableTracker):
|
|
"""Represents the c++ params struct for scaled dot product attention.
|
|
This is a read-only container."""
|
|
|
|
@staticmethod
|
|
def create(tx: "InstructionTranslator", value, source):
|
|
from torch.backends.cuda import SDPAParams
|
|
|
|
from .torch import TorchInGraphFunctionVariable
|
|
|
|
params = [
|
|
VariableTracker.build(tx, getattr(value, p), AttrSource(source, p))
|
|
for p in PARAM_NAMES
|
|
]
|
|
return TorchInGraphFunctionVariable(SDPAParams).call_function(tx, params, {})
|
|
|
|
def __init__(self, proxy, param_vars, **kwargs) -> None:
|
|
self.proxy = proxy
|
|
self.param_vars = param_vars
|
|
super().__init__(**kwargs)
|
|
|
|
def reconstruct(self, codegen):
|
|
assert self.source is None
|
|
assert self.param_vars is not None
|
|
codegen.add_push_null(
|
|
lambda: codegen.load_import_from("torch._C", "_SDPAParams")
|
|
)
|
|
codegen.foreach(self.param_vars)
|
|
codegen.extend_output(create_call_function(len(self.param_vars), False))
|
|
|
|
def as_proxy(self):
|
|
return self.proxy
|
|
|
|
def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
|
|
import torch._C
|
|
|
|
from .builder import wrap_fx_proxy
|
|
from .misc import GetAttrVariable
|
|
|
|
try:
|
|
getattr_static(torch._C._SDPAParams, name)
|
|
except AttributeError:
|
|
# Using raise from is too verbose here
|
|
raise Unsupported(
|
|
f"Unsupported torch._C._SDPAParams attribute {name}"
|
|
) from None
|
|
|
|
proxy = GetAttrVariable.create_getattr_proxy(self.as_proxy(), name)
|
|
if self.source is not None:
|
|
return wrap_fx_proxy(
|
|
tx=tx, proxy=proxy, source=AttrSource(self.source, name)
|
|
)
|
|
else:
|
|
return wrap_fx_proxy(tx=tx, proxy=proxy)
|
|
|
|
@staticmethod
|
|
def is_sdpa_params(value):
|
|
from torch.backends.cuda import SDPAParams
|
|
|
|
return value is SDPAParams
|