[ET Device Support] Add device tensor helper functions to TensorPtr API

Add clone_tensor_ptr_to_device and clone_tensor_ptr_to_cpu to tensor_ptr.h for cloning tensors between host and device memory via DeviceAllocatorRegistry.
Extend the existing make_tensor_ptr(const TensorPtr&, ...) overload with optional device_type/device_index parameters (default CPU/0) for seamless device placement.

Differential Revision: [D99913077](https://our.internmc.facebook.com/intern/diff/D99913077/)

ghstack-source-id: 364093618
Pull Request resolved: #18761

Author: Gasoonjia

extension/tensor/targets.bzl

@@ -24,6 +24,7 @@ def define_common_targets():
             ],
             visibility = ["PUBLIC"],
             deps = [
+                "//executorch/runtime/core:device_allocator",
                 "//executorch/runtime/core/exec_aten/util:dim_order_util" + aten_suffix,
                 "//executorch/runtime/core/exec_aten/util:tensor_util" + aten_suffix,
             ],

extension/tensor/tensor_ptr.cpp

@@ -10,6 +10,7 @@
 
 #include <numeric>
 
+#include <executorch/runtime/core/device_allocator.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 
 namespace executorch {
@@ -248,5 +249,179 @@ runtime::Error resize_tensor_ptr(
           sizes.data(), sizes.size()));
 }
 
