From ff8042bcfb518127c86ad5b4af4fa9171a499904 Mon Sep 17 00:00:00 2001
From: Huamin Li <huaminli@meta.com>
Date: Fri, 31 May 2024 23:56:11 +0000
Subject: [PATCH] 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
---
 buckbuild.bzl | 2 ++
 build.bzl     | 7 +++++++
 2 files changed, 9 insertions(+)

diff --git a/buckbuild.bzl b/buckbuild.bzl
index 649ebe66836..1d668117e91 100644
--- a/buckbuild.bzl
+++ b/buckbuild.bzl
@@ -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,
diff --git a/build.bzl b/build.bzl
index 5ab9f92acec..8fd15f4e9c4 100644
--- a/build.bzl
+++ b/build.bzl
@@ -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",
+]