frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[AOTI] Add missing ops to set of C-shim ops which can have nullptr returns (#158073)

Browse Source 
Most added ops are backwards ops, which have not been well-tested previously (thus why they were missed). Necessary ops were identified by manual examination of torch/_meta_registrations.py return values.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/158073
Approved by: https://github.com/desertfire
This commit is contained in:
Benjamin Glass
2025-07-10 22:28:27 +00:00
committed by PyTorch MergeBot
parent ff7dd1776f
commit f44a9eee47
1 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
torchgen/gen_aoti_c_shim.py
View File

@ -249,13 +249,19 @@ def gen_returns(schema: FunctionSchema) -> tuple[list[str], list[str]]:
ret_pointer_can_be_null = False
unambiguous_name = schema.name.unambiguous_name()
for name in [
"_scaled_dot_product_flash_attention",
"_scaled_dot_product_efficient_attention",
for name in (
"_functional_sym_constrain_range",
"_scaled_dot_product_cudnn_attention",
"_scaled_dot_product_efficient_attention_backward",
"_scaled_dot_product_efficient_attention",
"_scaled_dot_product_flash_attention",
"_scaled_dot_product_fused_attention_overrideable",
"_thhn_fused_lstm_cell_backward_impl",
"convolution_backward",
]:
"grid_sampler_2d_backward",
"grid_sampler_3d_backward",
"linear_backward",
):
if name in unambiguous_name:
ret_pointer_can_be_null = True
break