Skip to content

Commit 1d37332

Browse files
shoumikhinfacebook-github-bot
shoumikhin
authored and
facebook-github-bot
committed
Add APIs to make from an existing Tensor and clone by copying and owning the data. (#5304)
Summary: Pull Request resolved: #5304 . Reviewed By: kirklandsign Differential Revision: D62557645 fbshipit-source-id: 6789ab225465824f522344c0f55ae70194fee45c
1 parent 40e6e52 commit 1d37332

File tree

2 files changed

+207
-0
lines changed

2 files changed

+207
-0
lines changed

‎extension/tensor/tensor_ptr.h‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,35 @@ inline TensorPtr make_tensor_ptr(const TensorPtr& tensor) {
104104
#endif // USE_ATEN_LIB
105105
}
106106

107+
/**
108+
* Creates a TensorPtr that manages a new Tensor with the same properties
109+
* as the given Tensor, sharing the same data without owning it.
110+
*
111+
* @param tensor The Tensor whose properties are to be used to create a new
112+
* TensorPtr.
113+
* @return A new TensorPtr that manages a Tensor with the same properties as the
114+
* original.
115+
*/
116+
inline TensorPtr make_tensor_ptr(const exec_aten::Tensor& tensor) {
117+
return make_tensor_ptr(make_tensor_impl_ptr(
118+
tensor.scalar_type(),
119+
std::vector<exec_aten::SizesType>(
120+
tensor.sizes().begin(), tensor.sizes().end()),
121+
tensor.mutable_data_ptr(),
122+
#ifndef USE_ATEN_LIB
123+
std::vector<exec_aten::DimOrderType>(
124+
tensor.dim_order().begin(), tensor.dim_order().end()),
125+
std::vector<exec_aten::StridesType>(
126+
tensor.strides().begin(), tensor.strides().end()),
127+
tensor.shape_dynamism()
128+
#else // USE_ATEN_LIB
129+
{},
130+
std::vector<exec_aten::StridesType>(
131+
tensor.strides().begin(), tensor.strides().end())
132+
#endif // USE_ATEN_LIB
133+
));
134+
}
135+
107136
/**
108137
* Creates a TensorPtr that manages a Tensor with the specified properties.
109138
*
@@ -241,6 +270,49 @@ inline TensorPtr make_tensor_ptr(
241270
dynamism));
242271
}
243272

273+
/**
274+
* Creates a TensorPtr that manages a new Tensor with the same properties
275+
* as the given Tensor, but with a copy of the data owned by the returned
276+
* TensorPtr.
277+
*
278+
* @param tensor The Tensor to clone.
279+
* @return A new TensorPtr that manages a Tensor with the same properties as the
280+
* original but with copied data.
281+
*/
282+
inline TensorPtr clone_tensor_ptr(const exec_aten::Tensor& tensor) {
283+
return make_tensor_ptr(make_tensor_impl_ptr(
284+
tensor.scalar_type(),
285+
std::vector<exec_aten::SizesType>(
286+
tensor.sizes().begin(), tensor.sizes().end()),
287+
std::vector<uint8_t>(
288+
(uint8_t*)tensor.const_data_ptr(),
289+
(uint8_t*)tensor.const_data_ptr() + tensor.nbytes()),
290+
#ifndef USE_ATEN_LIB
291+
std::vector<exec_aten::DimOrderType>(
292+
tensor.dim_order().begin(), tensor.dim_order().end()),
293+
std::vector<exec_aten::StridesType>(
294+
tensor.strides().begin(), tensor.strides().end()),
295+
tensor.shape_dynamism()
296+
#else // USE_ATEN_LIB
297+
{},
298+
std::vector<exec_aten::StridesType>(
299+
tensor.strides().begin(), tensor.strides().end())
300+
#endif // USE_ATEN_LIB
301+
));
302+
}
303+
304+
/**
305+
* Creates a new TensorPtr by cloning the given TensorPtr, copying the
306+
* underlying data.
307+
*
308+
* @param tensor The TensorPtr to clone.
309+
* @return A new TensorPtr that manages a Tensor with the same properties as the
310+
* original but with copied data.
311+
*/
312+
inline TensorPtr clone_tensor_ptr(const TensorPtr& tensor) {
313+
return clone_tensor_ptr(*tensor);
314+
}
315+
244316
/**
245317
* Resizes the Tensor managed by the given TensorPtr to the new sizes provided.
246318
*

‎extension/tensor/test/tensor_ptr_test.cpp‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,3 +315,138 @@ TEST_F(TensorPtrTest, TensorSharingImplResizingAffectsBothVector) {
315315
EXPECT_EQ(tensor1->size(0), 4);
316316
EXPECT_EQ(tensor1->size(1), 3);
317317
}
318+
319+
TEST_F(TensorPtrTest, MakeTensorPtrFromExistingTensorInt32) {
320+
std::vector<int32_t> data = {1, 2, 3, 4};
321+
auto tensor = make_tensor_ptr({2, 2}, data);
322+
auto new_tensor = make_tensor_ptr(*tensor);
323+
324+
EXPECT_EQ(new_tensor->dim(), tensor->dim());
325+
EXPECT_EQ(new_tensor->size(0), tensor->size(0));
326+
EXPECT_EQ(new_tensor->size(1), tensor->size(1));
327+
EXPECT_EQ(
328+
new_tensor->const_data_ptr<int32_t>(), tensor->const_data_ptr<int32_t>());
329+
EXPECT_EQ(new_tensor->scalar_type(), exec_aten::ScalarType::Int);
330+
}
331+
332+
TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorInt32) {
333+
std::vector<int32_t> data = {1, 2, 3, 4};
334+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
335+
auto cloned_tensor = clone_tensor_ptr(*tensor);
336+
337+
EXPECT_EQ(cloned_tensor->dim(), tensor->dim());
338+
EXPECT_EQ(cloned_tensor->size(0), tensor->size(0));
339+
EXPECT_EQ(cloned_tensor->size(1), tensor->size(1));
340+
EXPECT_NE(
341+
cloned_tensor->const_data_ptr<int32_t>(),
342+
tensor->const_data_ptr<int32_t>());
343+
EXPECT_EQ(cloned_tensor->const_data_ptr<int32_t>()[0], 1);
344+
EXPECT_EQ(cloned_tensor->const_data_ptr<int32_t>()[3], 4);
345+
EXPECT_EQ(cloned_tensor->scalar_type(), exec_aten::ScalarType::Int);
346+
}
347+
348+
TEST_F(TensorPtrTest, CloneTensorPtrFromTensorPtrInt32) {
349+
std::vector<int32_t> data = {1, 2, 3, 4};
350+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
351+
auto cloned_tensor = clone_tensor_ptr(tensor);
352+
353+
EXPECT_EQ(cloned_tensor->dim(), tensor->dim());
354+
EXPECT_EQ(cloned_tensor->size(0), tensor->size(0));
355+
EXPECT_EQ(cloned_tensor->size(1), tensor->size(1));
356+
EXPECT_NE(
357+
cloned_tensor->const_data_ptr<int32_t>(),
358+
tensor->const_data_ptr<int32_t>());
359+
EXPECT_EQ(cloned_tensor->const_data_ptr<int32_t>()[0], 1);
360+
EXPECT_EQ(cloned_tensor->const_data_ptr<int32_t>()[3], 4);
361+
EXPECT_EQ(cloned_tensor->scalar_type(), exec_aten::ScalarType::Int);
362+
}
363+
364+
TEST_F(TensorPtrTest, MakeTensorPtrFromExistingTensorDouble) {
365+
std::vector<double> data = {1.0, 2.0, 3.0, 4.0};
366+
auto tensor = make_tensor_ptr({2, 2}, data);
367+
auto new_tensor = make_tensor_ptr(*tensor);
368+
369+
EXPECT_EQ(new_tensor->dim(), tensor->dim());
370+
EXPECT_EQ(new_tensor->size(0), tensor->size(0));
371+
EXPECT_EQ(new_tensor->size(1), tensor->size(1));
372+
EXPECT_EQ(
373+
new_tensor->const_data_ptr<double>(), tensor->const_data_ptr<double>());
374+
EXPECT_EQ(new_tensor->scalar_type(), exec_aten::ScalarType::Double);
375+
}
376+
377+
TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorDouble) {
378+
std::vector<double> data = {1.0, 2.0, 3.0, 4.0};
379+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
380+
auto cloned_tensor = clone_tensor_ptr(*tensor);
381+
382+
EXPECT_EQ(cloned_tensor->dim(), tensor->dim());
383+
EXPECT_EQ(cloned_tensor->size(0), tensor->size(0));
384+
EXPECT_EQ(cloned_tensor->size(1), tensor->size(1));
385+
EXPECT_NE(
386+
cloned_tensor->const_data_ptr<double>(),
387+
tensor->const_data_ptr<double>());
388+
EXPECT_EQ(cloned_tensor->const_data_ptr<double>()[0], 1.0);
389+
EXPECT_EQ(cloned_tensor->const_data_ptr<double>()[3], 4.0);
390+
EXPECT_EQ(cloned_tensor->scalar_type(), exec_aten::ScalarType::Double);
391+
}
392+
393+
TEST_F(TensorPtrTest, CloneTensorPtrFromTensorPtrDouble) {
394+
std::vector<double> data = {1.0, 2.0, 3.0, 4.0};
395+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
396+
auto cloned_tensor = clone_tensor_ptr(tensor);
397+
398+
EXPECT_EQ(cloned_tensor->dim(), tensor->dim());
399+
EXPECT_EQ(cloned_tensor->size(0), tensor->size(0));
400+
EXPECT_EQ(cloned_tensor->size(1), tensor->size(1));
401+
EXPECT_NE(
402+
cloned_tensor->const_data_ptr<double>(),
403+
tensor->const_data_ptr<double>());
404+
EXPECT_EQ(cloned_tensor->const_data_ptr<double>()[0], 1.0);
405+
EXPECT_EQ(cloned_tensor->const_data_ptr<double>()[3], 4.0);
406+
EXPECT_EQ(cloned_tensor->scalar_type(), exec_aten::ScalarType::Double);
407+
}
408+
409+
TEST_F(TensorPtrTest, MakeTensorPtrFromExistingTensorInt64) {
410+
std::vector<int64_t> data = {100, 200, 300, 400};
411+
auto tensor = make_tensor_ptr({2, 2}, data);
412+
auto new_tensor = make_tensor_ptr(*tensor);
413+
414+
EXPECT_EQ(new_tensor->dim(), tensor->dim());
415+
EXPECT_EQ(new_tensor->size(0), tensor->size(0));
416+
EXPECT_EQ(new_tensor->size(1), tensor->size(1));
417+
EXPECT_EQ(
418+
new_tensor->const_data_ptr<int64_t>(), tensor->const_data_ptr<int64_t>());
419+
EXPECT_EQ(new_tensor->scalar_type(), exec_aten::ScalarType::Long);
420+
}
421+
422+
TEST_F(TensorPtrTest, CloneTensorPtrFromExistingTensorInt64) {
423+
std::vector<int64_t> data = {100, 200, 300, 400};
424+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
425+
auto cloned_tensor = clone_tensor_ptr(*tensor);
426+
427+
EXPECT_EQ(cloned_tensor->dim(), tensor->dim());
428+
EXPECT_EQ(cloned_tensor->size(0), tensor->size(0));
429+
EXPECT_EQ(cloned_tensor->size(1), tensor->size(1));
430+
EXPECT_NE(
431+
cloned_tensor->const_data_ptr<int64_t>(),
432+
tensor->const_data_ptr<int64_t>());
433+
EXPECT_EQ(cloned_tensor->const_data_ptr<int64_t>()[0], 100);
434+
EXPECT_EQ(cloned_tensor->const_data_ptr<int64_t>()[3], 400);
435+
EXPECT_EQ(cloned_tensor->scalar_type(), exec_aten::ScalarType::Long);
436+
}
437+
438+
TEST_F(TensorPtrTest, CloneTensorPtrFromTensorPtrInt64) {
439+
std::vector<int64_t> data = {100, 200, 300, 400};
440+
auto tensor = make_tensor_ptr({2, 2}, std::move(data));
441+
auto cloned_tensor = clone_tensor_ptr(tensor);
442+
443+
EXPECT_EQ(cloned_tensor->dim(), tensor->dim());
444+
EXPECT_EQ(cloned_tensor->size(0), tensor->size(0));
445+
EXPECT_EQ(cloned_tensor->size(1), tensor->size(1));
446+
EXPECT_NE(
447+
cloned_tensor->const_data_ptr<int64_t>(),
448+
tensor->const_data_ptr<int64_t>());
449+
EXPECT_EQ(cloned_tensor->const_data_ptr<int64_t>()[0], 100);
450+
EXPECT_EQ(cloned_tensor->const_data_ptr<int64_t>()[3], 400);
451+
EXPECT_EQ(cloned_tensor->scalar_type(), exec_aten::ScalarType::Long);
452+
}

0 commit comments

Comments
 (0)