Skip to content

Add a miniature backend select implementation for prims #82311

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

Add a miniature backend select implementation for prims #82311

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e52a6c8
Add a miniature backend select implementation for prims
ezyang Jul 27, 2022
59d859c
Update on "Add a miniature backend select implementation for prims"
ezyang Jul 27, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion torch/_prims/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@

prim = torch.library.Library("prims", "DEF")
prim_impl = torch.library.Library("prims", "IMPL", "CompositeExplicitAutograd")
prim_backend_select_impl = torch.library.Library("prims", "IMPL", "BackendSelect")
prim_autograd_impl = torch.library.Library("prims", "IMPL", "Autograd")
prim_meta_impl = torch.library.Library("prims", "IMPL", "Meta")

Expand Down Expand Up @@ -453,14 +454,27 @@ def _autograd_impl(*args, **kwargs):
flat_args, args_spec = tree_flatten((args, kwargs))
return BackwardsNotSupported.apply(args_spec, *flat_args)

_meta_impl = _wrap_tensor_meta(meta)

def _backend_select_impl(*args, **kwargs):
if kwargs.get("device") and kwargs["device"].type == "meta":
return _meta_impl(*args, **kwargs)
else:
return _prim_impl(*args, **kwargs)

name = schema.split("(")[0]
prim_impl.impl(name, _prim_impl)
prim_autograd_impl.impl(name, _autograd_impl)
prim_meta_impl.impl(name, _wrap_tensor_meta(meta))
prim_meta_impl.impl(name, _meta_impl)

_prim_packet = getattr(torch.ops.prims, name)
_prim = _prim_packet.default

from torch._subclasses.fake_tensor import contains_tensor_types

if not any(contains_tensor_types(a.type) for a in _prim._schema.arguments):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this map 1-to-1 with the logic that codegen uses to determine when to generate backend select kernels?

pytorch/torchgen/gen.py

Line 859 in 1231194

def needs_backend_select(f: NativeFunction, selector: SelectiveBuilder) -> bool:

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hypothetically yes, but prims are simpler than native_functions.yaml so I guessed a simplified version would work.

prim_backend_select_impl.impl(name, _backend_select_impl)

for p in (_prim_packet, _prim):
p.__doc__ = doc
p.impl_nvfuser = impl_nvfuser # type: ignore[attr-defined]
Expand Down