Tell the runtime which device a Python input lives on

extension/aten_util/aten_bridge.cpp

@@ -134,6 +134,22 @@ c10::ScalarType executorch_to_torch_scalar_type(
   return static_cast<c10::ScalarType>(intermediate);
 }
 
+executorch::runtime::etensor::Device torch_to_executorch_device(
+    c10::Device device) {
+  using executorch::runtime::etensor::DeviceIndex;
+  using executorch::runtime::etensor::DeviceType;
+
+  if (device.is_cpu()) {
+    return executorch::runtime::etensor::Device(DeviceType::CPU);
+  }
+  ET_CHECK_MSG(
+      device.is_cuda(),
+      "Device type %s is not supported by the ExecuTorch runtime",
+      c10::DeviceTypeName(device.type()).c_str());
+  return executorch::runtime::etensor::Device(
+      DeviceType::CUDA, static_cast<DeviceIndex>(device.index()));
+}
+
 /*
  * Following makes two assumptions:
  * 1. aten_tensor's lifetime is longer than the liftime within which mutable_et

extension/aten_util/aten_bridge.h

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -28,6 +28,17 @@
 c10::ScalarType executorch_to_torch_scalar_type(
     torch::executor::ScalarType type);
 
+/**
+ * Translates a PyTorch device into the runtime's own device description.
+ *
+ * Needed so a tensor handed in from Python is described as living where it actually lives. A device
+ * tensor labelled as host memory is copied with a host memcpy, which is undefined.
+ *
+ * Throws for a device type the runtime does not model, rather than silently reporting host memory.
+ */
+executorch::runtime::etensor::Device torch_to_executorch_device(
+    c10::Device device);
+
 /*
  * @param[in] aten_tensor Input at::Tensor
  * @param[in,out] mutable_et ETensor whose underlying memory now will alias to

extension/pybindings/pybindings.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -103,6 +103,7 @@
 #ifndef USE_ATEN_LIB
 using ::executorch::extension::alias_attensor_to_etensor;
 using ::executorch::extension::alias_etensor_to_attensor;
+using ::executorch::extension::torch_to_executorch_device;
 using ::executorch::extension::torch_to_executorch_scalar_type;
 #endif // !USE_ATEN_LIB
 
@@ -1109,11 +1110,11 @@
         auto type =
             torch_to_executorch_scalar_type(at_tensor.options().dtype());
         size_t dim = at_tensor.dim();
-        // cant directly alias at::Tensor sizes and strides due to int64 vs
-        // int32 typing conflict
-        std::vector<int> sizes(
+        // Rebuilt rather than aliased, because an at::Tensor reports 64-bit sizes and strides while
+        // the runtime uses narrower types.
+        std::vector<aten::SizesType> sizes(
             at_tensor.sizes().begin(), at_tensor.sizes().end());
-        std::vector<int> strides(
+        std::vector<aten::StridesType> strides(
             at_tensor.strides().begin(), at_tensor.strides().end());
 
         // Only works for MemoryFormat::Contiguous or MemoryFormat::ChannelsLast
@@ -1133,14 +1134,18 @@
               " should be contiguous or channels-last.";
           throw std::runtime_error(error_msg);
         }
-        TensorPtr tensor = for_blob(
-                               mutable_tensor_data_ptr_no_cow(at_tensor),
-                               std::move(sizes),
-                               type)
-                               .strides(std::move(strides))
-                               .dim_order(std::move(dim_order))
-                               .dynamism(aten::TensorShapeDynamism::STATIC)
-                               .make_tensor_ptr();
+        // The device comes from the tensor the caller passed, so a program whose activations stay on
+        // an accelerator sees its input described as device memory. Defaulting to CPU labelled a GPU
+        // input as host memory, and the runtime then copied it with a host memcpy.
+        TensorPtr tensor = make_tensor_ptr(
+            std::move(sizes),
+            mutable_tensor_data_ptr_no_cow(at_tensor),
+            std::move(dim_order),
+            std::move(strides),
+            type,
+            aten::TensorShapeDynamism::STATIC,
+            nullptr,
+            torch_to_executorch_device(at_tensor.device()));
         input_tensors.push_back(tensor);
         EValue evalue(input_tensors.back());
 #endif