Make torch.library.triton_op public (#141880)

We've been using it privately for half a year and everything's been
good. This PR:
1. Makes torch.library.triton_op public
2. Renames capture_triton -> wrap_triton. We got feedback that no one
   knew what "capture triton" does.
3. Makes torch.library.wrap_triton public.

triton_op is used to construct a Python custom operator that may call 1+
triton kernels. Each of those triton kernels must be annotated with
wrap_triton.

Test Plan:
- existing tests

Pull Request resolved: https://github.com/pytorch/pytorch/pull/141880
Approved by: https://github.com/albanD
ghstack dependencies: #141894
This commit is contained in:
rzou
2024-12-02 13:22:01 -08:00
committed by PyTorch MergeBot
parent ac600fdce6
commit 827c322290
5 changed files with 68 additions and 40 deletions

View File

@ -18,6 +18,7 @@ from torch._library.custom_ops import (
device_types_t,
)
from torch._library.infer_schema import infer_schema # noqa: F401
from torch._library.triton import triton_op, wrap_triton
from torch._ops import OpOverload
@ -32,6 +33,8 @@ __all__ = [
"register_vmap",
"get_ctx",
"custom_op",
"triton_op",
"wrap_triton",
"infer_schema",
]