frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-03 07:24:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "Python Dispatcher integration with C++ dispatcher (#84826)"

Browse Source 
This reverts commit 35f6a69191ef762cf22b6cbfe94b8d9406e16674.

Reverted https://github.com/pytorch/pytorch/pull/84826 on behalf of https://github.com/malfet due to Broke dynamo, see 35f6a69191
This commit is contained in:
PyTorch MergeBot
2022-09-14 14:07:58 +00:00
parent 74ead61944
commit 706b990306
33 changed files with 160 additions and 686 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
torch/_dispatch/python.py
View File

@ -1,88 +0,0 @@
import torch
from contextlib import contextmanager
__all__ = ['enable_python_dispatcher', 'no_python_dispatcher']
DispatchKey = torch._C.DispatchKey # type: ignore[attr-defined]
def has_key(op, k):
return (
torch._C._dispatch_has_kernel_for_dispatch_key(op.name(), k)
or k in op.py_kernels
)
is_included_in_alias = torch._C._dispatch_is_included_in_alias
# Equivalent to computeDispatchTableEntryWithDebug
# TODO: memoize this or something
def resolve_key(op: torch._ops.PyOperatorABC, k: DispatchKey): # type: ignore[valid-type]
# 1. (Direct) operator registration
if has_key(op, k):
return k
# 2.1 Use CompositeExplicitAutogradNonFunctional kernel if available
cand = DispatchKey.CompositeExplicitAutogradNonFunctional
if (k == DispatchKey.Undefined or is_included_in_alias(k, cand)) and has_key(op, cand):
return cand
# 2.2 Use CompositeExplicitAutograd kernel if available
cand = DispatchKey.CompositeExplicitAutograd
if (k == DispatchKey.Undefined or is_included_in_alias(k, cand)) and has_key(op, cand):
return cand
has_backend_kernel = (
torch._C._dispatch_has_kernel_for_any_dispatch_key(op.name(), torch._C._dispatch_get_backend_keyset_from_autograd(k))
or has_key(op, DispatchKey.CompositeExplicitAutograd)
)
# 2.3. Use CompositeImplicitAutograd kernel if available
cand = DispatchKey.CompositeImplicitAutogradNestedTensor
if (
(k != DispatchKey.Undefined and is_included_in_alias(k, cand)) # type: ignore[attr-defined]
and has_key(op, cand) and not has_backend_kernel):
return cand
cand = DispatchKey.CompositeImplicitAutograd
if (k == DispatchKey.Undefined or is_included_in_alias(k, cand)) and has_key(op, cand):
if (
k == DispatchKey.AutogradOther
and torch._C._dispatch_has_kernel_for_any_dispatch_key(op.name(), torch._C._dispatch_autogradother_backends) # type: ignore[attr-defined] # noqa: B950
):
raise RuntimeError("ambiguous autogradother kernel")
elif not has_backend_kernel:
return cand
# 2.4. For autograd backend keys, use kernel from DispatchKey::Autograd if available
cand = DispatchKey.Autograd
if is_included_in_alias(k, cand) and has_key(op, cand):
return cand
# Backend fallback
if torch._C._dispatch_has_backend_fallback(k):
# The dispatch key itself will implicitly route to backend fallback.
# This is probably not great for the pure Python implementation.
return k
raise RuntimeError("could not find kernel")
@contextmanager
def no_python_dispatcher():
g = torch._C._DisablePythonDispatcher()
try:
yield
finally:
del g
@contextmanager
def enable_python_dispatcher():
g = torch._C._EnablePythonDispatcher()
try:
yield
finally:
del g
# The Python dispatcher
def python_dispatcher(op, ks, args, kwargs):
"""
with no_python_dispatcher():
print(op, ks, args, kwargs)
"""
k = resolve_key(op, ks.highestPriorityTypeId())
source = f'torch.ops.{op}.dispatch(k, *args, **kwargs)'
filename = f'{op}[{torch._C._dispatch_key_name(k)}]'
compiled = compile(source, filename, 'eval') # TODO: maybe cache?
return eval(compiled, {'torch': torch, 'k': k, 'args': args, 'kwargs': kwargs})
torch._C._set_python_dispatcher(python_dispatcher)