Make lazy codegen honor per-operator-headers flag (#74450)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74450

- per-operator-headers is a strict build mode where compulation units aren't allowed
to depend on bulk headers like ATen/Functions.h, but must instead depend only on the
specific operator headers used.  (In other configurations, the reverse is required).

Test Plan: CI to make sure nothing breaks for existing backends, and rebased next diff manual test to make sure it actually helps

Reviewed By: ezyang, bdhirsh

Differential Revision: D35002666

fbshipit-source-id: 712445f8d146cf026759444fbd42a20705be9bef
(cherry picked from commit f13e5522d49a6edcb6aed4431b1ec8e2b50a98fc)

tools/setup_helpers/generate_code.py

@@ -167,6 +167,11 @@ def main() -> None:
         action='store_true',
         help='Enable generation of the torch::lazy TorchScript backend'
     )
+    parser.add_argument(
+        '--per_operator_headers',
+        action='store_true',
+        help='Build lazy tensor ts backend with per-operator ATen headers, must match how ATen was built'
+    )
     options = parser.parse_args()
 
     generate_code(
@@ -201,7 +206,8 @@ def main() -> None:
                             impl_path=None,
                             gen_ts_lowerings=True,
                             node_base="TsNode",
-                            node_base_hdr="torch/csrc/lazy/ts_backend/ts_node.h")
+                            node_base_hdr="torch/csrc/lazy/ts_backend/ts_node.h",
+                            per_operator_headers=options.per_operator_headers)
 
 
 if __name__ == "__main__":