Custom ops support arbitrary input types by migrating to python dispatcher #147927

yanboliang · 2025-02-26T05:37:58Z

Test case:

@torch.library.custom_op("mylib::foo", mutates_args
10000
=())
def foo(d: dict, t: torch.Tensor) -> torch.Tensor:
    return torch.sin(d["x"] - d["y"] + t)


@foo.register_fake
def _(d: dict, t: torch.Tensor) -> torch.Tensor:
    return torch.empty_like(d["x"])

d = {"x": torch.randn(2, 3, requires_grad=True), "y": torch.randn(2, 3, requires_grad=True)}
t = torch.randn(2, 3, requires_grad=True)

@torch.compile(backend="eager", fullgraph=True)
def fn(d, t):
    return torch.sin(torch.ops.mylib.foo.default(d, t) + 1.5)

y = fn(d, t)
print(y)
y.sum().backward()
print(d["x"].grad)
print(d["y"].grad)
print(t.grad)

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames

pytorch-bot · 2025-02-26T05:38:02Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/147927

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 14c0439 with merge base 84e60ee ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

torch/_dynamo/variables/torch.py

torch/_library/infer_schema.py

torch/_ops.py

torch/_dynamo/variables/torch.py

torch/_ops.py

zou3519 · 2025-03-05T15:42:11Z

test/test_custom_ops.py

+
+        x = Point(x=torch.randn(2, 3), y=torch.randn(2, 3))
+        y = torch.ops.mylib.foo(x)
+        self.assertEqual(y, torch.sqrt(torch.sum((x.x - x.y) ** 2)))


Some high-level comments:

Let's wait until after the branch cut (Monday) to merge this, assuming its ready before then. We don't want this feature to be partially in PyTorch 2.7.

eager-mode performance is pretty important. Can you do some benchmarking comparing e.g. custom ops with dict input (uses the new path) to custom ops with list inputs (uses the C++ dispatcher)?

Agree! I'll do the perf benchmark while you conduct code review.

zou3519

My main comments are:

eager-mode performance. Let's make sure this is good.
the fallthrough mechanism isn't faithful to the C++ pytorch dispatcher. We should make it more faithful
the FakeTensor registration mechanism is on the sketchy side. In particular, the py_impl(FakeTensorMode) completely bypasses FakeTensor caching. Maybe something to handle as a follow-up.

zou3519 · 2025-03-06T15:16:29Z

torch/_library/infer_schema.py

-            if annotation_type.__origin__ is tuple:
+            if annotation_type in torch.utils._pytree.SUPPORTED_NODES:
+                # TODO: Move to a separate schema type for pytrees.
+                schema_type = "Any"


when are we going to change this?

tbh I kind of like having "Any" in the type, we could keep it for now

zou3519 · 2025-03-06T15:17:59Z

torch/_ops.py

@@ -46,6 +47,23 @@ def dl_open_guard():
        sys.setdlopenflags(old_flags)


+class Kernel:


put this somewhere else. Maybe in torch/_dispatch or torch/_library. Otherwise I think this gets exposed as torch.ops.Kernel (things in torch._ops behave a little weirdly)

zou3519 · 2025-03-06T15:18:33Z

torch/_ops.py

