| 135 | } |
| 136 | |
| 137 | StatusOr<std::vector<GlobalDataHandle>> AllocationTracker::DeconstructTuple( |
| 138 | const GlobalDataHandle& data) { |
| 139 | tensorflow::mutex_lock lock(mutex_); |
| 140 | |
| 141 | TF_ASSIGN_OR_RETURN(std::vector<const ShapedBuffer*> replicated_buffers, |
| 142 | ResolveInternal(data)); |
| 143 | // We only need to care about replica id 0 here, since the GlobalDataHandle is |
| 144 | // the same for all buffers across replicas. |
| 145 | const ShapedBuffer* shaped_buffer = replicated_buffers[0]; |
| 146 | if (!shaped_buffer->on_host_shape().IsTuple()) { |
| 147 | return InvalidArgument("global data handle %d is not a tuple", |
| 148 | data.handle()); |
| 149 | } |
| 150 | // If the on-host representation is a tuple, then the on-device one should be |
| 151 | // as well. |
| 152 | TF_RET_CHECK(shaped_buffer->on_device_shape().IsTuple()); |
| 153 | |
| 154 | if (ShapeUtil::IsNestedTuple(shaped_buffer->on_device_shape())) { |
| 155 | return Unimplemented("Deconstructing nested tuples is not implemented."); |
| 156 | } |
| 157 | |
| 158 | std::vector<GlobalDataHandle> element_handles; |
| 159 | for (int i = 0; |
| 160 | i < ShapeUtil::TupleElementCount(shaped_buffer->on_device_shape()); |
| 161 | ++i) { |
| 162 | auto element_buffer = ShapedBuffer( |
| 163 | ShapeUtil::GetTupleElementShape(shaped_buffer->on_host_shape(), i), |
| 164 | ShapeUtil::GetTupleElementShape(shaped_buffer->on_device_shape(), i), |
| 165 | shaped_buffer->platform(), shaped_buffer->device_ordinal()); |
| 166 | element_buffer.set_buffer(shaped_buffer->buffer(/*index=*/{i}), |
| 167 | /*index=*/{}); |
| 168 | std::vector<ShapedBuffer> replicated_buffers; |
| 169 | replicated_buffers.push_back(std::move(element_buffer)); |
| 170 | TF_ASSIGN_OR_RETURN( |
| 171 | GlobalDataHandle element_handle, |
| 172 | RegisterInternal(std::move(replicated_buffers), "deconstructed tuple")); |
| 173 | |
| 174 | element_handles.push_back(element_handle); |
| 175 | } |
| 176 | return std::move(element_handles); |
| 177 | } |
| 178 | |
| 179 | StatusOr<std::vector<const ShapedBuffer*>> AllocationTracker::Resolve( |
| 180 | const GlobalDataHandle& data) const { |
| 181 | tensorflow::mutex_lock lock(mutex_); |
| 182 | return AllocationTracker::ResolveInternal(data); |
| 183 | } |
| 184 | |
| 185 | StatusOr<const ShapedBuffer*> AllocationTracker::ResolveForReplica( |
| 186 | const GlobalDataHandle& data, int replica_id) const { |
| 187 | tensorflow::mutex_lock lock(mutex_); |
| 188 | TF_ASSIGN_OR_RETURN(std::vector<const ShapedBuffer*> replicated_buffers, |
| 189 | ResolveInternal(data)); |
| 190 | if (replica_id >= replicated_buffers.size()) { |
| 191 | return InvalidArgument( |
| 192 | "Requesting buffer for replica %d, but found buffers only for %lu " |
| 193 | "replicas.", |
| 194 | replica_id, replicated_buffers.size()); |
nothing calls this directly
no test coverage detected