diff --git a/extension/aten_util/aten_bridge.cpp b/extension/aten_util/aten_bridge.cpp index d307e88bedd..7506327d4ea 100644 --- a/extension/aten_util/aten_bridge.cpp +++ b/extension/aten_util/aten_bridge.cpp @@ -134,6 +134,22 @@ c10::ScalarType executorch_to_torch_scalar_type( return static_cast(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(device.index())); +} + /* * Following makes two assumptions: * 1. aten_tensor's lifetime is longer than the liftime within which mutable_et diff --git a/extension/aten_util/aten_bridge.h b/extension/aten_util/aten_bridge.h index 62b07eee51d..18ffe8404e4 100644 --- a/extension/aten_util/aten_bridge.h +++ b/extension/aten_util/aten_bridge.h @@ -28,6 +28,17 @@ torch::executor::ScalarType torch_to_executorch_scalar_type( 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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index c71d91e5933..e65a1f05d5c 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -103,6 +103,7 @@ using torch::executor::ETDumpGen; #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 @@ struct PyMethod final { 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 sizes( + // Rebuilt rather than aliased, because an at::Tensor reports 64-bit sizes and strides while + // the runtime uses narrower types. + std::vector sizes( at_tensor.sizes().begin(), at_tensor.sizes().end()); - std::vector strides( + std::vector strides( at_tensor.strides().begin(), at_tensor.strides().end()); // Only works for MemoryFormat::Contiguous or MemoryFormat::ChannelsLast @@ -1133,14 +1134,18 @@ struct PyMethod final { " 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