Make JIT Serialization support arbitrary std::function<> IO #27586

jjlilley · 2019-10-09T00:40:40Z

Summary:
Right now, torch::save() uses std::ostream, which results in unnecessary
data copies in practice. Similar for torch::load().

Adding a std::function<size_t(const void*, size_t)> as an output option,
parallel to the existing filename and std::ostream apis, gives users the
flexibility to emit directly to a backing store.

For a simple case of appending the output to a std::string, we observe
significant benchmark savings (on order of -50%), even with the
minor std::function<> dispatch overhead. The main reason is that
std::ostringstream effectively requires 2 extra copies of the data
beyond a simple string.append lambda.

We also provide a parallel api for the load(), though this one is
slightly more complex due to the need to do arbitrary position reads.

facebook-github-bot · 2019-10-09T00:40:54Z

This pull request was exported from Phabricator. Differential Revision: D17822962

caffe2/serialize/inline_container.cc

resistor · 2019-10-09T16:44:38Z

caffe2/serialize/inline_container.h

Is it really providing any value to have both paths? Could we not just convert all callers to use the std::function path?

That sounds good, I'll unify behind std::function<>.

resistor · 2019-10-09T16:45:32Z

test/cpp/api/serialize.cpp

For API orthogonality, we should probably make the same changes to load as well.

Yes, totally agree.
My intention was partly to gauge whether there would be objections that I hadn't thought of to this type of approach before doing the parallel load path. But I'll go ahead and change load as well.

zdevito · 2019-10-09T17:23:07Z

This is a good idea, Ill let Owen mark when it is ready to land.

facebook-github-bot · 2019-10-09T22:36:45Z

This pull request was exported from Phabricator. Differential Revision: D17822962

jjlilley · 2019-10-09T22:41:46Z

I just uploaded an updated version:

There is std::function<> support for both torch::load() and torch::save() codepaths now
I did simplify/unify the lowest layer PyTorchStreamWriter impl a bit, but chose not to remove the filename-based constructor api because of the python binding to that constructor.

facebook-github-bot · 2019-10-10T00:20:07Z

This pull request was exported from Phabricator. Differential Revision: D17822962

resistor · 2019-10-10T04:37:31Z

caffe2/serialize/inline_container.h

Why does this need to be a unique_ptr? I believe ofstream can be default constructed.

It didn't need to be a unique_ptr<>, was mostly trying to save work in the case where it's not needed. Reverted this part back.

facebook-github-bot · 2019-10-10T17:27:04Z

This pull request was exported from Phabricator. Differential Revision: D17822962

facebook-github-bot · 2019-10-10T17:47:29Z

This pull request was exported from Phabricator. Differential Revision: D17822962

jjlilley · 2019-10-10T17:48:06Z

(also fixed a few issues to make clang-tidy happy)

jjlilley · 2019-10-10T21:47:04Z

Looked into the failing tests - related to a bug in c10/util/tempfile.h, which was constructing a temp file name with a trailing '\0' in the std::string. Uploading the fix.

…#27586) Summary: Pull Request resolved: pytorch#27586 Right now, torch::save() uses std::ostream, which results in unnecessary data copies in practice. Similar for torch::load(). Adding a std::function<size_t(const void*, size_t)> as an output option, parallel to the existing filename and std::ostream apis, gives users the flexibility to emit directly to a backing store. For a simple case of appending the output to a std::string, we observe significant benchmark savings (on order of -50%), even with the minor std::function<> dispatch overhead. The main reason is that std::ostringstream effectively requires 2 extra copies of the data beyond a simple string.append lambda. We also provide a parallel api for the load(), though this one is slightly more complex due to the need to do arbitrary position reads. Test Plan: buck test mode/dev-nosan caffe2/test/... (Basic serialization test in caffe2/test/cpp/api/serialize.cpp) Benchmark in experimental/jeremyl/c2/SerializationBench.cpp, with D17823443 (1M time goes from 90ms -> 40ms, albeit with crc patch applied) Differential Revision: D17822962 fbshipit-source-id: 2cf2b24ad7d4212f2e4a49d7abc257cbdf33d206

facebook-github-bot · 2019-10-10T22:11:14Z

This pull request was exported from Phabricator. Differential Revision: D17822962

facebook-github-bot

@jjlilley is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

jjlilley · 2019-10-15T18:03:25Z

Thanks for looking at this!

yf225 · 2019-10-15T20:44:25Z

It seems that this PR is breaking macOS builds on master:

Oct 10 15:32:13 ../torch/csrc/api/src/serialize/input-archive.cpp:115:22: error: no matching function for call to 'min'
Oct 10 15:32:13       size_t nread = std::min(pos + n, size_) - pos;
Oct 10 15:32:13                      ^~~~~~~~

which also shows up in the CI for this PR.

yf225 · 2019-10-15T20:46:37Z

Sorry that to unbreak master, I will have to revert this PR.

jjlilley · 2019-10-15T21:01:08Z

Sorry about this! Thanks for reverting.
How do I notice this issue in the future, since this didn't show up on Sandcastle?

yf225 · 2019-10-15T21:25:16Z