+// ---- Device tensor helpers ----
+
+namespace {
+
+#ifndef USE_ATEN_LIB
+struct DeviceStorage final {
+  executorch::aten::TensorImpl tensor_impl;
+  executorch::aten::Tensor tensor;
+  std::vector<executorch::aten::SizesType> sizes;
+  std::vector<executorch::aten::DimOrderType> dim_order;
+  std::vector<executorch::aten::StridesType> strides;
+  std::function<void(void*)> deleter;
+
+  DeviceStorage(
+      executorch::aten::TensorImpl&& tensor_impl,
+      std::vector<executorch::aten::SizesType>&& sizes,
+      std::vector<executorch::aten::DimOrderType>&& dim_order,
+      std::vector<executorch::aten::StridesType>&& strides,
+      std::function<void(void*)>&& deleter)
+      : tensor_impl(std::move(tensor_impl)),
+        tensor(&this->tensor_impl),
+        sizes(std::move(sizes)),
+        dim_order(std::move(dim_order)),
+        strides(std::move(strides)),
+        deleter(std::move(deleter)) {}
+
+  ~DeviceStorage() {
+    if (deleter) {
+      deleter(tensor_impl.mutable_data());
+    }
+  }
+};
+#endif // USE_ATEN_LIB
+
+TensorPtr make_tensor_ptr_with_device(
+    std::vector<executorch::aten::SizesType> sizes,
+    void* data,
+    executorch::aten::ScalarType type,
+    runtime::etensor::DeviceType device_type,
+    runtime::etensor::DeviceIndex device_index,
+    std::function<void(void*)> deleter) {
+  const auto dim = sizes.size();
+  std::vector<executorch::aten::DimOrderType> dim_order(dim);
+  std::iota(dim_order.begin(), dim_order.end(), 0);
+
+  std::vector<executorch::aten::StridesType> strides(dim);
+  if (dim > 0) {
+    auto error = runtime::dim_order_to_stride(
+        sizes.data(), dim_order.data(), dim, strides.data());
+    ET_CHECK_MSG(error == runtime::Error::Ok, "Failed to compute strides.");
+  }
+
+#ifndef USE_ATEN_LIB
+  executorch::aten::TensorImpl tensor_impl(
+      type,
+      dim,
+      sizes.data(),
+      data,
+      dim_order.data(),
+      strides.data(),
+      dim > 0 ? executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND
+              : executorch::aten::TensorShapeDynamism::STATIC,
+      device_type,
+      device_index);
+  auto storage = std::make_shared<DeviceStorage>(
+      std::move(tensor_impl),
+      std::move(sizes),
+      std::move(dim_order),
+      std::move(strides),
+      std::move(deleter));
+  const auto tensor_ptr = &storage->tensor;
+  return std::shared_ptr<executorch::aten::Tensor>(
+      std::move(storage), tensor_ptr);
+#else
+  (void)device_type;
+  (void)device_index;
+  return make_tensor_ptr(
+      std::move(sizes),
+      data,
+      type,
+      executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND,
+      std::move(deleter));
+#endif // USE_ATEN_LIB
+}
+
+} // namespace
+
+TensorPtr clone_tensor_ptr_to_device(
+    const TensorPtr& cpu_tensor,
+    runtime::etensor::DeviceType device_type,
+    runtime::etensor::DeviceIndex device_index) {
+  ET_CHECK_MSG(
+      device_type != runtime::etensor::DeviceType::CPU,
+      "Target device must not be CPU; use clone_tensor_ptr for CPU-to-CPU copies.");
+
+  auto* allocator = runtime::get_device_allocator(device_type);
+  ET_CHECK_MSG(
+      allocator != nullptr,
+      "No device allocator registered for device type %d",
+      static_cast<int>(device_type));
+
+  const auto nbytes = cpu_tensor->nbytes();
+  const auto* cpu_data = cpu_tensor->const_data_ptr();
+  ET_CHECK_MSG(cpu_data != nullptr, "Source tensor has no data.");
+
+  auto result = allocator->allocate(nbytes, device_index);
+  ET_CHECK_MSG(result.ok(), "Failed to allocate device memory.");
+  void* device_data = result.get();
+
+  auto err = allocator->copy_host_to_device(
+      device_data, cpu_data, nbytes, device_index);
+  ET_CHECK_MSG(err == runtime::Error::Ok, "Host-to-device copy failed.");
+
+  std::vector<executorch::aten::SizesType> sizes(
+      cpu_tensor->sizes().begin(), cpu_tensor->sizes().end());
+
+  return make_tensor_ptr_with_device(
+      std::move(sizes),
+      device_data,
+      cpu_tensor->scalar_type(),
+      device_type,
+      device_index,
+      [allocator, device_index](void* ptr) {
+        allocator->deallocate(ptr, device_index);
+      });
+}
+
+TensorPtr clone_tensor_ptr_to_cpu(const TensorPtr& device_tensor) {
+  const auto nbytes = device_tensor->nbytes();
+  const auto* device_data = device_tensor->const_data_ptr();
+  ET_CHECK_MSG(device_data != nullptr, "Source device tensor has no data.");
+
+#ifndef USE_ATEN_LIB
+  const auto device_type = device_tensor->unsafeGetTensorImpl()->device_type();
+  const auto device_index =
+      device_tensor->unsafeGetTensorImpl()->device_index();
+#else
+  const auto& aten_device = device_tensor->device();
+  ET_CHECK_MSG(!aten_device.is_cpu(), "Source tensor is already on CPU.");
+  auto device_type = runtime::etensor::DeviceType::CPU;
+  if (aten_device.is_cuda()) {
+    device_type = runtime::etensor::DeviceType::CUDA;
+  }
+  const auto device_index =
+      static_cast<runtime::etensor::DeviceIndex>(aten_device.index());
+#endif
+
+  ET_CHECK_MSG(
+      device_type != runtime::etensor::DeviceType::CPU,
+      "Source tensor is already on CPU.");
+
+  auto* allocator = runtime::get_device_allocator(device_type);
+  ET_CHECK_MSG(
+      allocator != nullptr,
+      "No device allocator registered for device type %d",
+      static_cast<int>(device_type));
+
+  std::vector<uint8_t> cpu_data(nbytes);
+
+  auto err = allocator->copy_device_to_host(
+      cpu_data.data(), device_data, nbytes, device_index);
+  ET_CHECK_MSG(err == runtime::Error::Ok, "Device-to-host copy failed.");
+
+  std::vector<executorch::aten::SizesType> sizes(
+      device_tensor->sizes().begin(), device_tensor->sizes().end());
+
+  return make_tensor_ptr(
+      std::move(sizes),
+      std::move(cpu_data),
+      {},
+      {},
+      device_tensor->scalar_type());
+}
+
 } // namespace extension
 } // namespace executorch

extension/tensor/tensor_ptr.h

