Update vmap codegen based on changes to torchgen (#898)

torchgen was changed sometime in the last month. This PR updates our
codegen to work with the new changes.

Coming soon:
- we should be able to autogenerate our codegen and remove the checked
in version.

Author: zou3519

codegen/gen.py

@@ -2,7 +2,9 @@
 import argparse
 import pathlib
 
+import torchgen
 from torchgen.gen import FileManager, parse_native_yaml
+from torchgen.gen import get_torchgen_root
 from gen_vmap_plumbing import gen_all_vmap_plumbing
 
 """
@@ -14,12 +16,7 @@
 
 Step 2: Run this script.
 
-# Replace the last argument with your path to native_functions.yaml
-python codegen/gen.py -s /scratch/rzou/pt/debug-cpu/aten/src/ATen
-
-NB: PyTorch's `tools` module is a giant hack (it somehow gets installed into your
-environment when one does python setup.py develop), but it's highly likely that
-PyTorch won't change it anytime soon because it's very messy to modify.
+python codegen/gen.py
 """
 
 
@@ -29,42 +26,31 @@ def main() -> None:
         '-s',
         '--source-path',
         help='path to source directory for ATen',
-        default='/scratch/rzou/pt/debug-cpu/aten/src/ATen')
-    parser.add_argument(
-        '-o',
-        '--output-dependencies',
-        help='output a list of dependencies into the given file and exit')
-    parser.add_argument(
-        '--dry-run', action='store_true',
-        help='run without writing any files (still updates outputs)')
+        default=None)
     parser.add_argument(
         '-d', '--install_dir', help='output directory',
         default='functorch/csrc')
     options = parser.parse_args()
+    generate_code(options.install_dir, options.source_path)
+
 
-    native_yaml_path = os.path.join(options.source_path, 'native/native_functions.yaml')
-    parsed_yaml = parse_native_yaml(native_yaml_path)
+def generate_code(install_dir='functorch/csrc', source_path=None):
+    if source_path is None:
+        # infer the source path via torchgen
+        source_path = os.path.join(get_torchgen_root(), "packaged/ATen")
+
+    native_yaml_path = os.path.join(source_path, 'native/native_functions.yaml')
+    tags_path = os.path.join(source_path, 'native/tags.yaml')
+    parsed_yaml = parse_native_yaml(native_yaml_path, tags_path)
     native_functions, _ = parsed_yaml.native_functions, parsed_yaml.backend_indices
-    template_dir = os.path.join(options.source_path, "templates")
+    template_dir = os.path.join(source_path, "templates")
 
     def make_file_manager(install_dir: str) -> FileManager:
-        return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=options.dry_run)
+        return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=False)
 
-    cpu_fm = make_file_manager(options.install_dir)
+    cpu_fm = make_file_manager(install_dir)
     cpu_fm.write('VmapGeneratedPlumbing.h', lambda: gen_all_vmap_plumbing(native_functions))
 
-    if options.output_dependencies:
-        depfile_path = pathlib.Path(options.output_dependencies).resolve()
-        depfile_name = depfile_path.name
-        depfile_stem = depfile_path.stem
-
-        for fm, prefix in [
-                (cpu_fm, ""),
-        ]:
-            varname = prefix + depfile_stem
-            path = depfile_path.parent / (prefix + depfile_name)
-            fm.write_outputs(varname, str(path))
-
 
 if __name__ == '__main__':
     main()

codegen/gen_vmap_plumbing.py

@@ -3,7 +3,7 @@
 )
 from torchgen.model import (
     BaseTy, OptionalType, BaseType, ListType, NativeFunction, Type,
-    Argument, Return, SchemaKind, Tag
+    Argument, Return, SchemaKind
 )
 from torchgen.api.translate import translate
 from torchgen.context import method_with_native_function
@@ -198,7 +198,7 @@ def gen_vmap_plumbing(native_function: NativeFunction) -> str:
     if not accepts_at_least_one_tensor_input(schema):
         return None
     # in-place views need special handling
-    if native_function.tag == Tag.inplace_view:
+    if 'inplace_view' in native_function.tags:
         return None
 
     if schema.kind() == SchemaKind.inplace: