Skip to content

Python runtime API for operators #10726

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.

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

Open
3 tasks
larryliu0820 opened this issue May 6, 2025 · 0 comments
Open
3 tasks

Python runtime API for operators #10726

larryliu0820 opened this issue May 6, 2025 · 0 comments
Labels
module: runtime Issues related to the core runtime and code under runtime/ triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Comments

@larryliu0820
Copy link
Contributor

larryliu0820 commented May 6, 2025

🚀 The feature, motivation and pitch

Goal

Allow users to directly call kernels (potentially delegates) in python runtime. Supports dynamic kernel registration/deregistration, kernel metadata (op schema, selected dtype/dim order) inspection.

Motivation

Simplify kernel development workflow.
Leverage existing PyTorch op unit test framework for kernel coverage.
Microbenchmarks for kernels.

API

import torch
from typing import Callable
from executorch.runtime import Verification, Runtime, Program, Method

et_runtime: Runtime = Runtime.get()

# New API to retrieve kernel
op: Callable = et_runtime.operator_registry.aten.add.out

# Can also do:
op: Callable = et_runtime.operator_registry.get_kernel("aten::add.out")

a = torch.ones([2, 2])
b = torch.ones([2, 2])
c = torch.empty_like(a)

op(a, b, out=c) # calls into ExecuTorch kernel instead of PyTorch kernel.

# Print out schema
print(op._schema)

# No such kernel registered
et_runtime.operator_registry.aten.add.Tensor # Throws exception

Dynamically register a custom kernel

import torch
from executorch.runtime import Verification, Runtime, Program, Method

et_runtime: Runtime = Runtime.get()

# Register a custom kernel
op: Callable = et_runtime.operator_registry.register_kernel(
	"aten::add.out",
	"<kernel string>",
	kernel_key,
)

Task Breakdown

  • Add pybindings for OperatorRegistry and Kernel. For Kernel pybind, it should expose an operator() in python, that takes in a list of arguments and inside the pybind implementation we wrap them with EValues and call the kernel.
  • Connect Kernel with EdgeOpOverload. EdgeOpOverload is an AOT operator concept and if we call it directly it will trigger the ATen kernel. Add another API in EdgeOpOverload so that it can call the Kernel in ET.
  • Expose custom kernel registration API in python. To do this we need to leverage Ninja to compile the custom kernel code and register it. pytorch/pytorch has existing tools to do that.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @JacobSzwejbka @lucylq

@github-project-automation github-project-automation bot moved this to To triage in ExecuTorch Core May 6, 2025
@Jack-Khuu Jack-Khuu added module: runtime Issues related to the core runtime and code under runtime/ triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module labels May 9, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
module: runtime Issues related to the core runtime and code under runtime/ triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
Status: To triage
Development

No branches or pull requests

2 participants