Skip to content

[JIT] Build mechanism for custom operators #10226

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
goldsborough wants to merge 6 commits into from
Closed

[JIT] Build mechanism for custom operators #10226

goldsborough wants to merge 6 commits into from

Conversation

@goldsborough
Copy link
Contributor

@goldsborough goldsborough commented Aug 3, 2018 •
edited
Loading

This is the last step in the custom operator implementation: providing a way to build from C++ and Python. For this I:

  1. Created a FindTorch.cmake taken largely from @ebetica with a CMake function to easily create simple custom op libraries
  2. Created a torch/op.h header for easy inclusion of necessary headers,
  3. Created a test directory pytorch/test/custom_operator which includes the basic setup for a custom op.
    1. It defines an op in op.{h,cpp}
    2. Registers it with the JIT using RegisterOperators
    3. Builds it into a shared library via a CMakeLists.txt
    4. Binds it into Python using a setup.py. This step makes use of our C++ extension setup that we already have. No work, yey!

The pure C++ and the Python builds are separate and not coupled in any way.

@zdevito @soumith @dzhulgakov

@goldsborough goldsborough force-pushed the cmake-for-custom-op branch 2 times, most recently from 9b8ab0f to e43ebb2 Compare August 6, 2018 22:31
@goldsborough goldsborough changed the title [WIP][JIT] Build mechanism for custom operators [JIT] Build mechanism for custom operators Aug 6, 2018
@goldsborough
Copy link
Contributor Author

goldsborough commented Aug 6, 2018

Tested in Python, now ready for review! Also added better docstrings to automatically bound ops:

screen shot 2018-08-06 at 15 30 37

facebook-github-bot
facebook-github-bot reviewed Aug 13, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

goldsborough
goldsborough commented Aug 13, 2018
View reviewed changes
test/custom_operator/test.py Outdated

This comment was marked as off-topic.

Sign in to view

dzhulgakov
dzhulgakov reviewed Aug 14, 2018
View reviewed changes
test/custom_operator/op.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

test/custom_operator/op.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/custom_operator.h Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/custom_operator.h Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

@zou3519 zou3519 added the oncall: jit Add this issue/PR to JIT oncall triage queue label Aug 14, 2018
@goldsborough
Copy link
Contributor Author

goldsborough commented Aug 14, 2018

Thanks for the review! I added code to dynamically load a shared library in the way that caffe2 does it. This avoid having to make the user .so Python-specific in any way, and doesn't even require a setup.py at all. I also added a test for using a registered op from a @torch.jit.script-ed function.

@goldsborough
Copy link
Contributor Author

goldsborough commented Aug 14, 2018 •
edited
Loading

One last TODO is to try the end-to-end flow of defining a script module in C++ that uses a custom operator registered inside Python, then exporting that to something that is loaded from C++, and then executing the module from C++.

Ok, actually, I will do this in a separate PR. There is more work to be done here. Let's keep this PR as the build integration.

dzhulgakov
dzhulgakov reviewed Aug 15, 2018
View reviewed changes
Copy link
Collaborator

@dzhulgakov dzhulgakov left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, but please double check the CMake part

cmake/Modules/FindTorch.cmake Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

test/custom_operator/test.py Outdated

This comment was marked as off-topic.

Sign in to view

test/test_jit.py Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

zdevito
zdevito approved these changes Aug 15, 2018
View reviewed changes
Copy link
Contributor

@zdevito zdevito left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The JIT changes look good. I don't know enough about CMake best practices to comment there.

test/custom_operator/CMakeLists.txt Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/custom_operator.h Outdated

This comment was marked as off-topic.

Sign in to view

torch/csrc/jit/custom_operator.h Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

@goldsborough
Create torch/op.h for custom operators
ab7fb4a 
Create CMake setup for custom ops

Fixes to enable custom op build

Figure out Python integration

Fix Windows

Fix python test for custom op and give bound operators better docstrings
@goldsborough
Dynamically load library with ctypes
158717b
@goldsborough
Add test for custom ops in script
a8a394f
@goldsborough
Uncomment tests
2429356
@goldsborough
Create TorchConfig.cmake
b017f0a
@goldsborough
Copy link
Contributor Author

goldsborough commented Aug 16, 2018

I think if tests pass I will land this diff. I know the CMake stuff can be improved, but this diff has code changes unrelated to the cmake. I will then create a new PR that integrates the custom operator tests with CI, in which we can discuss the proper way of doing the TorchConfig.cmake. I spent a little time looking at how Caffe2Config.cmake works, but there may be some open questions

facebook-github-bot
facebook-github-bot reviewed Aug 16, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

facebook-github-bot
facebook-github-bot reviewed Aug 16, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

facebook-github-bot
facebook-github-bot reviewed Aug 17, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@goldsborough goldsborough deleted the cmake-for-custom-op branch August 17, 2018 02:03
zdevito pushed a commit to zdevito/ATen that referenced this pull request Aug 17, 2018
@goldsborough @facebook-github-bot
Build mechanism for custom operators (#10226)
98bc979 
Summary:
This is the last step in the custom operator implementation: providing a way to build from C++ and Python. For this I:

1. Created a `FindTorch.cmake` taken largely from ebetica with a CMake function to easily create simple custom op libraries
2. Created a ` torch/op.h` header for easy inclusion of necessary headers,
3. Created a test directory `pytorch/test/custom_operator` which includes the basic setup for a custom op.
    1. It defines an op in `op.{h,cpp}`
    2. Registers it with the JIT using `RegisterOperators`
    3. Builds it into a shared library via a `CMakeLists.txt`
    4. Binds it into Python using a `setup.py`. This step makes use of our C++ extension setup that we already have. No work, yey!

The pure C++ and the Python builds are separate and not coupled in any way.

zdevito soumith dzhulgakov
Pull Request resolved: pytorch/pytorch#10226

Differential Revision: D9296839

Pulled By: goldsborough

fbshipit-source-id: 32f74cafb6e3d86cada8dfca8136d0dfb1f197a0
@ezyang ezyang added the merged label Jun 26, 2019
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@facebook-github-bot facebook-github-bot facebook-github-bot left review comments

@dzhulgakov dzhulgakov dzhulgakov left review comments

@zdevito zdevito zdevito approved these changes

@apaszke apaszke Awaiting requested review from apaszke

@colesbury colesbury Awaiting requested review from colesbury

@ezyang ezyang Awaiting requested review from ezyang

@gchanan gchanan Awaiting requested review from gchanan

@soumith soumith Awaiting requested review from soumith

Assignees

No one assigned

Labels

oncall: jit Add this issue/PR to JIT oncall triage queue

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

7 participants

@goldsborough @zdevito @apaszke @facebook-github-bot @dzhulgakov @ezyang @zou3519