1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_ONLY_SERVICE_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_ONLY_SERVICE_H_ 18 19 #include "tensorflow/compiler/xla/service/backend.h" 20 #include "tensorflow/compiler/xla/service/compiler.h" 21 #include "tensorflow/compiler/xla/service/service.h" 22 #include "tensorflow/compiler/xla/statusor.h" 23 #include "tensorflow/compiler/xla/xla_data.pb.h" 24 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 25 26 namespace xla { 27 28 // An XLA Service specialization for ahead-of-time compilation. This only 29 // instantiates a Compiler object for the relevant platform; it does not 30 // instantiate or require an execution backend. 31 class CompileOnlyService : public Service { 32 public: 33 // Factory for creating a CompileOnlyService. The parameter platform is the 34 // platform that the service should target. If platform is null then the 35 // default platform is used. 36 static StatusOr<std::unique_ptr<CompileOnlyService>> NewService( 37 se::Platform* platform); 38 static StatusOr<std::unique_ptr<CompileOnlyService>> NewService( 39 const ServiceOptions& options); 40 41 // A description of a xla computation to compile using CompileAheadOfTime. 42 struct AotXlaComputationInstance { 43 HloModuleProto computation; 44 std::vector<const Shape*> argument_layouts; 45 const Shape* result_layout = nullptr; 46 }; 47 48 // Compiles a list of xla computations for ahead-of-time execution. This is 49 // intended for use in static compilation. See 50 // |CompileOnlyClient::CompileAheadOfTime| for additional details. 51 StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>> 52 CompileAheadOfTime( 53 const absl::Span<const AotXlaComputationInstance> computations, 54 const AotCompilationOptions& options); 55 56 StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>> 57 CompileAheadOfTime( 58 const absl::Span<const AotXlaComputationInstance> computations, 59 const AotCompilationOptions& options, 60 std::unique_ptr<AotCompilationMetadata>* metadata); 61 GetDeviceHandles(const GetDeviceHandlesRequest * arg,GetDeviceHandlesResponse * result)62 Status GetDeviceHandles(const GetDeviceHandlesRequest* arg, 63 GetDeviceHandlesResponse* result) override { 64 return Unimplemented("CompileOnlyService does not support devices."); 65 } WaitForExecution(const WaitForExecutionRequest * arg,WaitForExecutionResponse * result)66 Status WaitForExecution(const WaitForExecutionRequest* arg, 67 WaitForExecutionResponse* result) override { 68 return Unimplemented("CompileOnlyService does not support execution."); 69 } TransferToServer(const TransferToServerRequest * arg,TransferToServerResponse * result)70 Status TransferToServer(const TransferToServerRequest* arg, 71 TransferToServerResponse* result) override { 72 return Unimplemented( 73 "CompileOnlyService does not support device data transfers."); 74 } TransferToInfeed(const TransferToInfeedRequest * arg,TransferToInfeedResponse * result)75 Status TransferToInfeed(const TransferToInfeedRequest* arg, 76 TransferToInfeedResponse* result) override { 77 return Unimplemented( 78 "CompileOnlyService does not support device data transfers."); 79 } TransferFromOutfeed(const TransferFromOutfeedRequest * arg,TransferFromOutfeedResponse * result)80 Status TransferFromOutfeed(const TransferFromOutfeedRequest* arg, 81 TransferFromOutfeedResponse* result) override { 82 return Unimplemented( 83 "CompileOnlyService does not support device data transfers."); 84 } ResetDevice(const ResetDeviceRequest * arg,ResetDeviceResponse * result)85 Status ResetDevice(const ResetDeviceRequest* arg, 86 ResetDeviceResponse* result) override { 87 return Unimplemented("CompileOnlyService does not support devices."); 88 } 89 90 private: 91 explicit CompileOnlyService(const ServiceOptions& options, 92 Compiler* compiler); 93 CompileOnlyService(const CompileOnlyService&) = delete; 94 void operator=(const CompileOnlyService&) = delete; 95 96 // The compiler for the target platform. This is included in place of 97 // the Service::execute_backend_'s compiler, since execute_backend_ is a 98 // nullptr in CompileOnlyService. 99 Compiler* compiler_; 100 }; 101 102 } // namespace xla 103 104 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_ONLY_SERVICE_H_ 105