xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/local_service.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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 #include <vector>
21 
22 #include "absl/types/span.h"
23 #include "tensorflow/compiler/xla/client/executable_build_options.h"
24 #include "tensorflow/compiler/xla/client/xla_computation.h"
25 #include "tensorflow/compiler/xla/service/backend.h"
26 #include "tensorflow/compiler/xla/service/compiler.h"
27 #include "tensorflow/compiler/xla/service/executable.h"
28 #include "tensorflow/compiler/xla/service/service.h"
29 #include "tensorflow/compiler/xla/service/shaped_buffer.h"
30 #include "tensorflow/compiler/xla/statusor.h"
31 #include "tensorflow/compiler/xla/xla_data.pb.h"
32 #include "tensorflow/core/platform/stream_executor_no_cuda.h"
33 #include "tensorflow/stream_executor/device_memory_allocator.h"
34 
35 namespace xla {
36 
37 // Service implementation that extends the XLA Service to leverage running
38 // in the same process as the client.
39 class LocalService : public Service {
40  public:
41   // Factory for creating a LocalService.
42   static StatusOr<std::unique_ptr<LocalService>> NewService(
43       const ServiceOptions& options);
44 
45   // Builds Executables with the given XlaComputation, argument layouts and
46   // options. If result_layout is non-null, then the executable is compiled to
47   // produce a result of the given layout.  If device_allocator is non-null,
48   // then the compiler may use it to allocate temp space on the device.  The
49   // compiler is responsible for freeing any memory it allocates this way.
50   StatusOr<std::vector<std::unique_ptr<Executable>>> CompileExecutables(
51       const XlaComputation& computation,
52       const absl::Span<const Shape* const> argument_layouts,
53       const ExecutableBuildOptions& build_options);
54 
55   // Same as CompileExecutables() above, but return AotCompilationResult objects
56   // (instead of Executable objects), which can be persisted to later load
57   // Executable objects.
58   StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>>
59   CompileAotResults(const XlaComputation& computation,
60                     const absl::Span<const Shape* const> argument_layouts,
61                     const ExecutableBuildOptions& build_options);
62 
63   // Returns the device ordinal that corresponds to the given replica number.
64   //
65   // This returns an error if there is not a one-to-one correspondence of
66   // replicas to device ordinals, but is useful as a short term mechanism for
67   // the "easy" case where a single replica is a single device.
68   StatusOr<int> ReplicaNumberToDeviceOrdinal(int replica_number);
69 
70   // Converts a GlobalDataHandle into a pointer to a ShapedBuffer that's valid
71   // as long as the handle is valid.
72   StatusOr<const ShapedBuffer*> GlobalDataToShapedBuffer(
73       const GlobalDataHandle& data, int replica_number);
74 
75   // Registers a vector of shaped buffers of device memory, one per replica, and
76   // returns a corresponding handle that can be used for talking to XLA clients.
77   StatusOr<GlobalDataHandle> RegisterReplicatedBuffers(
78       std::vector<ScopedShapedBuffer> replicated_buffers,
79       const std::string& tag);
80 
81  private:
82   explicit LocalService(const ServiceOptions& options,
83                         std::unique_ptr<Backend> backend);
84   LocalService(const LocalService&) = delete;
85   void operator=(const LocalService&) = delete;
86 
87   // Validates the computation argument layouts, and returns the corresponding
88   // HloModuleConfig.
89   StatusOr<std::unique_ptr<HloModuleConfig>> GetHloModuleConfig(
90       const XlaComputation& computation,
91       const absl::Span<const Shape* const> argument_layouts,
92       const ExecutableBuildOptions& build_options);
93 };
94 
95 }  // namespace xla
96 
97 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_LOCAL_SERVICE_H_
98