xref: /aosp_15_r20/external/federated-compute/fcp/aggregation/core/input_tensor_list.cc (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2023 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #include "fcp/aggregation/core/input_tensor_list.h"
17 
18 #include <cstddef>
19 #include <initializer_list>
20 #include <utility>
21 #include <vector>
22 
23 #include "fcp/aggregation/core/tensor.h"
24 #include "fcp/base/new.h"
25 
26 namespace fcp {
27 namespace aggregation {
28 
InputTensorList(std::initializer_list<const Tensor * > list)29 InputTensorList::InputTensorList(std::initializer_list<const Tensor*> list)
30     : InputTensorList(list.size()) {
31   size_t i = 0;
32   for (const Tensor* t : list) {
33     data_ptr_[i++] = t;
34   }
35 }
36 
InputTensorList(size_t size)37 InputTensorList::InputTensorList(size_t size)
38     : size_(size), is_allocated_(size > kInlinedSize) {
39   if (is_allocated_) {
40     // Since the `allocated` union member is a class with user-defined
41     // constructors and destructors, to switch the active member, explicit
42     // placement new is needed. See
43     // https://en.cppreference.com/w/cpp/language/union.
44     new (&data_storage_.allocated) std::vector<const Tensor*>(size);
45     data_ptr_ = data_storage_.allocated.data();
46   } else {
47     // Use the new syntax to initialize elements to nullptr.
48     new (&data_storage_.inlined) int[size]();
49     data_ptr_ = data_storage_.inlined;
50   }
51 }
52 
InputTensorList(InputTensorList && other)53 InputTensorList::InputTensorList(InputTensorList&& other)
54     : size_(other.size_), is_allocated_(other.is_allocated_) {
55   MoveData(std::move(other));
56 }
57 
operator =(InputTensorList && other)58 InputTensorList& InputTensorList::operator=(InputTensorList&& other) {
59   // Destroy any existing allocated storage.
60   if (is_allocated_) {
61     data_storage_.allocated.~vector();
62   }
63   size_ = other.size_;
64   is_allocated_ = other.is_allocated_;
65   MoveData(std::move(other));
66   return *this;
67 }
68 
MoveData(InputTensorList && other)69 void InputTensorList::MoveData(InputTensorList&& other) {
70   if (is_allocated_) {
71     new (&data_storage_.allocated) std::vector<const Tensor*>;
72     data_storage_.allocated = std::move(other.data_storage_.allocated);
73     data_ptr_ = data_storage_.allocated.data();
74     other.data_storage_.allocated.~vector();
75   } else {
76     // If the storage is inlined copy the data; this is cheap since
77     // size_ < kInlinedSize.
78     for (size_t i = 0; i < size_; ++i) {
79       data_storage_.inlined[i] = other.data_storage_.inlined[i];
80     }
81     data_ptr_ = data_storage_.inlined;
82   }
83   new (&other.data_storage_.inlined) int[0]();
84   other.size_ = 0;
85   other.is_allocated_ = false;
86 }
87 
~InputTensorList()88 InputTensorList::~InputTensorList() {
89   // Since the `allocated` union member is a class with user-defined
90   // constructors and destructors, explicit destruction is needed. See
91   // https://en.cppreference.com/w/cpp/language/union.
92   if (is_allocated_) {
93     data_storage_.allocated.~vector();
94   }
95 }
96 
97 }  // namespace aggregation
98 }  // namespace fcp
99