+    It is the thing that is called when you call the operator.
+    """
+
+    def __init__(self, func, with_keyset=False):


nit: make with_keyset kwarg-only for readability

zou3519 · 2025-03-06T15:21:00Z

torch/_ops.py

+    if _has_pytree_type_in_args_or_kwargs(args, kwargs):
+        op_overload = getattr(op, op.overloads()[0])
+        return op_overload(*args, **kwargs)


this is probably slow and needs caching on the OpOverload

this is wrong if someone tries to register multiple OpOverloads for one OpOverloadPacket, so we should error out in that situation

zou3519 · 2025-03-06T15:28:25Z

torch/_library/custom_ops.py

+        need_python_dispatch = isinstance(
+            self._opoverload, torch._ops.PythonDispatcherOpOverload
+        )


nit: should probably make this a helper function

zou3519 · 2025-03-06T15:30:50Z

torch/_library/custom_ops.py

+        if need_python_dispatch:
+            self._opoverload.py_impl(torch._subclasses.fake_tensor.FakeTensorMode)(
+                fake_impl
+            )


NB: this is kind of a hack

zou3519 · 2025-03-06T15:31:44Z

torch/_library/custom_ops.py

+                self._opoverload.py_impl(_C.DispatchKey.ADInplaceOrView)(
+                    adinplaceorview_impl
+                )
+            else:
+                lib.impl(
+                    self._name,
+                    adinplaceorview_impl,
+                    "ADInplaceOrView",
+                    with_keyset=True,
+                )


I feel like we should align the API of py_impl and lib.impl. Otherwise it's kind of annoying

Doesn't need to happen in this PR, just a suggestion

zou3519 · 2025-03-06T15:34:59Z

torch/_ops.py

+        ]
+
+        def _may_use_fallthrough_instead_of_fallback(key: DispatchKey):
+            if torch._C._dispatch_has_kernel_for_dispatch_key(self.name(), key):


If there's a kernel in py_impl as well, then we should not use the fallthrough either

torch/_ops.py

zou3519 · 2025-03-06T16:09:44Z

torch/_ops.py

+        # TODO: we should be calling the fallback for these, but a fallthrough is almost close
+        # enough to the fallback in most cases that we care about.
+        _DEFAULT_FALLTHROUGH_KEYS = [
+            DispatchKey.ADInplaceOrView,
+            DispatchKey.BackendSelect,
+            DispatchKey.PythonTLSSnapshot,
+            DispatchKey.PythonDispatcher,
+        ]


we should try to model fallthroughs closer to how the C++ dispatcher does it. That is, the operator doesn't come with fallthrough keys, instead these are fallbacks.

(maybe in a follow-up)

zou3519 · 2025-03-06T16:14:42Z

torch/_ops.py

+
+    def redispatch(self, /, keyset, *args, **kwargs):
+        return self._dispatch_in_python(args, kwargs, self._fallthrough_keys())
+


At some point I want to try an exercise of "refactor the python dispatcher to look more like the C++ dispatcher". Concretely:

each operator has an OperatorEntry

there is a mapping from DispatchKey->kernel on the OperatorEntry

OperatorEntry has a DispatchKeyExtractor

registering a fallback as a py_kernel should modify the DispatchKeyExtractor

one can use the DispatchKeyExtractor to get the next DispatchKey

DispatchKeys can have fallbacks in Python

etc

Now my question for you is, are you interested in doing this? It could be a good learning experience. If not I'm happy to try this refactor

XuehaiPan · 2025-03-06T16:26:01Z

torch/_ops.py

+def _has_pytree_type_in_args_or_kwargs(args, kwargs) -> bool:
+    return any(
+        not isinstance(x, (list, tuple))
+        and type(x) in torch.utils._pytree.SUPPORTED_NODES
+        for x in itertools.chain(args, kwargs.values())
+    )


Suggested change

def _has_pytree_type_in_args_or_kwargs(args, kwargs) -> bool:

return any(

not isinstance(x, (list, tuple))

and type(x) in torch.utils._pytree.SUPPORTED_NODES

for x in itertools.chain(args, kwargs.values())

)

def _has_pytree_type_in_args_or_kwargs(args, kwargs) -> bool:

def is_list_or_tuple(x):

return isinstance(x, (list, tuple))

return any(

not torch.utils._pytree.tree_is_leaf(x, is_leaf=is_list_or_tuple)

for x in itertools.chain(args, kwargs.values())

)

You can use tree_is_leaf after #113257 get merged into main.

XuehaiPan · 2025-03-06T16:29:39Z

torch/_ops.py

+        self.func = func
+        self.with_keyset = with_keyset
+
+    def __call__(self, *args, **kwargs):


Suggested change

def __call__(self, *args, **kwargs):

def __call__(self, /, *args, **kwargs):

Allow 'self' as a key in kwargs.

github-actions · 2025-05-16T15:36:24Z

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

yanboliang requested a review from XuehaiPan as a code owner February 26, 2025 05:37

pytorch-bot bot added ciflow/inductor module: dynamo labels Feb 26, 2025

yanboliang added the topic: not user facing topic category label Feb 26, 2025

yanboliang commented Feb 26, 2025

View reviewed changes

torch/_dynamo/variables/torch.py Show resolved Hide resolved

zou3519 reviewed Feb 26, 2025

View reviewed changes

torch/_dynamo/variables/torch.py Outdated Show resolved Hide resolved

zou3519 reviewed Feb 26, 2025

View reviewed changes

torch/_library/infer_schema.py Show resolved Hide resolved

zou3519 reviewed Feb 26, 2025

View reviewed changes

torch/_ops.py Outdated Show resolved Hide resolved

zou3519 reviewed Feb 26, 2025

View reviewed changes

torch/_dynamo/variables/torch.py Outdated Show resolved Hide resolved

[CustomOps] Support pytree input

cb34942

yanboliang force-pushed the custom1 branch from 35f93db to cb34942 Compare February 27, 2025 19:11

Update

cd1dc0c

zou3519 reviewed Feb 28, 2025

View reviewed changes

torch/_ops.py Outdated Show resolved Hide resolved

zou3519 reviewed Feb 28, 2025

View reviewed changes

torch/_ops.py Outdated Show resolved Hide resolved

yanboliang added 6 commits February 28, 2025 19:25

Update

858b0b0

Update

db58378

Update

8156410

Update

2f51c4a

Update

71fcd73

Update

d84db1a

yanboliang changed the title ~~[Prototype] Custom ops support arbitrary input types by migrating to python dispatcher~~ Custom ops support arbitrary input types by migrating to python dispatcher Mar 5, 2025

zou3519 reviewed Mar 5, 2025

View reviewed changes

yanboliang added 3 commits March 5, 2025 15:08

Update

5fc5fa2

Update

d883874

Update

14c0439

zou3519 reviewed Mar 6, 2025

View reviewed changes

XuehaiPan reviewed Mar 6, 2025

View reviewed changes

pytorchbot added the open source label Mar 14, 2025

zou3519 added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Mar 17, 2025

github-actions bot added the Stale label May 16, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Custom ops support arbitrary input types by migrating to python dispatcher #147927

Custom ops support arbitrary input types by migrating to python dispatcher #147927

		@@ -46,6 +47,23 @@ def dl_open_guard():
		sys.setdlopenflags(old_flags)


		class Kernel:


		def redispatch(self, /, keyset, args, *kwargs):
		return self._dispatch_in_python(args, kwargs, self._fallthrough_keys())

	def __call__(self, args, *kwargs):
	def __call__(self, /, args, *kwargs):

Custom ops support arbitrary input types by migrating to python dispatcher #147927

Are you sure you want to change the base?

Custom ops support arbitrary input types by migrating to python dispatcher #147927

Conversation

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/147927

✅ No Failures

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment