Summary:
This PR implements the necessary hooks/stubs/enums/etc for complete ONNX Runtime (ORT) Eager Mode integration. The actual extension will live out of tree at https://github.com/pytorch/ort.
We have been [working on this at Microsoft](https://github.com/microsoft/onnxruntime-pytorch/tree/eager-ort/torch_onnxruntime) for the last few months, and are finally ready to contribute the PyTorch core changes upstream (nothing major or exciting, just the usual boilerplate for adding new backends).
The ORT backend will allow us to ferry [almost] all torch ops into granular ONNX kernels that ORT will eagerly execute against any devices it supports (therefore, we only need a single ORT backend from a PyTorch perspective).
Pull Request resolved: https://github.com/pytorch/pytorch/pull/58248
Reviewed By: astaff
Differential Revision: D30344992
Pulled By: albanD
fbshipit-source-id: 69082b32121246340d686e16653626114b7714b2
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/55294
Some static checkers like pyre have difficulties with types like `builtings.type`, so we strip the `builtins` prefix from autogened proto type stubs.
Test Plan: Let CI run.
Reviewed By: d4l3k
Differential Revision: D27477699
fbshipit-source-id: 45e19835974200a030817d37aec785e3ecb23e8b
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53735
Add an option to BlobSerializationOptions to request that float data be
serialized as bfloat16. This reduces the serialized data size at the expense
of some loss in precision.
ghstack-source-id: 124317910
Test Plan: Included a new unit test.
Reviewed By: mraway
Differential Revision: D26658205
fbshipit-source-id: 74521ed161059066355a3f208488ed01a344dbb5
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53403
This updates the `TensorProto` field to independently track the data type of
the in-memory (deserialized) data from the serialized data format.
This will allow us to support multiple different serialization formats in the
future. For instance, we could choose to perform quantization of floating
point data types, or varint encoding for integer fields.
For now this diff does not actually change the serialization code path yet,
and does not introduce any new serialization formats, but only refactors the
deserialization code path to make it easier to introduce new formats.
I'm not really that thrilled with the heavy use of macros and templates here,
but I didn't really see better alternatives that made it as simple to specify
new deserialization function implementations.
ghstack-source-id: 123594220
Test Plan:
Confirmed that the existing unit tests pass. This diff only touches the
deserialization code path and not the serialization code to help ensure that
the deserialization code works with the existing serialization logic, and that
there are no changes to the current serialization format.
Reviewed By: mraway
Differential Revision: D26658206
fbshipit-source-id: d7297d600aee28b92fd9f4ece437b7f519060942
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53402
Add an `options` field to the `Save` operator which accepts options for how to
serialize different blobs. At the moment this simply allows controlling the
existing `chunk_size` behavior, but in the future we can add other options,
such as the ability to control compression settings or other serialization
formats.
ghstack-source-id: 123567034
Test Plan:
Added a new test to `load_save_test.py` that passes in options and verifies
that blobs were serialized with the expected number of chunks.
buck test caffe2/caffe2:caffe2_test_cpu \
caffe2/caffe2/core:serialization_test \
caffe2/caffe2/python/operator_test:load_save_test
Reviewed By: mraway
Differential Revision: D26502577
fbshipit-source-id: 6e302e530bb96990517c2e35c505db7f14a56284
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52521
The `storage_type` and `external_data` fields were added a few years ago in
D10246743 (30aaa07594) but don't appear to have been used anywhere. Let's remove them to
help simplify the `TensorProto` message definition.
ghstack-source-id: 122110201
Test Plan: Confirmed the code still builds.
Reviewed By: dzhulgakov
Differential Revision: D26500028
fbshipit-source-id: 1e188f98f59e0b8673ea342ad9aaf7e5ba9b5fac
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52341
Add type stubs for caffe2 protos and scripts for generating them.
It's worth calling special attention to the following. In order to make `DeviceType`s like `CPU`, `CUDA`, etc. directly accessible from the `caffe2_pb2` module, they are currently freedom-patched into it in `caffe2/python/__init__.py`. This is not ideal: it would be better if these were autogenerated when the protobuf definitions were created by using `allow_alias = true` in the `DeviceTypeProto` definition in `caffe2.proto`.
However, it is impossible to do this currently without significant effort. The issue is that the generated proto constants would conflict with various constants defined in the C++ caffe2 codebase in `caffe2_pb.h`. We cannot simply remove these constants and replace them with the caffe2 DeviceTypeProto constants, because a huge portion of code expects `at::DeviceType` constants defined in `core/DeviceType.h` (apparently duplicated to avoid having to figure out how to autogenerate the protobuf definitions using cmake for ATen).
Instead, we make a best-effort to add additional definitions in `caffe2_pb2.py` by looking for any freedom-patched constants in `caffe2/python/__init__.py` and making sure they have corresponding stubs in the pyi (see `gen_proto_typestubs_helper.py`).
Test Plan: Make sure CI is green; we're just adding stubs.
Reviewed By: d4l3k
Differential Revision: D26331875
fbshipit-source-id: 2eea147e5bf393542f558ff8cf6385c47624b770