Register et_copy device-copy ops for ATen-mode runtimes#20902
Register et_copy device-copy ops for ATen-mode runtimes#20902shoumikhin wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20902
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 34 Pending, 1 Unclassified FailureAs of commit 7fc0c28 with merge base ae754e9 ( UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@shoumikhin has exported this pull request. If you are a Meta employee, you can view the originating Diff in D111752990. |
This PR needs a
|
There was a problem hiding this comment.
Pull request overview
This PR fixes ATen-mode runtime failures for et_copy::_h2d_copy.out / et_copy::_d2h_copy.out by adding an ATen-mode registration path (custom-op schema + generated registration lib) and extending op__device_copy with an ATen-mode kernel implementation.
Changes:
- Add
et_copy_ops.yamlplusdevice_copy_ops_aten_libto register the device-copy ops in ATen-mode runtimes. - Adjust operator target wiring so
op__device_copy_atenis produced and can be linked by the ATen-mode registration library. - Implement and test an ATen-mode kernel branch for device copies (direction determined by op identity; uses registered CUDA allocator on current device).
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl | Moves op__device_copy into the list that produces an _aten target as well. |
| kernels/test/targets.bzl | Adds a dedicated ATen-mode device-copy unit test target. |
| kernels/test/op__device_copy_aten_test.cpp | New ATen-mode tests validating direction selection, resize behavior, and guardrails. |
| kernels/portable/targets.bzl | Exports et_copy_ops.yaml, adds an operator library for it, and adds device_copy_ops_aten_lib for ATen-mode registration. |
| kernels/portable/et_copy_ops.yaml | Defines custom-op schemas for et_copy device-copy ops for ATen-mode codegen/registration. |
| kernels/portable/cpu/op__device_copy.cpp | Adds ATen-mode kernel path and refactors device metadata access helpers. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
0c3ce1e to
6154fc2
Compare
Summary: The device-placement pass inserts `et_copy::_h2d_copy.out` / `_d2h_copy.out` at CPU<->accelerator boundaries (e.g. CUDA-delegated subgraphs). These portable ops are registered for the portable runtime via `functions.yaml`, but ATen-mode consumers hit "Missing operator: et_copy::_h2d_copy.out" at runtime: ATen-mode codegen only registers ATen `native_functions.yaml` ops, so `generated_lib_aten` never emits these custom ops and no `op__device_copy_aten` target existed. This adds an ATen-mode registration path: - `et_copy_ops.yaml` custom-ops schema + `device_copy_ops_aten_lib` (+`_static`) that registers the ops for ATen-mode runtimes, reusing the `op__device_copy` `_aten` kernel. - `op__device_copy` moved into `CUSTOM_OPS` to produce the `_aten` target; it stays registered for the portable runtime via `functions.yaml`. - `op__device_copy.cpp` gains an ATen-mode kernel branch: the planned `at::Tensor` carries no device metadata in ATen mode, so direction comes from op identity (_h2d/_d2h) and the copy uses the registered CUDA DeviceAllocator on the current device. The portable branch is factored to share device-accessor helpers but keeps its existing behavior. Scope/limitations: - ATen-mode device copy is CUDA-only and single-GPU (current device): the planned device type/index are not recoverable from the at::Tensor, so the kernel targets the registered CUDA allocator's current device. Multi-GPU / non-CUDA support is a follow-up (see the TODO in op__device_copy.cpp). fbcode and xplat copies are kept identical. Reviewed By: digantdesai Differential Revision: D111752990
6154fc2 to
91c178d
Compare
91c178d to
6e6bed5
Compare
6e6bed5 to
92d38b4
Compare
Summary: The device-placement pass inserts `et_copy::_h2d_copy.out` / `_d2h_copy.out` at CPU<->accelerator boundaries (e.g. CUDA-delegated subgraphs). These portable ops are registered for the portable runtime via `functions.yaml`, but ATen-mode consumers hit "Missing operator: et_copy::_h2d_copy.out" at runtime: ATen-mode codegen only registers ATen `native_functions.yaml` ops, so `generated_lib_aten` never emits these custom ops and no `op__device_copy_aten` target existed. This adds an ATen-mode registration path: - `et_copy_ops.yaml` custom-ops schema + `device_copy_ops_aten_lib` (+`_static`) that registers the ops for ATen-mode runtimes, reusing the `op__device_copy` `_aten` kernel. - `op__device_copy` moved into `CUSTOM_OPS` to produce the `_aten` target; it stays registered for the portable runtime via `functions.yaml`. - `op__device_copy.cpp` gains an ATen-mode kernel branch: the planned `at::Tensor` carries no device metadata in ATen mode, so direction comes from op identity (_h2d/_d2h) and the copy uses the registered CUDA DeviceAllocator on the current device. The portable branch is factored to share device-accessor helpers but keeps its existing behavior. Scope/limitations: - ATen-mode device copy is CUDA-only and single-GPU (current device): the planned device type/index are not recoverable from the at::Tensor, so the kernel targets the registered CUDA allocator's current device. Multi-GPU / non-CUDA support is a follow-up (see the TODO in op__device_copy.cpp). fbcode and xplat copies are kept identical. Reviewed By: digantdesai Differential Revision: D111752990
92d38b4 to
7fc0c28
Compare
digantdesai
left a comment
There was a problem hiding this comment.
Do we need any CMake changes?
Summary:
The device-placement pass inserts
et_copy::_h2d_copy.out/_d2h_copy.outatCPU<->accelerator boundaries (e.g. CUDA-delegated subgraphs). These portable ops
are registered for the portable runtime via
functions.yaml, but ATen-modeconsumers hit "Missing operator: et_copy::_h2d_copy.out" at runtime: ATen-mode
codegen only registers ATen
native_functions.yamlops, sogenerated_lib_atennever emits these custom ops and no
op__device_copy_atentarget existed.This adds an ATen-mode registration path:
et_copy_ops.yamlcustom-ops schema +device_copy_ops_aten_lib(+_static)that registers the ops for ATen-mode runtimes, reusing the
op__device_copy_atenkernel.op__device_copymoved intoCUSTOM_OPSto produce the_atentarget; itstays registered for the portable runtime via
functions.yaml.op__device_copy.cppgains an ATen-mode kernel branch: the plannedat::Tensorcarries no device metadata in ATen mode, so direction comes from op identity
(_h2d/_d2h) and the copy uses the registered CUDA DeviceAllocator on the
current device. The portable branch is factored to share device-accessor
helpers but keeps its existing behavior.
Scope/limitations:
planned device type/index are not recoverable from the at::Tensor, so the
kernel targets the registered CUDA allocator's current device. Multi-GPU /
non-CUDA support is a follow-up (see the TODO in op__device_copy.cpp).
fbcode and xplat copies are kept identical.
Differential Revision: D111752990