mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
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:
@ -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",
|
||||
]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user