@@ -17,6 +17,7 @@
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/portable_type/device.h>
 
 C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")
 
@@ -388,27 +389,59 @@ inline TensorPtr make_tensor_ptr(
       std::move(deleter));
 }
 
+/**
+ * Clones a CPU TensorPtr to a device TensorPtr.
+ *
+ * Allocates memory on the specified device and copies the tensor data from
+ * host to device using the DeviceAllocator registered for the given device
+ * type. The returned TensorPtr owns the device memory and will free it via
+ * the allocator when destroyed.
+ *
+ * Forward declaration to support make_tensor_ptr below usage.
+ *
+ * @param cpu_tensor The source CPU tensor whose data will be copied.
+ * @param device_type The target device type (e.g., DeviceType::CUDA).
+ * @param device_index The target device index (default 0).
+ * @return A TensorPtr backed by device memory containing the copied data.
+ */
+TensorPtr clone_tensor_ptr_to_device(
+    const TensorPtr& cpu_tensor,
+    runtime::etensor::DeviceType device_type,
+    runtime::etensor::DeviceIndex device_index = 0);
+
 /**
  * Convenience overload identical to make_tensor_ptr(*tensor_ptr, ...).
  * Keeps the original TensorPtr alive until the returned TensorPtr is destroyed.
+ * When device_type is not CPU, the tensor data is additionally copied to the
+ * specified device.
  *
  * @param tensor_ptr The source tensor pointer to alias.
  * @param sizes Optional sizes override.
  * @param dim_order Optional dimension order override.
  * @param strides Optional strides override.
- * @return A TensorPtr aliasing the same storage with requested metadata.
+ * @param device_type The target device type (default CPU, meaning no copy).
+ * @param device_index The target device index (default 0).
+ * @return A TensorPtr aliasing the same storage with requested metadata, or a
+ * device TensorPtr if device_type is not CPU.
  */
 inline TensorPtr make_tensor_ptr(
     const TensorPtr& tensor_ptr,
     std::vector<executorch::aten::SizesType> sizes = {},
     std::vector<executorch::aten::DimOrderType> dim_order = {},
-    std::vector<executorch::aten::StridesType> strides = {}) {
-  return make_tensor_ptr(
+    std::vector<executorch::aten::StridesType> strides = {},
+    runtime::etensor::DeviceType device_type =
+        runtime::etensor::DeviceType::CPU,
+    runtime::etensor::DeviceIndex device_index = 0) {
+  auto result = make_tensor_ptr(
       *tensor_ptr,
       std::move(sizes),
       std::move(dim_order),
       std::move(strides),
       [tensor_ptr](void*) {});
+  if (device_type != runtime::etensor::DeviceType::CPU) {
+    return clone_tensor_ptr_to_device(result, device_type, device_index);
+  }
+  return result;
 }
 
 /**
@@ -479,6 +512,18 @@ runtime::Error resize_tensor_ptr(
     TensorPtr& tensor,
     const std::vector<executorch::aten::SizesType>& sizes);
 
+/**
+ * Clones a device TensorPtr to a CPU TensorPtr.
+ *
+ * Allocates host memory and copies the tensor data from device to host using
+ * the DeviceAllocator registered for the source tensor's device type. The
+ * device type is determined from the source tensor's metadata.
+ *
+ * @param device_tensor The source device tensor whose data will be copied.
+ * @return A TensorPtr backed by CPU memory containing the copied data.
+ */
+TensorPtr clone_tensor_ptr_to_cpu(const TensorPtr& device_tensor);
+
 } // namespace extension
 } // namespace executorch
 

extension/tensor/test/targets.bzl

@@ -21,3 +21,14 @@ def define_common_targets():
                 "//executorch/extension/tensor:tensor" + aten_suffix,
             ],
         )
+
+        runtime.cxx_test(
+            name = "tensor_ptr_device_test" + aten_suffix,
+            srcs = [
+                "tensor_ptr_device_test.cpp",
+            ],
+            deps = [
+                "//executorch/extension/tensor:tensor" + aten_suffix,
+                "//executorch/runtime/core:device_allocator",
+            ],
+        )