11e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 21e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 31e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsLicensed under the Apache License, Version 2.0 (the "License"); 41e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsyou may not use this file except in compliance with the License. 51e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsYou may obtain a copy of the License at 61e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 71e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins http://www.apache.org/licenses/LICENSE-2.0 81e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 91e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsUnless required by applicable law or agreed to in writing, software 101e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsdistributed under the License is distributed on an "AS IS" BASIS, 111e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 121e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsSee the License for the specific language governing permissions and 131e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinslimitations under the License. 141e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins==============================================================================*/ 151e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 161e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/compiler/xla/client/computation.h" 171e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 181e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/compiler/xla/ptr_util.h" 191e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/compiler/xla/status_macros.h" 201e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins#include "tensorflow/core/lib/core/errors.h" 211e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 221e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsnamespace xla { 231e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 241e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsComputation::Computation() : parent_(nullptr) {} 251e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 261e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsComputation::Computation(ServiceInterface* parent, 271e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins const ComputationHandle& handle) 281e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins : handle_(handle), parent_(parent) {} 291e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 301e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsComputation::Computation(Computation&& computation) 312e7e1d57bd27a0acb7b46bbccf227a9e4325e38aA. Unique TensorFlower : handle_(std::move(computation.handle_)), parent_(computation.parent_) { 321e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins computation.ResetWithoutFreeing(); 331e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins} 341e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 351e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsvoid Computation::Reset() { 368e46bc5420217b6460f6aaae0a3616fe8a04cff8A. Unique TensorFlower // TODO(b/34469253) deallocate any owned computation. 371e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins ResetWithoutFreeing(); 381e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins} 391e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 401e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsStatusOr<std::unique_ptr<SessionModule>> Computation::Snapshot() const { 411e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins SnapshotComputationRequest request; 421e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins *request.mutable_computation() = handle_; 431e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins SnapshotComputationResponse response; 441e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 451e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins TF_RETURN_IF_ERROR(parent_->SnapshotComputation(&request, &response)); 461e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 471e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins return WrapUnique(response.release_module()); 481e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins} 491e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 501e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsComputation::~Computation() { Reset(); } 511e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 521e67c90e2caceeff82d09793d1ef5fa0300d219bPeter HawkinsComputation& Computation::operator=(Computation&& computation) { 531e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins if (&computation != this) { 541e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins Reset(); 551e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins handle_ = computation.handle_; 561e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins parent_ = computation.parent_; 571e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins computation.ResetWithoutFreeing(); 581e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins } 591e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins return *this; 601e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins} 611e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 621e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkinsvoid Computation::ResetWithoutFreeing() { 631e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins handle_.Clear(); 641e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins parent_ = nullptr; 651e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins} 661e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins 675b984d1bd80768a954063f9bb220373ff459bbfdChris LearyStatusOr<ProgramShape> Computation::GetProgramShape() const { 685b984d1bd80768a954063f9bb220373ff459bbfdChris Leary GetComputationShapeRequest request; 695b984d1bd80768a954063f9bb220373ff459bbfdChris Leary *request.mutable_computation() = handle_; 705b984d1bd80768a954063f9bb220373ff459bbfdChris Leary GetComputationShapeResponse response; 715b984d1bd80768a954063f9bb220373ff459bbfdChris Leary 725b984d1bd80768a954063f9bb220373ff459bbfdChris Leary TF_RETURN_IF_ERROR(parent_->GetComputationShape(&request, &response)); 735b984d1bd80768a954063f9bb220373ff459bbfdChris Leary 745b984d1bd80768a954063f9bb220373ff459bbfdChris Leary return std::move(*response.mutable_program_shape()); 755b984d1bd80768a954063f9bb220373ff459bbfdChris Leary} 765b984d1bd80768a954063f9bb220373ff459bbfdChris Leary 771e67c90e2caceeff82d09793d1ef5fa0300d219bPeter Hawkins} // namespace xla 78