Make functorch codegen more robust (#693)

Previously the functorch codegen would only work if you didn't have a
PyTorch develop install in your environment. This PR changes it so that
the functorch codegen works when you have a PyTorch develop install in
the environment.

The reason for this change is that the PyTorch develop install adds a
`tools` module into the environment. It turns out we can just rely on
the tools module and this makes our codegen more robust to changes to
pytorch/pytorch codegen (when compared to what we were doing before,
which was keeping a copy of the PyTorch codegen inside of the functorch
repo).

Test Plan:
- wait for tests

Author: zou3519

codegen/gen.py

@@ -0,0 +1,70 @@
+import os
+import argparse
+import pathlib
+
+from tools.codegen.gen import FileManager, parse_native_yaml
+from gen_vmap_plumbing import gen_all_vmap_plumbing
+
+"""
+INSTRUCTIONS
+
+Step 1: You must have a PyTorch installation (in develop mode, i.e.
+installed with python setup.py develop) in your current environment.
+This script relies on the `tools` module from the PyTorch develop installation.
+
+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.
+"""
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description='functorch codegen')
+    parser.add_argument(
+        '-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)')
+    parser.add_argument(
+        '-d', '--install_dir', help='output directory',
+        default='functorch/csrc')
+    options = parser.parse_args()
+
+    native_yaml_path = os.path.join(options.source_path, 'native/native_functions.yaml')
+    parsed_yaml = parse_native_yaml(native_yaml_path)
+    native_functions, _ = parsed_yaml.native_functions, parsed_yaml.backend_indices
+    template_dir = os.path.join(options.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)
+
+    cpu_fm = make_file_manager(options.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()

tools/codegen/gen_vmap_plumbing.py renamed to codegen/gen_vmap_plumbing.py

@@ -2,14 +2,14 @@
     DispatcherSignature,
 )
 from tools.codegen.model import (
-    BaseTy, Variant, OptionalType, BaseType, ListType, NativeFunction, Type,
+    BaseTy, OptionalType, BaseType, ListType, NativeFunction, Type,
     Argument, Return, SchemaKind, Tag
 )
 from tools.codegen.api.translate import translate
 from tools.codegen.context import method_with_native_function
 from tools.codegen.utils import mapMaybe
 from dataclasses import dataclass
-from typing import List, Dict, Optional, Tuple, Set, Any, Union, Sequence, TypeVar
+from typing import List, Optional, Tuple
 import textwrap
 
 

functorch/csrc/VmapGeneratedPlumbing.h

@@ -597,6 +597,7 @@ class Support(enum.Enum):
 }
 
 JVP_EXEMPTIONS = {
+    'nn.functional.dropout',  # not actually problem, randomness testing artifact
     'nn.functional.dropout2d',  # not actually problem, randomness testing artifact
     'nn.functional.rrelu',  # not actually problem, randomness testing artifact
     # 'normal',
@@ -808,10 +809,10 @@ def summary(self):
 
 print("=" * 30 + " Top 60 Summary " + "=" * 30)
 opset = OperatorSet.from_top_ops_threshold(35, 25)
-result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
-pprint.pprint(result)
-result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN))
-pprint.pprint(result)
+# result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
+# pprint.pprint(result)
+# result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN))
+# pprint.pprint(result)
 # kresult = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
 # kpprint.pprint(result)
 # result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
@@ -823,12 +824,14 @@ def summary(self):
 
 print("=" * 30 + " Top 125 Summary " + "=" * 30)
 opset = OperatorSet.from_top125()
-result = opset.query(Operator.supports_vmap, (Support.NO, Support.UNKNOWN))
-pprint.pprint(result)
-# kresult = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
-# kpprint.pprint(result)
-# result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
+# result = opset.query(Operator.supports_vmap, (Support.NO, Support.UNKNOWN))
+# pprint.pprint(result)
+# result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
 # pprint.pprint(result)
+result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN))
+pprint.pprint(result)
+result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
+pprint.pprint(result)
 # result = opset.query(Operator.supports_fast_vmapjvp, (Support.NO, Support.UNKNOWN))
 # pprint.pprint(result)
 # pprint.pprint(result)

test/discover_coverage.py

@@ -381,8 +381,6 @@ def wrapped_fn(*args, **kwargs):
         xfail('std_mean'),
         # https://gist.github.com/zou3519/f62a167fb46cda01d7f238f61dd9ccf9
         xfail('linalg.eigvalsh'),
-        # https://gist.github.com/zou3519/b86616d01ca375a4bd17403277f49225
-        xfail('nn.functional.dropout', device_type='cuda'),
 
         # =============================================
         # NB: The above failures also fail using PyTorch core's