8000 [inductor] remove hardcoded mapping to resolve ops from ExternKernelSchedulerNode · Issue #147137 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[inductor] remove hardcoded mapping to resolve ops from ExternKernelSchedulerNode #147137

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.

Sign up for GitHub

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

Jump to bottom
Open
xmfan opened this issue Feb 13, 2025 · 2 comments
Open

[inductor] remove hardcoded mapping to resolve ops from ExternKernelSchedulerNode #147137

xmfan opened this issue Feb 13, 2025 · 2 comments
Labels
module: inductor oncall: pt2 triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Comments

@xmfan
Copy link
Member
xmfan commented Feb 13, 2025 •
edited by pytorch-bot bot
Loading

🐛 Describe the bug

https://github.com/pytorch/pytorch/pull/146992/files#r1953070064

during runtime estimation, we use this reverse map to lookup ops contained in ExternKernelSchedulerNode

pytorch/torch/_inductor/scheduler.py

Lines 924 to 930 in b0553ce

# TODO(xmfan): reuse: an existing mapping for this if it exists, or formalize this into ir.py:ExternKernel
kernel_name_to_op = {
"extern_kernels.convolution": torch.ops.aten.convolution,
"extern_kernels.mm": torch.ops.aten.mm,
"extern_kernels.bmm": torch.ops.aten.bmm,
"extern_kernels.addmm": torch.ops.aten.addmm,
}

i've been wanting to find an existing registration system, per @eellison one alternative can be FlopCounter registrations

Versions

main

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @amjames @aakhundov
@xmfan xmfan added module: inductor oncall: pt2 triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module labels Feb 13, 2025
@eellison
8000 Copy link
Contributor

Just to clarify: we do already look up the FlopCounter registration just by virtue of running it. And we should be able to look up ExternKernel -> op mappings when we have op_overload specified on the ExternKernel:

pytorch/torch/_inductor/ir.py

Lines 4784 to 4786 in cdbf677

op_overload: Optional[
Union[torch._ops.OpOverload, torch._ops.HigherOrderOperator]
] = None

@xmfan
Copy link
Member Author
xmfan commented Feb 25, 2025 •
edited
Loading

Oh this runtime estimation can be ran without FlopCounter. Looks like the op_overload is empty:

ExternKernelOut(name='buf0',
    layout=FixedLayout('cuda:0', torch.float32, size=[10, 10], stride=[10, 1]),
    inputs=[InputBuffer(name='arg1_1',
                        layout=FixedLayout('cuda:0', torch.float32, size=[10, 10], stride=[10, 1])),
            InputBuffer(name='arg0_1',
                        layout=FixedLayout('cuda:0', torch.float32, size=[10, 10], stride=[10, 1]))],
    constant_args=(),
    kwargs={},
    output_view=None,
    python_kernel_name='extern_kernels.mm',
    cpp_kernel_name='at::mm_out',
    ordered_kwargs_for_cpp_kernel=(),
    op_overload=None,
    arg_properties=[{}, {}],
    kwarg_properties=None,
    unbacked_bindings={},
    mutation_outputs=[])

@eellison eellison mentioned this issue May 2, 2025
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
module: inductor oncall: pt2 triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@xmfan @eellison
0