[nativert] Move auto_functionalize_kernel (#156454)

Summary:
Torch Native Runtime RFC: https://github.com/pytorch/rfcs/pull/72

As part of the effort to open source TorchNativeRuntime (or what we call Sigmoid), we are moving the Pytree implementation to torch/:

fbcode/sigmoid/kernels -> fbcode/caffe2/torch/nativert/kernels

Copied from original auto_functionalize Diff Summary D53776805:

This is a non-functional kernel implementation for auto_functionalize

In AutoFunctionalizeKernel, I directly call the underlying target without making a clone of mutating inputs.

This would mutates the input tensors inplace, which is unsafe in general.

However, Sigmoid is not doing any graph optimization, or node reordering at the moment, so it's ok do take this short cut.

In the proper functional implementation, it will

make a clone of the mutating input tensor

return these new instance of tensors as AutoFunctionalizeKernel output.

If the original exported program has some "bufferMutation" or "userInputMutation" fields, it will also need to honor such mutations in Sigmoid.

Test Plan: See internal for test plan

Differential Revision: D76926383

Pull Request resolved: https://github.com/pytorch/pytorch/pull/156454
Approved by: https://github.com/zhxchen17
This commit is contained in:
Shangdi Yu
2025-06-20 19:53:12 +00:00
committed by PyTorch MergeBot
parent eb331b59fe
commit e5ea24fb27
3 changed files with 94 additions and 0 deletions

View File

@ -608,6 +608,7 @@ libtorch_nativert_sources = [
"torch/nativert/common/FileUtil.cpp",
"torch/nativert/detail/ITree.cpp",
"torch/nativert/kernels/C10Kernel.cpp",
"torch/nativert/kernels/AutoFunctionalizeKernel.cpp",
]
torch_mobile_tracer_sources = [