Update doc to indicate how a CUDA custom op is synchronized with respect to ORT's CUDA kernels (#6981)

* Update docs

* Formatting

* Formatting

* update
This commit is contained in:
Hariharan Seshadri 2021-03-11 18:43:44 -08:00 committed by GitHub
parent 2da5895b3a
commit 35eb8e70ee
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -23,16 +23,18 @@ Use the custom operator C/C++ API (onnxruntime_c_api.h)
* Create an OrtCustomOp structure for each op and add them to the OrtCustomOpDomain with OrtCustomOpDomain_Add
* Call OrtAddCustomOpDomain to add the custom domain of ops to the session options
See [this](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/test/shared_lib/test_inference.cc) for examples called MyCustomOp and SliceCustomOp that use the C++ helper API (onnxruntime_cxx_api.h).
See [this](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/test/shared_lib/test_inference.cc) for examples called `MyCustomOp` and `SliceCustomOp` that use the C++ helper API (onnxruntime_cxx_api.h).
You can also compile the custom ops into a shared library and use that to run a model via the C++ API. The same test file contains an example.
The source code for a sample custom op shared library containing two custom kernels is [here](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc).
See [this](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/test/python/onnxruntime_test_python.py) for an example called testRegisterCustomOpsLibrary that uses the Python API to register a shared library that contains custom op kernels. Currently, the only supported Execution Providers (EPs) for custom ops registered via this approach are the CUDA and the CPU EPs.
See [this](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/test/python/onnxruntime_test_python.py) for an example called `testRegisterCustomOpsLibrary` that uses the Python API to register a shared library that contains custom op kernels. Currently, the only supported Execution Providers (EPs) for custom ops registered via this approach are the CUDA and the CPU EPs.
Note that when a model being inferred on gpu, onnxruntime will insert MemcpyToHost op before a cpu custom op and append MemcpyFromHost after to make sure tensor(s) are accessible throughout calling, meaning there are no extra efforts required from custom op developer for the case.
When using CUDA custom ops, to ensure synchronization between ORT's CUDA kernels and the custom CUDA kernels, they must all use the same CUDA compute stream. To ensure this, you may first create a CUDA stream and pass it to the underlying Session via SessionOptions (use `OrtCudaProviderOptions` struct). This will ensure ORT's CUDA kernels use that stream and if the custom CUDA kernels are launched using the same stream, synchronization is now taken care of implicitly. For a sample, please see how the afore-mentioned `MyCustomOp` is being launched and how the Session using this custom op is created.
## Use RegisterCustomRegistry API
Implement your kernel and schema (if required) using the OpKernel and OpSchema APIs (headers are in the include folder).