Let custom ops registration code only include ATen headers

larryliu0820 · facebook-github-bot · commit 5f8aa6db8a3f · 2023-08-12T16:47:17.000-07:00
Summary: X-link: pytorch/pytorch#107064 Basically we generate `CustomOpsNativeFunctions.h` for registering custom ops into PyTorch JIT runtime. This header needs to hookup with the C++ kernel implementation of all the custom ops. For this reason it should include ATen headers instead of Executorch headers. This PR changes it. Reviewed By: kirklandsign Differential Revision: D48282828 fbshipit-source-id: cf4ed96a7fb899b85a71d7efbc3fcd6c5a38632d
diff --git a/codegen/templates/NativeFunctions.h b/codegen/templates/NativeFunctions.h
@@ -11,9 +11,7 @@
 
 #include <tuple>
 
-#include <executorch/runtime/core/exec_aten/exec_aten.h> // at::Tensor etc.
-#include <executorch/codegen/macros.h> // TORCH_API
-#include <executorch/runtime/kernel/kernel_runtime_context.h>
+${headers}
 
 // ${generated_comment}
 
diff --git a/shim/xplat/executorch/codegen/codegen.bzl b/shim/xplat/executorch/codegen/codegen.bzl
@@ -258,13 +258,7 @@ def exir_custom_ops_aot_lib(
             # @lint-ignore BUCKLINT link_whole
             link_whole = True,
             visibility = visibility,
-            deps = [
-                "//executorch/runtime/core/exec_aten:lib_aten",
-                "//executorch/codegen:macros",
-            ] + kernels + deps,
-            exported_deps = [
-                "//executorch/runtime/kernel:kernel_runtime_context_aten",
-            ],
+            deps = kernels + deps,
             external_deps = ["libtorch"],
             define_static_target = define_static_target,
             # Relax visibility restrictions since deps may include targets
