[executorch] Propagate device metadata from partitioner result onto TensorSpecs

[Gasoonjia]

Stack from ghstack (oldest at bottom):

• [ET Device Support] CUDA-native Qwen 3.5 MoE inference with device tensor pipeline #18788
• [ET Device Support] Extract shared device test utilities to reduce redundancy #18762
• [ET Device Support] Add device tensor helper functions to TensorPtr API #18761
• [ET Device Support] Add ExecutorchBackendConfig flags for skipping H2D/D2H copies #18760
• [ET Device Support] PropagateDevicePass inserts H2D/D2H copy ops at delegate boundaries #18730
• [ET Device Support] Define et_copy runtime h2d and d2h copy ops #18729
• [ET Device Support] Define AOT device copy ops registry #18728
• [ET Device Support] CudaAllocator: device memory allocator for CUDA backend #18477
• [ET Device Support] Module: allocate device memory for planned buffers #18476
• [ET Device Support] MemoryManager: add per-buffer device metadata #18475
• [ET Device Support] MethodMeta: expose per-buffer device placement API #18474
• [ET Device Support] DeviceMemoryBuffer RAII class for device memory lifetime management #18473
• [ET Device Support] Emitter reads non_const_buffer_device from graph meta #18472
• [ET Device Support] Device-aware memory planning: separate buffers per device type #18375
• [ET Device Support] Add NonConstBufferDevice schema for per-buffer device mapping #18330
• [ET Device Support] DeviceAllocator interface and DeviceAllocatorRegistry #17535
• [ET Device Support] Annotate device attributes of CUDA backend IO tensors cuda device #18080
• [ET Device Support] Parse device info from serialized tensor in tensor_parser #18328
• [ET Device Support] Propagate device info from TensorSpec into serialized Tensor #18079
• -> [executorch] Propagate device metadata from partitioner result onto TensorSpecs #18078

Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph

The overall pipeline is:
a. Partitioner use compile_spec to determine which device the partitoned blob is runing on
b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device.

Differential Revision: D95842511

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/18078

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 2 Unrelated Failures

As of commit 4708adb with merge base b5ae0b9 ():

NEW FAILURE - The following job has failed:

• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t ecec9814515d4a075c02a4e6e38a7b018d073a9d4f0d8f8ae346e9914b673bd8 /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

This effectively assumes that we know the device 'name' AoT. In theory, we can have a multi-device delegate then the runtime might interpret this name differently and that can cause some confusion i.e cuda:0 device on Metal.

I am not sure about using generic names like 'gpu' but also not sure about following PyTorch's eager/jit style naming convention where you won't switch devices underneath.

[Gasoonjia]

May I have your suggestions on the executorch device name?

Currently we set up the device name AOT and intentionally decouple dour device attribute with pytorch/pytorch device concept; we created a enum in the etensor schema for all devices we are supporting right now. In this way we can support as much as device as we want.

For the situaton you mentioned, if other backend like vulken need its own gpu device, they should add a new one to the enum. We should avoid using generic names like 'gpu'.

[JacobSzwejbka]

Multi-device graph serialization will necessitate multiple graphs. We can maybe make an exception for input tensors, but for any intermediate the runtime needs to know what the device its loading intermediates onto.

Device is fixed at export aot. If you want to have some generic shader style lib where the gpu type is decided lazily then you will have to use a generic key like gpu.

[JacobSzwejbka]

nit: this type of util really should be placed in a single spot. There are other things like this in the passes. Lets take it as a follow up to have claude just search for generic utils like this and centralize them

[JacobSzwejbka]

Are these fields already in the TensorSpec class definition? Are they initialized to just cpu and 0?

[JacobSzwejbka]

We should throw here no?

[JacobSzwejbka]

Why does it not return cpu by default

[JacobSzwejbka]

Can we just do 1 pass. You can look at users of the delegate node to find the getitem nodes.