Spaces:
Running
Running
File size: 4,862 Bytes
c61ccee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
import dataclasses
import inspect
import sys
from typing import Any, Callable, Tuple
import torch
@dataclasses.dataclass
class Kernel:
"""Models a (function, source location)"""
func: Callable
source: str
def __call__(self, *args, **kwargs):
return self.func(*args, **kwargs)
class RegistrationHandle:
"""Does something when someone calls .destroy() on it"""
def __init__(self, on_destroy: Callable):
self._on_destroy = on_destroy
def destroy(self) -> None:
self._on_destroy()
def get_source(stacklevel: int) -> str:
"""Get a string that represents the caller.
Example: "/path/to/foo.py:42"
Use stacklevel=1 to get the caller's source
Use stacklevel=2 to get the caller's caller's source
etc.
"""
frame = inspect.getframeinfo(sys._getframe(stacklevel))
source = f"{frame.filename}:{frame.lineno}"
return source
def parse_namespace(qualname: str) -> Tuple[str, str]:
splits = qualname.split("::")
if len(splits) != 2:
raise ValueError(
f"Expected `qualname` to be of the form "
f'"namespace::name", but got {qualname}. '
f"The qualname passed to the torch.library APIs must consist "
f"of a namespace and a name, e.g. aten::sin"
)
return splits[0], splits[1]
def lookup_op(qualname: str) -> torch._ops.OpOverloadPacket:
namespace, name = parse_namespace(qualname)
if "." in name:
name, overload = name.split(".")
else:
overload = "default"
ns = getattr(torch.ops, namespace)
packet = getattr(ns, name)
return getattr(packet, overload)
def is_builtin(op: torch._ops.OpOverload) -> bool:
assert isinstance(op, torch._ops.OpOverload)
return op.namespace in {"aten", "prim", "prims"}
def is_functional_schema(schema: Any) -> bool:
"""Check if the schema is functional.
An operator is functional if:
- it does not mutate any of its inputs
- it does not return a view on any of its inputs
- it has at least one return
"""
# Lazy import because not all PyTorch builds have torchgen
from torchgen.model import FunctionSchema, SchemaKind
assert isinstance(schema, (str, FunctionSchema))
if isinstance(schema, str):
schema = FunctionSchema.parse(schema)
if schema.kind() != SchemaKind.functional:
return False
rets = schema.returns
is_non_mutating_view = len(rets) > 0 and any(
r.annotation is not None and not r.annotation.is_write for r in rets
)
if is_non_mutating_view:
return False
if not schema.returns:
return False
return True
def mutates_and_returns_first_arg(op: torch._ops.OpOverload):
"""Check if an op is an inplace aten op, i.e. it mutates and returns the first arg.
TODO: torchgen/model.py's FunctionSchema.parse is the source of truth for this,
but not all PyTorch builds have torchgen (due to the yaml dependency being weird).
Figure this out.
Example: add_(Tensor(a!) x, Tensor y) -> Tensor(a)
"""
if op.namespace != "aten":
return False
schema = op._schema
if not len(schema.returns) == 1:
return False
if schema.returns[0].alias_info is None:
return False
alias_set = schema.returns[0].alias_info.after_set
if len(alias_set) != 1:
return False
loc = next(iter(alias_set))
if len(schema.arguments) < 1:
return False
first_arg = schema.arguments[0]
if first_arg.alias_info is None:
return False
if not first_arg.alias_info.is_write:
return False
alias_set = first_arg.alias_info.after_set
if len(alias_set) != 1:
return False
if loc != next(iter(alias_set)):
return False
for arg in schema.arguments[1:]:
if arg.alias_info is not None:
return False
return True
def zip_schema(schema, args, kwargs):
"""zips schema.arguments and (args, kwargs) together.
Assumes that (args, kwargs) were the inputs to some torch._ops.OpOverload:
that is, kwargs must be keyword-only arguments and default values may be omitted.
"""
assert len(schema.arguments) >= len(args) + len(kwargs)
for i in range(len(schema.arguments)):
info = schema.arguments[i]
if info.kwarg_only:
if info.name in kwargs:
yield info, kwargs[info.name]
continue
if i >= len(args):
# args that are equal to their default values are not populated
# if they are followed by args that are equal to their defaults.
# Skip these.
continue
yield info, args[i]
return
|