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_LOCAL_SERVICE_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_LOCAL_SERVICE_H_ 18 19 #include <memory> 20 21 #include "tensorflow/compiler/xla/client/executable_build_options.h" 22 #include "tensorflow/compiler/xla/service/backend.h" 23 #include "tensorflow/compiler/xla/service/compiler.h" 24 #include "tensorflow/compiler/xla/service/device_memory_allocator.h" 25 #include "tensorflow/compiler/xla/service/executable.h" 26 #include "tensorflow/compiler/xla/service/service.h" 27 #include "tensorflow/compiler/xla/service/shaped_buffer.h" 28 #include "tensorflow/compiler/xla/statusor.h" 29 #include "tensorflow/compiler/xla/xla_data.pb.h" 30 #include "tensorflow/core/lib/gtl/array_slice.h" 31 #include "tensorflow/core/platform/stream_executor_no_cuda.h" 32 33 namespace xla { 34 35 // Service implementation that extends the XLA Service to leverage running 36 // in the same process as the client. 37 class LocalService : public Service { 38 public: 39 // Factory for creating a LocalService. 40 static StatusOr<std::unique_ptr<LocalService>> NewService( 41 const ServiceOptions& options); 42 43 // Builds an Executable with the given argument layouts and options. If 44 // result_layout is non-null, then the executable is compiled to produce a 45 // result of the given layout. If device_allocator is non-null, then the 46 // compiler may use it to allocate temp space on the device. The compiler is 47 // responsible for freeing any memory it allocates this way. 48 StatusOr<std::unique_ptr<Executable>> CompileExecutable( 49 const ComputationHandle& computation, 50 const tensorflow::gtl::ArraySlice<const Shape*> argument_layouts, 51 const ExecutableBuildOptions& options); 52 53 // Returns the device ordinal that corresponds to the given replica number. 54 // 55 // This returns an error if there is not a one-to-one correspondence of 56 // replicas to device ordinals, but is useful as a short term mechanism for 57 // the "easy" case where a single replica is a single device. 58 StatusOr<int> ReplicaNumberToDeviceOrdinal(int replica_number); 59 60 private: 61 explicit LocalService(const ServiceOptions& options, 62 std::unique_ptr<Backend> backend); 63 LocalService(const LocalService&) = delete; 64 void operator=(const LocalService&) = delete; 65 }; 66 67 } // namespace xla 68 69 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_LOCAL_SERVICE_H_ 70
