Enable AOTI shim v2 build and add into libtorch (#125211)

Summary:
Follow up of https://github.com/pytorch/pytorch/pull/125087

This diff will create shim v2 header and cpp file and corresponding build

Differential Revision: D56617546

Pull Request resolved: https://github.com/pytorch/pytorch/pull/125211
Approved by: https://github.com/desertfire
This commit is contained in:
Huamin Li 2024-05-31 23:56:11 +00:00 committed by PyTorch MergeBot
parent a8c9b26534
commit ff8042bcfb
2 changed files with 9 additions and 0 deletions

View File

@ -383,6 +383,7 @@ def get_aten_generated_files(enabled_backends):
"core/TensorMethods.cpp",
"core/aten_interned_strings.h",
"core/enum_tag.h",
"torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.cpp",
] + get_aten_derived_type_srcs(enabled_backends)
# This is tiresome. A better strategy would be to unconditionally
@ -467,6 +468,7 @@ def gen_aten_files(
cmd = "$(exe {}torchgen:gen) ".format(ROOT_PATH) + " ".join([
"--source-path $(location {}:aten_src_path)/aten/src/ATen".format(ROOT),
"--install_dir $OUT",
"--aoti_install_dir $OUT/torch/csrc/inductor/aoti_torch/generated"
] + extra_params),
visibility = visibility,
compatible_with = compatible_with,

View File

@ -73,6 +73,7 @@ def define_targets(rules):
"$(execpath //torchgen:gen)",
"--install_dir=$(RULEDIR)",
"--source-path aten/src/ATen",
"--aoti_install_dir=$(RULEDIR)/torch/csrc/inductor/aoti_torch/generated"
] + (["--static_dispatch_backend CPU"] if rules.is_cpu_static_dispatch_build() else []))
gen_aten_outs_cuda = (
@ -83,6 +84,7 @@ def define_targets(rules):
gen_aten_outs = (
GENERATED_H + GENERATED_H_CORE +
GENERATED_CPP + GENERATED_CPP_CORE +
GENERATED_AOTI_CPP +
aten_ufunc_generated_cpu_sources() +
aten_ufunc_generated_cpu_kernel_sources() + [
"Declarations.yaml",
@ -316,3 +318,8 @@ GENERATED_AUTOGRAD_CPP = [
"torch/csrc/lazy/generated/RegisterAutogradLazy.cpp",
"torch/csrc/lazy/generated/RegisterLazy.cpp",
] + _GENERATED_AUTOGRAD_CPP_HEADERS + GENERATED_LAZY_H
GENERATED_AOTI_CPP = [
"torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.cpp",
"torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.cpp",
]