We'll need to look at the Github CI on this PR - specifically pytorch_macos_10_13_py3_build and pytorch_macos_10_13_cuda9_2_cudnn7_py3_build was failing.

facebook-github-bot · 2019-10-15T23:47:44Z

@jjlilley merged this pull request in cbe5ab1.

…27586) Summary: Right now, torch::save() uses std::ostream, which results in unnecessary data copies in practice. Similar for torch::load(). Adding a std::function<size_t(const void*, size_t)> as an output option, parallel to the existing filename and std::ostream apis, gives users the flexibility to emit directly to a backing store. For a simple case of appending the output to a std::string, we observe significant benchmark savings (on order of -50%), even with the minor std::function<> dispatch overhead. The main reason is that std::ostringstream effectively requires 2 extra copies of the data beyond a simple string.append lambda. We also provide a parallel api for the load(), though this one is slightly more complex due to the need to do arbitrary position reads. Pull Request resolved: pytorch#27586 Test Plan: buck test mode/dev-nosan caffe2/test/... (Basic serialization test in caffe2/test/cpp/api/serialize.cpp) Benchmark in experimental/jeremyl/c2/SerializationBench.cpp, with D17823443 (1M time goes from 90ms -> 40ms, albeit with crc patch applied) Differential Revision: D17822962 Pulled By: jjlilley fbshipit-source-id: d344a7e59707f3b30d42280fbab78f87399e4d10

jjlilley requested review from apaszke, ebetica, goldsborough and yf225 as code owners October 9, 2019 00:40

pytorchbot added caffe2 oncall: jit Add this issue/PR to JIT oncall triage queue module: cpp Related to C++ API labels Oct 9, 2019

jjlilley requested review from resistor and zdevito October 9, 2019 00:42

resistor reviewed Oct 9, 2019

View reviewed changes

zdevito removed their request for review October 9, 2019 17:22

jjlilley force-pushed the export-D17822962 branch from 648da20 to 7061fea Compare October 9, 2019 22:36

jjlilley force-pushed the export-D17822962 branch from 7061fea to aaf28be Compare October 10, 2019 00:20

resistor reviewed Oct 10, 2019

View reviewed changes

jjlilley force-pushed the export-D17822962 branch from aaf28be to 7d01b5e Compare October 10, 2019 17:26

jjlilley force-pushed the export-D17822962 branch from 7d01b5e to 1ce2d05 Compare October 10, 2019 17:47

jjlilley force-pushed the export-D17822962 branch from 1ce2d05 to 9818c04 Compare October 10, 2019 22:11

pytorchbot added the module: internals Related to internal abstractions in c10 and ATen label Oct 10, 2019

jjlilley changed the title ~~Make JIT Serialization support arbitrary std::function<> output~~ Make JIT Serialization support arbitrary std::function<> IO Oct 11, 2019

jjlilley requested review from suo and zdevito October 11, 2019 23:59

resistor approved these changes Oct 12, 2019

View reviewed changes

zdevito removed their request for review October 15, 2019 03:48

facebook-github-bot reviewed Oct 15, 2019

View reviewed changes

facebook-github-bot closed this in cbe5ab1 Oct 15, 2019

facebook-github-bot added the merged label Oct 15, 2019

jjlilley mentioned this pull request Oct 16, 2019

[pytorch] Make JIT Serialization support arbitrary std::function<> IO #28039

Closed

mruberry added the Merged label Oct 28, 2020

Make JIT Serialization support arbitrary std::function<> IO #27586

Make JIT Serialization support arbitrary std::function<> IO #27586

Uh oh!

Conversation

jjlilley commented Oct 9, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Oct 9, 2019

Uh oh!

Uh oh!

resistor Oct 9, 2019

Choose a reason for hiding this comment

Uh oh!

jjlilley Oct 9, 2019

Choose a reason for hiding this comment

Uh oh!

resistor Oct 9, 2019

Choose a reason for hiding this comment

Uh oh!

jjlilley Oct 9, 2019

Choose a reason for hiding this comment

Uh oh!

zdevito commented Oct 9, 2019

Uh oh!

facebook-github-bot commented Oct 9, 2019

Uh oh!

jjlilley commented Oct 9, 2019

Uh oh!

facebook-github-bot commented Oct 10, 2019

Uh oh!

resistor Oct 10, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jjlilley Oct 10, 2019

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Oct 10, 2019

Uh oh!

facebook-github-bot commented Oct 10, 2019

Uh oh!

jjlilley commented Oct 10, 2019

Uh oh!

jjlilley commented Oct 10, 2019

Uh oh!

facebook-github-bot commented Oct 10, 2019

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

jjlilley commented Oct 15, 2019

Uh oh!

yf225 commented Oct 15, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yf225 commented Oct 15, 2019

Uh oh!

jjlilley commented Oct 15, 2019

Uh oh!

yf225 commented Oct 15, 2019

Uh oh!

facebook-github-bot commented Oct 15, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

jjlilley commented Oct 9, 2019 •

edited

Loading

resistor Oct 10, 2019 •

edited

Loading

yf225 commented Oct 15, 2019 •

edited

Loading