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_CLIENT_COMPILE_ONLY_CLIENT_H_ 17 #define TENSORFLOW_COMPILER_XLA_CLIENT_COMPILE_ONLY_CLIENT_H_ 18 19 #include "tensorflow/compiler/xla/client/client.h" 20 #include "tensorflow/compiler/xla/client/xla_computation.h" 21 #include "tensorflow/compiler/xla/service/compile_only_service.h" 22 #include "tensorflow/compiler/xla/service/compiler.h" 23 #include "tensorflow/compiler/xla/statusor.h" 24 #include "tensorflow/compiler/xla/xla_data.pb.h" 25 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 26 27 namespace xla { 28 29 // An XLA Client specialization for doing ahead-of-time compilation. This does 30 // not require (or attempt to instantiate) an execution-capable backend for the 31 // relevant platform. 32 class CompileOnlyClient : public Client { 33 public: CompileOnlyClient(CompileOnlyService * service)34 explicit CompileOnlyClient(CompileOnlyService* service) 35 : Client(service), compiler_service_(service) {} 36 37 CompileOnlyClient(const CompileOnlyClient&) = delete; 38 void operator=(const CompileOnlyClient&) = delete; 39 40 // A description of an xla computation to compile using CompileAheadOfTime. 41 struct AotXlaComputationInstance { 42 const XlaComputation* computation; 43 // Inform the compiler of the expected layout for arguments. 44 std::vector<const Shape*> argument_layouts; 45 // Specifies the expected result layout. 46 const Shape* result_layout; 47 }; 48 49 // Compiles a list of xla computations for ahead-of-time execution. 50 // This is intended for use in static compilation. The |options| 51 // parameter describes the target for which the compiler should emit 52 // code. |metadata|, if provided, is populated during compilation. 53 StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>> 54 CompileAheadOfTime( 55 const absl::Span<const AotXlaComputationInstance> computations, 56 const AotCompilationOptions& options, 57 std::unique_ptr<AotCompilationMetadata>* metadata = nullptr); 58 59 // Create a Hlo module config for the given program shape and arguments. 60 // execution_options is optional; if not given a default is used. 61 StatusOr<std::unique_ptr<HloModuleConfig>> CreateModuleConfig( 62 const ProgramShape& program_shape, 63 absl::Span<const Shape* const> argument_shapes, 64 const ExecutionOptions* execution_options); 65 66 // Returns the size of a pointer in bytes for a given triple. 67 static int64 PointerSizeForTriple(absl::string_view triple); 68 69 private: 70 CompileOnlyService* compiler_service_; 71 }; 72 73 } // namespace xla 74 75 #endif // TENSORFLOW_COMPILER_XLA_CLIENT_COMPILE_ONLY_CLIENT_H_ 76