[JIT] Make aot autograd decompositions usable in JIT, add script for serializing the decompositions #73938

eellison · 2022-03-08T20:32:57Z

Stack from ghstack:

This is a first step in porting and making usable all of the decompositions defined in functorch in core and in JIT as well as C++.

The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run tools/codegen/decompositions/gen_jit_decompositions.py.

Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like aten.var (shown here in the example).

This is just a first PR, i'm sure there will be many follows ups such as:

making these runnable in C++ with simple executor
porting over more decompositions from AOT Autograd
Using opinfos / more robust testing
Categorizing decompositions
Hooking in decompositions at various points of JIT execution

Differential Revision: D34938126

…izing the decompositions [ghstack-poisoned]

pytorch-bot · 2022-03-08T20:33:01Z

CI Flow Status

⚛️ CI Flow

Ruleset - Version: v1
Ruleset - File: https://github.com/pytorch/pytorch/blob/7771043abf7f559f55c272b821118c37272cfed9/.github/generated-ciflow-ruleset.json
PR ciflow labels: ciflow/default
Add ciflow labels to this PR to trigger more builds:

Workflows	Labels (bold enabled)	Status
Triggered Workflows
linux-binary-conda	`ciflow/binaries`, `ciflow/binaries_conda`, `ciflow/default`	✅ triggered
linux-binary-libtorch-cxx11-abi	`ciflow/binaries`, `ciflow/binaries_libtorch`, `ciflow/default`	✅ triggered
linux-binary-libtorch-pre-cxx11	`ciflow/binaries`, `ciflow/binaries_libtorch`, `ciflow/default`	✅ triggered
linux-binary-manywheel	`ciflow/binaries`, `ciflow/binaries_wheel`, `ciflow/default`	✅ triggered
linux-bionic-py3.7-clang9	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/noarch`, `ciflow/trunk`	✅ triggered
linux-bionic-rocm4.5-py3.7	`ciflow/all`, `ciflow/default`, `ciflow/linux`, `ciflow/rocm`, `ciflow/trunk`	✅ triggered
linux-docs	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/docs`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
linux-vulkan-bionic-py3.7-clang9	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`, `ciflow/vulkan`	✅ triggered
linux-xenial-cuda11.3-py3.7-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
linux-xenial-cuda11.3-py3.7-gcc7-bazel-test	`ciflow/all`, `ciflow/bazel`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
linux-xenial-py3-clang5-mobile-build	`ciflow/all`, `ciflow/default`, `ciflow/linux`, `ciflow/mobile`, `ciflow/trunk`	✅ triggered
linux-xenial-py3-clang5-mobile-custom-build-static	`ciflow/all`, `ciflow/default`, `ciflow/linux`, `ciflow/mobile`, `ciflow/trunk`	✅ triggered
linux-xenial-py3.7-clang7-asan	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/sanitizers`, `ciflow/trunk`	✅ triggered
linux-xenial-py3.7-clang7-onnx	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/onnx`, `ciflow/trunk`	✅ triggered
linux-xenial-py3.7-gcc5.4	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
linux-xenial-py3.7-gcc7	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
linux-xenial-py3.7-gcc7-no-ops	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
macos-arm64-binary-conda	`ciflow/binaries`, `ciflow/binaries_conda`, `ciflow/default`	✅ triggered
macos-arm64-binary-wheel	`ciflow/binaries`, `ciflow/binaries_wheel`, `ciflow/default`	✅ triggered
macos-binary-conda	`ciflow/binaries`, `ciflow/binaries_conda`, `ciflow/default`	✅ triggered
macos-binary-libtorch-cxx11-abi	`ciflow/binaries`, `ciflow/binaries_libtorch`, `ciflow/default`	✅ triggered
macos-binary-libtorch-pre-cxx11	`ciflow/binaries`, `ciflow/binaries_libtorch`, `ciflow/default`	✅ triggered
macos-binary-wheel	`ciflow/binaries`, `ciflow/binaries_wheel`, `ciflow/default`	✅ triggered
pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-custom-build-single	`ciflow/all`, `ciflow/android`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
pytorch-linux-xenial-py3-clang5-android-ndk-r19c-gradle-custom-build-single-full-jit	`ciflow/all`, `ciflow/android`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/trunk`	✅ triggered
win-vs2019-cpu-py3	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/trunk`, `ciflow/win`	✅ triggered
win-vs2019-cuda11.3-py3	`ciflow/all`, `ciflow/cuda`, `ciflow/default`, `ciflow/trunk`, `ciflow/win`	✅ triggered
windows-binary-libtorch-cxx11-abi	`ciflow/binaries`, `ciflow/binaries_libtorch`, `ciflow/default`	✅ triggered
windows-binary-libtorch-pre-cxx11	`ciflow/binaries`, `ciflow/binaries_libtorch`, `ciflow/default`	✅ triggered
windows-binary-wheel	`ciflow/binaries`, `ciflow/binaries_wheel`, `ciflow/default`	✅ triggered
Skipped Workflows
caffe2-linux-xenial-py3.7-gcc5.4	`ciflow/all`, `ciflow/cpu`, `ciflow/linux`, `ciflow/trunk`	🚫 skipped
docker-builds	`ciflow/all`, `ciflow/trunk`	🚫 skipped
ios-12-5-1-arm64	`ciflow/all`, `ciflow/ios`, `ciflow/macos`, `ciflow/scheduled`	🚫 skipped
ios-12-5-1-arm64-coreml	`ciflow/all`, `ciflow/ios`, `ciflow/macos`, `ciflow/scheduled`	🚫 skipped
ios-12-5-1-arm64-custom-ops	`ciflow/all`, `ciflow/ios`, `ciflow/macos`, `ciflow/scheduled`	🚫 skipped
ios-12-5-1-arm64-metal	`ciflow/all`, `ciflow/ios`, `ciflow/macos`, `ciflow/scheduled`	🚫 skipped
ios-12-5-1-x86-64	`ciflow/all`, `ciflow/ios`, `ciflow/macos`, `ciflow/trunk`	🚫 skipped
ios-12-5-1-x86-64-coreml	`ciflow/all`, `ciflow/ios`, `ciflow/macos`, `ciflow/trunk`	🚫 skipped
libtorch-linux-xenial-cuda10.2-py3.7-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/libtorch`, `ciflow/linux`, `ciflow/trunk`	🚫 skipped
libtorch-linux-xenial-cuda11.3-py3.7-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/libtorch`, `ciflow/linux`, `ciflow/trunk`	🚫 skipped
linux-bionic-cuda10.2-py3.9-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/slow`, `ciflow/trunk`	🚫 skipped
linux-docs-push	`ciflow/all`, `ciflow/cpu`, `ciflow/linux`, `ciflow/scheduled`	🚫 skipped
linux-xenial-cuda11.3-py3.7-gcc7-no-ops	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/trunk`	🚫 skipped
macos-10-15-py3-arm64	`ciflow/all`, `ciflow/macos`, `ciflow/trunk`	🚫 skipped
macos-10-15-py3-lite-interpreter-x86-64	`ciflow/all`, `ciflow/macos`, `ciflow/trunk`	🚫 skipped
macos-11-py3-x86-64	`ciflow/all`, `ciflow/macos`, `ciflow/trunk`	🚫 skipped
parallelnative-linux-xenial-py3.7-gcc5.4	`ciflow/all`, `ciflow/cpu`, `ciflow/linux`, `ciflow/trunk`	🚫 skipped
periodic-libtorch-linux-bionic-cuda11.5-py3.7-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/libtorch`, `ciflow/linux`, `ciflow/scheduled`	🚫 skipped
periodic-linux-bionic-cuda11.5-py3.7-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/scheduled`	🚫 skipped
periodic-linux-xenial-cuda10.2-py3-gcc7-slow-gradcheck	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/scheduled`, `ciflow/slow`, `ciflow/slow-gradcheck`	🚫 skipped
periodic-linux-xenial-cuda11.3-py3.7-gcc7-debug	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/scheduled`	🚫 skipped
periodic-win-vs2019-cuda11.5-py3	`ciflow/all`, `ciflow/cuda`, `ciflow/scheduled`, `ciflow/win`	🚫 skipped
pytorch-linux-xenial-py3-clang5-android-ndk-r19c-build	`ciflow/all`, `ciflow/android`, `ciflow/cpu`, `ciflow/linux`, `ciflow/trunk`	🚫 skipped
pytorch-xla-linux-bionic-py3.7-clang8	`ciflow/all`, `ciflow/cpu`, `ciflow/linux`, `ciflow/trunk`, `ciflow/xla`	🚫 skipped

facebook-github-bot · 2022-03-08T20:33:04Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/73938
Need help or want to give feedback on the CI? Visit our office hours

💊 CI failures summary and remediations

As of commit dcfa896 (more details on the Dr. CI page):

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution [ghstack-poisoned]

…izing the decompositions ghstack-source-id: 5dbf857 Pull Request resolved: #73938

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution [ghstack-poisoned]

…izing the decompositions ghstack-source-id: 6348f49 Pull Request resolved: #73938

Chillee

LGTM - modulo some of the comments we discussed offline.

Chillee · 2022-03-09T03:30:16Z

test/jit/test_op_decompositions.py

+        torch._C._jit_pass_run_decompositions(foo_s.graph)
+        inp = torch.rand([10, 10])
+        self.assertEqual(foo(inp), foo_s(inp))
+        FileCheck().check_not("aten::var").run(foo_s.graph)


:O - very cool!

Chillee · 2022-03-09T03:31:07Z

tools/codegen/decompositions/gen_jit_decompositions.py

+def main() -> None:
+    pytorch_dir = Path(__file__).resolve().parents[3]
+    upgrader_path = pytorch_dir / "torch" / "csrc" / "jit" / "runtime"
+    write_decomposition_util_file(str(upgrader_path))


I imagine this is just copied from the upgrader file? Should we rename it, or can we share infra?

There are like 5 shared lines.. i'll look into refactoring

If there's not that much started, dw about it.

Chillee · 2022-03-10T03:06:40Z

torch/csrc/jit/runtime/decomposition_registry.h

+TORCH_API c10::optional<std::shared_ptr<Graph>> DecompositionGraphForSchema(
+    const FunctionSchema& schema);
+
+TORCH_API void RunDecompositions(std::shared_ptr<Graph> g);


Would be good to have these allow selective use of decompositions.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution [ghstack-poisoned]

eellison · 2022-03-16T21:01:01Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-16T22:44:03Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-17T00:16:05Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-17T00:37:59Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-22T16:41:40Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-23T17:14:51Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-23T19:30:00Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-23T21:44:11Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…script for serializing the decompositions" This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Differential Revision: [D34938126](https://our.internmc.facebook.com/intern/diff/D34938126) [ghstack-poisoned]

eellison · 2022-03-25T19:43:59Z

@eellison has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

…serializing the decompositions (#73938) Summary: Pull Request resolved: #73938 This is a first step in porting and making usable all of the decompositions defined in [functorch](https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.py#L349) in core and in JIT as well as C++. The decompositions are defined in python, scripted and inlined, and then serialized as C++ code which TorchScript can parse. The workflow is edit python decomposition file then run [tools/codegen/decompositions/gen_jit_decompositions.py](https://github.com/pytorch/pytorch/pull/73938/files#diff-6adef2116be233c3524e3b583e373ab0ffc9169beb6c1f6d96b5d0385e75afa1). Decompositions are mapped to their corresponding aten schemas via the schema in their python def. This allows multiple decompositions for an overloaded op like `aten.var` (shown here in the example). This is just a first PR, i'm sure there will be many follows ups such as: - making these runnable in C++ with simple executor - porting over more decompositions from AOT Autograd - Using opinfos / more robust testing - Categorizing decompositions - Hooking in decompositions at various points of JIT execution Test Plan: Imported from OSS Reviewed By: gchanan Differential Revision: D34938126 Pulled By: eellison fbshipit-source-id: 9559a7cb731982e3a726f2f95af498b84fb09c13

…izing the decompositions ghstack-source-id: 55e0c48 Pull Request resolved: pytorch/pytorch#73938

Make aot autograd decompositions usable in JIT, add script for serial…

7771043

…izing the decompositions [ghstack-poisoned]

pytorch-bot bot added the ciflow/default label Mar 8, 2022

facebook-github-bot added the cla signed label Mar 8, 2022

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Mar 8, 2022

eellison changed the title ~~Make aot autograd decompositions usable in JIT, add script for serializing the decompositions~~ [JIT] Make aot autograd decompositions usable in JIT, add script for serializing the decompositions Mar 8, 2022

eellison requested a review from Chillee March 8, 2022 20:39

eellison pushed a commit that referenced this pull request Mar 8, 2022

Make aot autograd decompositions usable in JIT, add script for serial…

fd7d6cd

…izing the decompositions ghstack-source-id: 5dbf857 Pull Request resolved: #73938

This was referenced Mar 10, 2022

Clean up profiling mode and profiling executor strategy #73875

Closed

Get rid of DefaultNumBailouts #73689

Closed

Remove bailout logic #73876

Closed

eellison pushed a commit that referenced this pull request Mar 10, 2022

Make aot autograd decompositions usable in JIT, add script for serial…

a9ad177

…izing the decompositions ghstack-source-id: 6348f49 Pull Request resolved: #73938

eellison mentioned this pull request Mar 10, 2022

Introduce function-local settings for executor, expose in c++ #74012

Closed

Chillee approved these changes Mar 10, 2022

View reviewed changes

This was referenced Mar 10, 2022

Remove outdated logic #74043

Closed

Add api for changing function executor settings, hook up execution with decomposition registry #74186

Closed

Elias Ellison and others added 2 commits March 15, 2022 08:47

This was referenced Mar 16, 2022

Extend Graph Export to NNC, extend script to support CPU #74076

Closed

temp #74348

Closed

eellison and others added 2 commits March 16, 2022 16:21

Elias Ellison added 2 commits March 22, 2022 09:32

suo removed the ciflow/default label Mar 22, 2022

pytorchmergebot closed this in aacdf29 Mar 29, 2022

facebook-github-bot deleted the gh/eellison/281/head branch April 2, 2022 14:16

NesrineMHB pushed a commit to NesrineMHB/pytorch that referenced this pull request Apr 8, 2022

Make aot autograd decompositions usable in JIT, add script for serial…

4a68425

…izing the decompositions ghstack-source-id: 55e0c48 Pull Request resolved: pytorch/pytorch#73938

WBobby mentioned this pull request Aug 17, 2022

Add ROCm5.2.3/AMDGPU support for PyTorch WBobby/pytorch#2

Closed

[JIT] Make aot autograd decompositions usable in JIT, add script for serializing the decompositions #73938

[JIT] Make aot autograd decompositions usable in JIT, add script for serializing the decompositions #73938

Uh oh!

Conversation

eellison commented Mar 8, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Mar 8, 2022

⚛️ CI Flow

Uh oh!

facebook-github-bot commented Mar 8, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

💊 CI failures summary and remediations

Uh oh!

Chillee left a comment

Choose a reason for hiding this comment

Uh oh!

Chillee Mar 9, 2022

Choose a reason for hiding this comment

Uh oh!

Chillee Mar 9, 2022

Choose a reason for hiding this comment

Uh oh!

eellison Mar 10, 2022

Choose a reason for hiding this comment

Uh oh!

Chillee Mar 11, 2022

Choose a reason for hiding this comment

Uh oh!

Chillee Mar 10, 2022

Choose a reason for hiding this comment

Uh oh!

eellison commented Mar 16, 2022

Uh oh!

eellison commented Mar 16, 2022

Uh oh!

eellison commented Mar 17, 2022

Uh oh!

eellison commented Mar 17, 2022

Uh oh!

eellison commented Mar 22, 2022

Uh oh!

eellison commented Mar 23, 2022

Uh oh!

eellison commented Mar 23, 2022

Uh oh!

eellison commented Mar 23, 2022

Uh oh!

eellison commented Mar 25, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

eellison commented Mar 8, 2022 •

edited

Loading

facebook-github-bot commented Mar 8, 2022 •

edited

Loading