benchmark

c960e717 - Add a user-kernel HOP registry and register Triton with it

Add a user-kernel HOP registry and register Triton with it Summary: X-link: https://github.com/pytorch/pytorch/pull/203473 Seventeen places outside triton_kernel_wrap.py recognize a user kernel by comparing a node's target with Triton's HOPs, so a second kernel DSL's HOPs would be skipped, or would crash, at each of them. This PR adds a registry of user-kernel HOP pairs and registers Triton with its current behavior; the next PR, #202442, routes those sites through the registry. Only this PR's tests read the registry, and reinplacing, the one pass it touches, calls Triton's kernel-name function instead of the same code inline, so no pass behaves differently. - `torch/_higher_order_ops/user_kernel_hop.py`: each DSL registers a `UserKernelHOPInfo` with its mutation/functional HOP pair, a getter for the kernel's write set, a kernel name for logs, a getter for the layout constraint Inductor gives the mutation HOP's tensor arguments, the `ReInplaceTrigger` member that keys its reinplacing counters, and `serializable` (whether torch.export can write its nodes). The class docstring states the HOP kwargs layout the passes rely on. `register_user_kernel_hop` checks everything before writing anything (HOP types, the functional HOP's CompositeExplicitAutograd py_impl, two distinct HOPs that no other info registered, and a trigger that is `USER_KERNEL_OPS` or one `TRIGGER_DSLS` gives to this DSL, never `AUTO_FUNC_V1` or `AUTO_FUNC_V2`), so a refused registration leaves the registry unchanged. A lock is held across the ownership checks and both writes, and every lookup takes it, so DSLs registering from different threads cannot both claim a HOP and no reader sees half of a pair. Infos compare by identity (`eq=False`), so a DSL's callables need not be hashable. The sites will ask `get_user_kernel_hop_info` (optionally for one role), `is_user_kernel_mutation_hop`, `is_user_kernel_functional_hop` or `user_kernel_hop_infos`. - `torch/_higher_order_ops/triton_kernel_wrap.py`: Triton's registration (`TRITON_OPS`, `serializable=True`) and `get_triton_layout_constraint`, which checks `triton_kernel_default_layout_constraint` and, for an unknown value, raises the error graph.py raises today. The reinplacing pass's kernel-name code moves here as `get_triton_kernel_name`. - `ReInplaceTrigger` moves from `torch._dynamo.utils`, which re-exports it, to the new `torch/_higher_order_ops/reinplace_trigger.py`, which adds `USER_KERNEL_OPS` (for DSLs without a member of their own) and `TRIGGER_DSLS`. Importing triton_kernel_wrap still imports neither `torch._dynamo` nor `torch._inductor`. `triton_kernel_default_layout_constraint` stays Triton's: every other DSL returns its own constraint from its getter (the CuTeDSL registration in #203538 asks for `needs_fixed_stride_order`). The FX graph cache still keys only Triton kernels by source, so a DSL with cacheable HOPs must extend the cache key first; the CuTeDSL HOPs (#203538) are not cacheable. This PR was split out of #202442 to keep each PR under the 1500-line limit. ### Test plan New CPU tests in test/inductor/test_triton_kernels.py cover the lookups, each refused registration (which must leave the registry unchanged), the registry lock and Triton's layout-constraint getter, with a second HOP pair defined in the test file; test_not_import_dynamo_from_user_kernel_hops in test/test_testing.py checks the import claim above. ``` python test/inductor/test_triton_kernels.py -k user_kernel_hop \ -k layout_constraint_config python test/test_testing.py -k test_not_import_dynamo_from_user_kernel_hops ``` The split, the review fixes and this description were drafted with an AI assistant. ghstack-source-id: 466260746 exported-using-ghexport Reviewed By: desertfire Differential Revision: D129753992 fbshipit-source-id: f0522dc1f3c886560dc4a88245e08209dccdbdcd