1 /* Copyright 2020 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_PJRT_GPU_DEVICE_H_ 17 #define TENSORFLOW_COMPILER_XLA_PJRT_GPU_DEVICE_H_ 18 19 #include <memory> 20 21 #include "tensorflow/compiler/xla/pjrt/distributed/client.h" 22 #include "tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h" 23 #include "tensorflow/compiler/xla/statusor.h" 24 #include "tensorflow/core/common_runtime/bfc_allocator.h" 25 26 namespace xla { 27 28 class GpuDevice : public PjRtStreamExecutorDevice { 29 public: 30 GpuDevice(int id, std::unique_ptr<LocalDeviceState> local_device_state, 31 std::string device_kind, int node_id); 32 }; 33 34 struct GpuAllocatorConfig { 35 enum class Kind { 36 kDefault, // Client picks the best option for the platform. 37 kPlatform, // The platform's default. 38 kBFC, // Allocator using a "Best-Fit with Coalescing" algorithm. Currently 39 // only available for GPU. 40 }; 41 Kind kind = Kind::kDefault; 42 43 // Only used if kind == kBFC. The maximum fraction of available memory to 44 // allocate. 45 double memory_fraction = 0.9; 46 47 // Only used if kind == kBFC. If true, the allocator will immediately allocate 48 // the maximum amount allowed by `memory_fraction`. This reduces 49 // fragmentation, allowing more of the total memory to be used. If false, the 50 // allocator will allocate more memory as allocations are requested. 51 bool preallocate = true; 52 }; 53 54 // distributed_client may be nullptr in non-distributed settings. 55 // distributed_client should be in the connected state before calling this 56 // function. 57 StatusOr<std::unique_ptr<PjRtClient>> GetGpuClient( 58 bool asynchronous, const GpuAllocatorConfig& allocator_config, 59 std::shared_ptr<DistributedRuntimeClient> distributed_client, int node_id); 60 61 } // namespace xla 62 63 #endif // TENSORFLOW_COMPILER_XLA_PJRT_GPU_DEVICE_H_ 64