Tell the runtime which device a Python input lives on

[shoumikhin]

The problem

When you pass a tensor to a model from Python, ExecuTorch builds its own description of
that tensor: the shape, the data type, and where the memory lives. That last part was
never filled in, so it always said main memory (the host):

TensorPtr tensor = for_blob(data_ptr, sizes, type)
    .strides(...).dim_order(...).make_tensor_ptr();   // device silently defaults to CPU

For an ordinary model that is correct, since the tensor really is in main memory. It is
wrong for a model whose activations stay on the GPU, where you hand over memory that is
already on the GPU. The runtime reads this description to decide what to do with your
input, so it was being told the wrong thing, and nothing downstream could tell the
difference between a host tensor and a GPU one.

The change

Pass the tensor's device through instead of letting it default:

TensorPtr tensor = make_tensor_ptr(
    std::move(sizes), data_ptr, std::move(dim_order), std::move(strides), type,
    aten::TensorShapeDynamism::STATIC, nullptr,
    torch_to_executorch_device(at_tensor.device()));

The tensor factory already accepts a device, and its documentation describes exactly this
case, so nothing new had to be invented:

The `device` parameter sets the Tensor's device location only, no data is
allocated or copied. The caller is responsible for ensuring `data` already
lives on the requested device.

The one new piece is a small translator beside the existing scalar-type one, because
PyTorch and the ExecuTorch runtime describe devices with different types:

executorch::runtime::etensor::Device torch_to_executorch_device(c10::Device device);

It raises an error for a device the runtime does not model, rather than quietly reporting
host memory. Quietly reporting host memory is what made the original problem invisible.

The size and stride vectors now use the runtime's own types, which the factory takes
directly, so no extra conversion happens at the call.

Test plan

• ran a CPU model from Python and confirmed it still matches eager PyTorch exactly. Every
  existing caller goes through this code path, so this is the important check: the device
  of an ordinary tensor is the host, which is what the code used to assume.
• ran the Python binding test suite as a baseline. Two quantized-operator tests fail
  before this change as well, for an unrelated reason, and the remaining 35 pass.

Tested on Linux x86_64. Continuous integration builds and runs these bindings, which is
where the compiled result is exercised.

Note

This describes the input correctly; it does not by itself make a GPU-activation program
run from Python. Such a program also has to be exported so the runtime shares the
caller's memory rather than reserving its own buffer for it.

[shoumikhin]

Stack from ghstack (oldest at bottom):

• -> Tell the runtime which device a Python input lives on #21705

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 1 Cancelled Job

As of commit 403891f with merge base 730b77a ():

NEW FAILURE - The following job has failed:

• Lint / lintrunner (gh)
  >>> Lint for extension/pybindings/pybindings.cpp:

CANCELLED JOB - The following job was cancelled. Please retry:

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

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

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[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.