Skip to content

Commit c7debeb

Browse files
committed
[slimtensor] Add from_etensor factory function for ETensor to SlimTensor conversion
Add from_etensor() factory function that creates a SlimTensor from an ExecuTorch portable tensor (ETensor), copying data to a target device. Key features: - Handles int32_t to int64_t conversion for sizes/strides (ETensor uses int32_t, SlimTensor uses int64_t) - Supports CPU and CUDA target devices via storage()->copy_() - Preserves tensor strides (non-contiguous layouts) - Provides both reference and pointer overloads Differential Revision: [D90539554](https://our.internmc.facebook.com/intern/diff/D90539554/) ghstack-source-id: 333060891 Pull Request resolved: #16551
1 parent 83a5ece commit c7debeb

File tree

4 files changed

+449
-0
lines changed

4 files changed

+449
-0
lines changed
Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <executorch/backends/aoti/slim/core/slim_tensor.h>
12+
#include <executorch/backends/aoti/slim/factory/empty.h>
13+
#include <executorch/backends/aoti/slim/util/array_ref_util.h>
14+
#include <executorch/runtime/core/portable_type/tensor.h>
15+
16+
namespace executorch::backends::aoti::slim {
17+
18+
/// Creates a SlimTensor from an ETensor (ExecuTorch portable tensor).
19+
///
20+
/// This factory function converts an ETensor to a SlimTensor, optionally
21+
/// copying the data to a target device. The ETensor is assumed to always
22+
/// reside on CPU.
23+
///
24+
/// @param etensor The source ETensor (always on CPU).
25+
/// @param target_device The target device for the output SlimTensor.
26+
/// @return A new SlimTensor with data copied to the target device.
27+
///
28+
/// @note ETensor uses int32_t (SizesType/StridesType) for sizes and strides,
29+
/// while SlimTensor uses int64_t. This function handles the conversion.
30+
///
31+
/// Example usage:
32+
/// @code
33+
/// auto* cpu_tensor = &(args[i]->toTensor()); // ETensor from EValue
34+
/// SlimTensor gpu_tensor = from_etensor(*cpu_tensor, DEFAULT_CUDA_DEVICE);
35+
/// @endcode
36+
inline SlimTensor from_etensor(
37+
const executorch::runtime::etensor::Tensor& etensor,
38+
const c10::Device& target_device = CPU_DEVICE) {
39+
// Step 1: Extract metadata from ETensor
40+
const auto ndim = static_cast<size_t>(etensor.dim());
41+
42+
// Convert sizes from exec_aten::SizesType (int32_t) to int64_t
43+
std::vector<int64_t> sizes_vec(ndim);
44+
for (size_t i = 0; i < ndim; ++i) {
45+
sizes_vec[i] = static_cast<int64_t>(etensor.size(static_cast<ssize_t>(i)));
46+
}
47+
48+
// Convert strides from exec_aten::StridesType (int32_t) to int64_t
49+
std::vector<int64_t> strides_vec(ndim);
50+
auto etensor_strides = etensor.strides();
51+
for (size_t i = 0; i < ndim; ++i) {
52+
strides_vec[i] = static_cast<int64_t>(etensor_strides[i]);
53+
}
54+
55+
// Map ETensor ScalarType to SlimTensor ScalarType
56+
c10::ScalarType dtype = static_cast<c10::ScalarType>(etensor.scalar_type());
57+
58+
// Step 2: Create SlimTensor on target device
59+
SlimTensor result = empty_strided(
60+
makeArrayRef(sizes_vec), makeArrayRef(strides_vec), dtype, target_device);
61+
62+
// Step 3: Copy data from ETensor (CPU) to SlimTensor (target device)
63+
// ETensor is always on CPU, so this handles CPU→CPU or CPU→CUDA copy
64+
const void* src_data = etensor.const_data_ptr();
65+
void* dst_data = result.data_ptr();
66+
size_t nbytes = etensor.nbytes();
67+
68+
if (nbytes > 0) {
69+
// const_cast is safe here because copy_ only reads from src_data
70+
result.storage()->copy_(
71+
dst_data, const_cast<void*>(src_data), nbytes, CPU_DEVICE);
72+
}
73+
74+
return result;
75+
}
76+
77+
/// Creates a SlimTensor from an ETensor pointer.
78+
///
79+
/// Convenience overload that accepts a pointer instead of a reference.
80+
///
81+
/// @param etensor Pointer to the source ETensor (must not be null).
82+
/// @param target_device The target device for the output SlimTensor.
83+
/// @return A new SlimTensor with data copied to the target device.
84+
inline SlimTensor from_etensor(
85+
const executorch::runtime::etensor::Tensor* etensor,
86+
const c10::Device& target_device = CPU_DEVICE) {
87+
ET_CHECK_MSG(
88+
etensor != nullptr, "from_etensor: etensor pointer cannot be nullptr");
89+
return from_etensor(*etensor, target_device);
90+
}
91+
92+
} // namespace executorch::backends::aoti::slim

backends/aoti/slim/factory/targets.bzl

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,3 +29,16 @@ def define_common_targets():
2929
"//executorch/backends/aoti/slim/util:size_util",
3030
],
3131
)
32+
33+
runtime.cxx_library(
34+
name = "from_etensor",
35+
headers = [
36+
"from_etensor.h",
37+
],
38+
visibility = ["@EXECUTORCH_CLIENTS"],
39+
exported_deps = [
40+
"//executorch/backends/aoti/slim/factory:empty",
41+
"//executorch/backends/aoti/slim/util:array_ref_util",
42+
"//executorch/runtime/core/portable_type:portable_type",
43+
],
44+
)

backends/aoti/slim/factory/test/targets.bzl

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,3 +44,16 @@ def define_common_targets():
4444
],
4545
**backend_kwargs
4646
)
47+
48+
runtime.cxx_test(
49+
name = "test_from_etensor" + backend_suffix,
50+
srcs = [
51+
"test_from_etensor.cpp",
52+
],
53+
deps = [
54+
"//executorch/backends/aoti/slim/core:storage",
55+
"//executorch/backends/aoti/slim/factory:from_etensor",
56+
"//executorch/runtime/core/exec_aten/testing_util:tensor_util",
57+
],
58+
**backend_kwargs
59+
)

0 commit comments

Comments
 (0)