417 lines
16 KiB
Python
417 lines
16 KiB
Python
|
import inspect
|
||
|
|
import math
|
||
|
|
import operator
|
||
|
|
from collections.abc import Iterable
|
||
|
|
from typing import Any, Dict, final, List, Optional, Tuple, Type
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch._ops import HigherOrderOperator, OpOverload
|
||
|
|
from torch._subclasses.fake_tensor import FakeTensor
|
||
|
|
from torch.export.exported_program import ExportedProgram
|
||
|
|
from torch.export.graph_signature import (
|
||
|
|
CustomObjArgument,
|
||
|
|
InputKind,
|
||
|
|
SymIntArgument,
|
||
|
|
TensorArgument,
|
||
|
|
)
|
||
|
|
from torch.fx import GraphModule
|
||
|
|
from torch.fx.experimental.symbolic_shapes import SymBool, SymFloat, SymInt
|
||
|
|
|
||
|
|
|
||
|
|
class SpecViolationError(Exception):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def is_functional(op: OpOverload) -> bool:
|
||
|
|
return not op._schema.is_mutable
|
||
|
|
|
||
|
|
|
||
|
|
def _check_has_fake_tensor(node: torch.fx.Node) -> None:
|
||
|
|
# TODO(angelayi): remove this in favor of _check_val
|
||
|
|
return _check_val(node)
|
||
|
|
|
||
|
|
|
||
|
|
def _check_val(node: torch.fx.Node) -> None:
|
||
|
|
def _check_correct_val(val):
|
||
|
|
if val is None:
|
||
|
|
return True
|
||
|
|
elif isinstance(val, (int, bool, str, float)):
|
||
|
|
return True
|
||
|
|
elif isinstance(val, (torch.memory_format, torch.dtype, torch.device, torch.layout)):
|
||
|
|
return True
|
||
|
|
elif isinstance(val, (FakeTensor, torch.Tensor)): # TODO(zhxchen17) Remove Tensor.
|
||
|
|
return True
|
||
|
|
elif isinstance(val, (SymInt, SymFloat, SymBool)):
|
||
|
|
return True
|
||
|
|
elif isinstance(val, CustomObjArgument):
|
||
|
|
return True
|
||
|
|
elif isinstance(val, Iterable):
|
||
|
|
return all(_check_correct_val(x) for x in val)
|
||
|
|
return False
|
||
|
|
|
||
|
|
def _no_returns(op):
|
||
|
|
if not isinstance(op, OpOverload):
|
||
|
|
return False
|
||
|
|
return len(op._schema.returns) == 0
|
||
|
|
|
||
|
|
if "val" not in node.meta:
|
||
|
|
if node.op == "call_function" and _no_returns(node.target):
|
||
|
|
return
|
||
|
|
raise SpecViolationError(f"Node.meta {node.name} is missing val field.")
|
||
|
|
|
||
|
|
val = node.meta["val"]
|
||
|
|
if not _check_correct_val(val):
|
||
|
|
raise SpecViolationError(f"Node.meta {node.name} has invalid val field {val}")
|
||
|
|
|
||
|
|
|
||
|
|
class _VerifierMeta(type):
|
||
|
|
_registry: Dict[str, Type['Verifier']] = {}
|
||
|
|
|
||
|
|
def __new__(metacls, name, bases, attrs):
|
||
|
|
if bases:
|
||
|
|
if "check" in attrs or "_check_graph_module" in attrs:
|
||
|
|
raise SyntaxError("Overriding method check is not allowed.")
|
||
|
|
assert "dialect" in attrs and attrs["dialect"] != "ATEN"
|
||
|
|
else:
|
||
|
|
assert "check" in attrs
|
||
|
|
assert "_check_graph_module" in attrs
|
||
|
|
assert attrs["dialect"] == "ATEN"
|
||
|
|
|
||
|
|
assert isinstance(attrs["dialect"], str)
|
||
|
|
ret = type.__new__(metacls, name, bases, attrs)
|
||
|
|
metacls._registry[attrs["dialect"]] = ret # type: ignore[assignment]
|
||
|
|
return ret
|
||
|
|
|
||
|
|
def getattr_recursive(obj: Any, target: str) -> Any:
|
||
|
|
target_atoms = target.split('.')
|
||
|
|
attr_itr = obj
|
||
|
|
for i, atom in enumerate(target_atoms):
|
||
|
|
if not hasattr(attr_itr, atom):
|
||
|
|
raise RuntimeError(f"Node referenced nonexistent target {'.'.join(target_atoms[:i])}")
|
||
|
|
attr_itr = getattr(attr_itr, atom)
|
||
|
|
return attr_itr
|
||
|
|
|
||
|
|
|
||
|
|
class Verifier(metaclass=_VerifierMeta):
|
||
|
|
dialect = "ATEN"
|
||
|
|
|
||
|
|
def allowed_builtin_ops(self) -> List:
|
||
|
|
return [
|
||
|
|
operator.getitem,
|
||
|
|
operator.add,
|
||
|
|
operator.mul,
|
||
|
|
operator.sub,
|
||
|
|
operator.truediv,
|
||
|
|
operator.ge,
|
||
|
|
operator.le,
|
||
|
|
operator.gt,
|
||
|
|
operator.lt,
|
||
|
|
operator.eq,
|
||
|
|
operator.ne,
|
||
|
|
operator.floordiv,
|
||
|
|
operator.mod,
|
||
|
|
operator.and_,
|
||
|
|
operator.or_,
|
||
|
|
operator.not_,
|
||
|
|
operator.pow,
|
||
|
|
operator.neg,
|
||
|
|
operator.abs,
|
||
|
|
math.ceil,
|
||
|
|
math.floor,
|
||
|
|
]
|
||
|
|
|
||
|
|
def allowed_op_types(self) -> Tuple[Type[Any], ...]:
|
||
|
|
return (OpOverload, HigherOrderOperator)
|
||
|
|
|
||
|
|
def allowed_getattr_types(self) -> Tuple[Type[Any], ...]:
|
||
|
|
return (torch.fx.GraphModule,)
|
||
|
|
|
||
|
|
def check_valid_op(self, op):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def check_additional(self, gm: GraphModule) -> None:
|
||
|
|
"""
|
||
|
|
Additional checks that are specific to some dialects.
|
||
|
|
"""
|
||
|
|
pass
|
||
|
|
|
||
|
|
@final
|
||
|
|
def check(self, ep: ExportedProgram) -> None:
|
||
|
|
self._check_graph_module(ep.graph_module)
|
||
|
|
_verify_exported_program_signature(ep)
|
||
|
|
|
||
|
|
@final
|
||
|
|
def _check_graph_module(self, gm: torch.fx.GraphModule) -> None:
|
||
|
|
def _allowed_getattr_types() -> Tuple[Type[Any], ...]:
|
||
|
|
ret = self.allowed_getattr_types()
|
||
|
|
assert not any(t is object for t in ret)
|
||
|
|
return ret
|
||
|
|
|
||
|
|
def _check_valid_op(op) -> None:
|
||
|
|
def _allowed_builtin_ops() -> List:
|
||
|
|
ret = self.allowed_builtin_ops()
|
||
|
|
assert all(inspect.isbuiltin(op) for op in ret)
|
||
|
|
return ret
|
||
|
|
|
||
|
|
def _allowed_op_types() -> Tuple[Type[Any], ...]:
|
||
|
|
ret = self.allowed_op_types()
|
||
|
|
assert not any(t is object for t in ret)
|
||
|
|
return ret
|
||
|
|
|
||
|
|
# TODO Remove this allowlist.
|
||
|
|
_allowed_torch_functions = (
|
||
|
|
torch.autograd.grad_mode.set_grad_enabled,
|
||
|
|
torch.sym_int,
|
||
|
|
torch.sym_ite,
|
||
|
|
torch.sym_max,
|
||
|
|
torch.sym_min,
|
||
|
|
torch.sym_not,
|
||
|
|
torch.sym_sqrt,
|
||
|
|
# TODO (tmanlaibaatar)
|
||
|
|
# Predispatch export is able to contain autograd ops.
|
||
|
|
# These will be modeled as HOO later
|
||
|
|
torch._C._set_grad_enabled
|
||
|
|
|
||
|
|
)
|
||
|
|
|
||
|
|
if not isinstance(op, _allowed_op_types()):
|
||
|
|
if op not in _allowed_builtin_ops() and op not in _allowed_torch_functions:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Operator '{op}' is not an allowed operator type: {_allowed_op_types()}\n"
|
||
|
|
f"Valid builtin ops: {_allowed_builtin_ops()}"
|
||
|
|
f"Valid torch functions: {_allowed_torch_functions}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if isinstance(op, OpOverload):
|
||
|
|
# All ops functional
|
||
|
|
if not is_functional(op):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"operator '{op}' is not functional"
|
||
|
|
)
|
||
|
|
self.check_valid_op(op)
|
||
|
|
|
||
|
|
for mod in gm.modules():
|
||
|
|
if not isinstance(mod, torch.fx.GraphModule):
|
||
|
|
continue
|
||
|
|
|
||
|
|
mod.graph.lint()
|
||
|
|
for node in mod.graph.nodes:
|
||
|
|
# TODO(T140410192): should have fake tensor for all dialects
|
||
|
|
if node.op in {"call_module", "call_method"}:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"call_module is not valid: got a class '{node.target}' ",
|
||
|
|
)
|
||
|
|
|
||
|
|
elif node.op == "call_function":
|
||
|
|
_check_val(node)
|
||
|
|
|
||
|
|
_check_valid_op(node.target)
|
||
|
|
|
||
|
|
elif node.op == "get_attr":
|
||
|
|
if not isinstance(node.target, str):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Expected get_attr target to be string, but got {type(node.target)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
attr = getattr_recursive(mod, node.target)
|
||
|
|
if isinstance(attr, torch.nn.Module):
|
||
|
|
def _is_type(name, ty):
|
||
|
|
return isinstance(getattr(attr, name, None), ty)
|
||
|
|
if type(attr).__name__ == "LoweredBackendModule":
|
||
|
|
if _is_type("backend_id", str) \
|
||
|
|
and _is_type("processed_bytes", bytes) \
|
||
|
|
and _is_type("compile_specs", list) \
|
||
|
|
and hasattr(attr, "original_module"):
|
||
|
|
continue
|
||
|
|
else:
|
||
|
|
backend_id = getattr(attr, "backend_id", None)
|
||
|
|
processed_bytes = getattr(attr, "processed_bytes", None)
|
||
|
|
compile_specs = getattr(attr, "compile_specs", None)
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Invalid get_attr type {type(attr)}. \n"
|
||
|
|
f"LoweredBackendModule fields: "
|
||
|
|
f"backend_id(str) : {type(backend_id)}, "
|
||
|
|
f"processed_bytes(bytes) : {type(processed_bytes)}, "
|
||
|
|
f"compile_specs(list) : {type(compile_specs)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if not isinstance(attr, _allowed_getattr_types()):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Invalid get_attr type {type(attr)}. \n"
|
||
|
|
f"Valid get_attr types: {_allowed_getattr_types()}"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
elif node.op == "placeholder":
|
||
|
|
_check_val(node)
|
||
|
|
# TODO(zhxchen17)
|
||
|
|
# elif node.op == "output":
|
||
|
|
# _check_flattened_outputs()
|
||
|
|
|
||
|
|
self.check_additional(gm)
|
||
|
|
|
||
|
|
|
||
|
|
def _verify_exported_program_signature(exported_program) -> None:
|
||
|
|
# Check ExportedProgram signature matches
|
||
|
|
gs = exported_program.graph_signature
|
||
|
|
|
||
|
|
# Check every node in the signature exists in the graph
|
||
|
|
input_node_names = [node.name for node in exported_program.graph.nodes if node.op == "placeholder"]
|
||
|
|
|
||
|
|
if len(input_node_names) != len(gs.input_specs):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Number of graph inputs ({len(input_node_names)}) "
|
||
|
|
f"does not match number of inputs in the graph signature ({len(gs.user_inputs)})"
|
||
|
|
)
|
||
|
|
|
||
|
|
for input_spec, node in zip(gs.input_specs, input_node_names):
|
||
|
|
if isinstance(input_spec.arg, (TensorArgument, SymIntArgument)):
|
||
|
|
if input_spec.arg.name != node:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Input spec name {input_spec.arg.name} does not match node name {node}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if input_spec.kind == InputKind.USER_INPUT:
|
||
|
|
continue
|
||
|
|
|
||
|
|
elif input_spec.kind == InputKind.PARAMETER:
|
||
|
|
if not isinstance(input_spec.arg, TensorArgument):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Parameter {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
|
||
|
|
)
|
||
|
|
if input_spec.target is None:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"InputSpec for {input_spec.name} has no target."
|
||
|
|
)
|
||
|
|
|
||
|
|
param = input_spec.target
|
||
|
|
if param not in exported_program.state_dict:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Parameter {param} is not in the state dict."
|
||
|
|
)
|
||
|
|
|
||
|
|
if not isinstance(exported_program.state_dict[param], torch.nn.Parameter):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"State dict entry for parameter {param} is not an instance of torch.nn.Parameter."
|
||
|
|
)
|
||
|
|
|
||
|
|
elif input_spec.kind == InputKind.BUFFER:
|
||
|
|
if not isinstance(input_spec.arg, TensorArgument):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Buffer {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
|
||
|
|
)
|
||
|
|
if input_spec.target is None:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"InputSpec for {input_spec.name} has no target."
|
||
|
|
)
|
||
|
|
|
||
|
|
buffer = input_spec.target
|
||
|
|
if input_spec.persistent is None:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Buffer {buffer} is missing a persistence flag"
|
||
|
|
)
|
||
|
|
|
||
|
|
if input_spec.persistent is True and buffer not in exported_program.state_dict:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Buffer {buffer} is not in the state dict."
|
||
|
|
)
|
||
|
|
|
||
|
|
if input_spec.persistent is False and buffer in exported_program.state_dict:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Non-persistent buffer {buffer} is in the state dict, it should not be."
|
||
|
|
)
|
||
|
|
elif input_spec.kind == InputKind.CONSTANT_TENSOR:
|
||
|
|
if not isinstance(input_spec.arg, TensorArgument):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Constant tensor {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
|
||
|
|
)
|
||
|
|
if input_spec.target is None:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"InputSpec for {input_spec.name} has no target."
|
||
|
|
)
|
||
|
|
|
||
|
|
tensor_const = input_spec.target
|
||
|
|
if tensor_const not in exported_program.constants:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Constant tensor {tensor_const} is not in the constants dictionary."
|
||
|
|
)
|
||
|
|
elif input_spec.kind == InputKind.CUSTOM_OBJ:
|
||
|
|
if not isinstance(input_spec.arg, CustomObjArgument):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Custom object {input_spec.name} is not a custom object argument. Found {input_spec.arg} instead."
|
||
|
|
)
|
||
|
|
if input_spec.target is None:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"InputSpec for {input_spec.name} has no target."
|
||
|
|
)
|
||
|
|
|
||
|
|
custom_obj = input_spec.target
|
||
|
|
if custom_obj not in exported_program.constants:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Custom object {custom_obj} is not in the constants dictionary."
|
||
|
|
)
|
||
|
|
elif input_spec.kind == InputKind.TOKEN:
|
||
|
|
if not isinstance(input_spec.arg, TensorArgument):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Constant tensor {input_spec.name} is not a tensor argument. Found {input_spec.arg} instead."
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Unknown InputKind {input_spec.kind}."
|
||
|
|
)
|
||
|
|
|
||
|
|
# Check outputs
|
||
|
|
output_node = list(exported_program.graph.nodes)[-1]
|
||
|
|
assert output_node.op == "output"
|
||
|
|
output_nodes = [
|
||
|
|
arg.name if isinstance(arg, torch.fx.Node) else arg
|
||
|
|
for arg in output_node.args[0]
|
||
|
|
]
|
||
|
|
|
||
|
|
if len(output_nodes) != len(gs.output_specs):
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Number of output nodes {len(output_nodes)} is different "
|
||
|
|
"Than the number of outputs specified by the graph signature: \n"
|
||
|
|
f"Number of mutated buffers: {len(gs.buffers_to_mutate)}. \n"
|
||
|
|
f"Number of user outputs: {len(gs.user_outputs)}. \n"
|
||
|
|
)
|
||
|
|
|
||
|
|
num_tokens = len(gs.output_tokens)
|
||
|
|
end = len(gs.buffers_to_mutate) + len(gs.user_inputs_to_mutate) + num_tokens
|
||
|
|
mutate_nodes: List[str] = output_nodes[num_tokens:end]
|
||
|
|
user_output_nodes = output_nodes[end:end + len(gs.user_outputs)]
|
||
|
|
|
||
|
|
for mutation_node in mutate_nodes:
|
||
|
|
if mutation_node in gs.buffers_to_mutate:
|
||
|
|
if gs.buffers_to_mutate[mutation_node] not in gs.buffers:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Buffer output {mutation_node} does not point to a buffer that exists. \n"
|
||
|
|
f"Dict of buffers that are mutated, in order: {gs.buffers_to_mutate} \n"
|
||
|
|
f"Buffer nodes available: {gs.buffers} \n"
|
||
|
|
)
|
||
|
|
elif mutation_node in gs.user_inputs_to_mutate:
|
||
|
|
if gs.user_inputs_to_mutate[mutation_node] not in gs.user_inputs:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"User input output {mutation_node} does not point to a user input that exists. \n"
|
||
|
|
f"Dict of user inputs that are mutated, in order: {gs.user_inputs_to_mutate} \n"
|
||
|
|
f"User input nodes available: {gs.user_inputs} \n")
|
||
|
|
else:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"Mutation node {mutation_node} is neither a buffer nor a user input. "
|
||
|
|
f"Buffers to mutate: {gs.buffers_to_mutate}, User inputs to mutate: {gs.user_inputs_to_mutate}"
|
||
|
|
)
|
||
|
|
|
||
|
|
for user_output_node, user_output_name in zip(user_output_nodes, gs.user_outputs):
|
||
|
|
if user_output_node != user_output_name:
|
||
|
|
raise SpecViolationError(
|
||
|
|
f"User output {user_output_node} is not in the correct "
|
||
|
|
"order or is not found in the "
|
||
|
|
f"exported program's user_output list: {gs.user_outputs}. "
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def load_verifier(dialect: str) -> Optional[Type[Verifier]]:
|
||
|
|
if dialect == "ATEN":
|
||
|
|
return _VerifierMeta._registry.get(dialect)
|
||
|
|
return _VerifierMeta._registry[dialect]
|