# User-defined Kernels vs. \`torch.library\` custom op

**URL:** <https://dev-discuss.pytorch.org/t/user-defined-kernels-vs-torch-library-custom-op/2113>\
**Category:** compiler\
**Created:** [May 22, 2024, 1:26pm UTC](https://dev-discuss.pytorch.org/t/user-defined-kernels-vs-torch-library-custom-op/2113 "2024-05-22T13:26:07Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![jeromeku](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jeromeku/32/925_2.png) [@jeromeku](https://dev-discuss.pytorch.org/u/jeromeku)\
**Post date:** [May 22, 2024, 1:26pm UTC](https://dev-discuss.pytorch.org/t/user-defined-kernels-vs-torch-library-custom-op/2113/1 "2024-05-22T13:26:07Z")

</div>

If I have a user-defined `triton` kernel as part of a torch.nn.Module, what is the difference between:

1. Registering the user-defined `triton` kernel using the `torch.library` API and compiling the module  
vs.
2. Calling `torch.compile` on the module and relying on the `inductor` backend to inline the user-defined kernel?

- What are the performance differences to be expected between the two approaches? Are there guidelines for when to one API over the other?
- Will registering a custom op overcome the current limitations of user-defined kernels per the [official tutorial](https://pytorch.org/tutorials/recipes/torch_compile_user_defined_triton_kernel_tutorial.html#composibility-and-limitations)? Specifically, often need to define heuristics _after_ autotune.
- More generally, are there sharp bits to watch out for when compiling a function / module which includes a user-defined kernel that can result in worse performance than calling the function / module without compilation?

@oulgen @zou3519

---

<div class="post-metadata">

**Author:** ![oulgen](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/oulgen/32/1653_2.png) [@oulgen](https://dev-discuss.pytorch.org/u/oulgen)\
**Post date:** [May 22, 2024, 4:47pm UTC](https://dev-discuss.pytorch.org/t/user-defined-kernels-vs-torch-library-custom-op/2113/2 "2024-05-22T16:47:11Z")

</div>

The rule of thumb for picking between the raw triton kernels versus through custom ops is generally about composibility requirements. If your code needs to be composible with dtensor or other tensor subclasses, use custom ops. Otherwise, use raw triton kernels. The raw version provides a simpler user experience as there’s no registration etc, code that works without compile works with compile.

In terms of performance, dynamo/inductor have optimization passes that optimize the code around the triton kernel so you’ll likely see the equal or better performance using raw versions.